Why does np.select return 'nan' as a string instead of np.nan when np.nan is set as default value?

user_13765813 0
Viewed 1182

I am using np.select to create a new column based on multiple conditions applied to other columns. Here is a simple example:

df = pd.DataFrame({'A': [0, 3, 4], 'B': [10, 0, 2]})

mask1 = (df['A'] == 0)
mask2 = (df['A'] == 4)

df = df.assign(C = np.select([mask1, mask2], ['Cond1', 'Cond2'], default=np.nan))

Using np.select I am expecting that if neither of the conditions are met I should get an np.nan value in the 'C' column, but when I look for NaNs in the dataframe using df.isna().sum() I don't get any.

I've tried looking this up but I cannot seem to find the answer. What am I missing here?

pandas dataframe numpy python
2 Answers

It is worth knowing that pandas has it's own arrays (which are built on top of numpy arrays), but have a general missing values indicator pd.NA (note: it is still experimental).

So to make use of the pandas arrays, we can do:

mask1 = (df['A'] == 0)
mask2 = (df['A'] == 4)

c = np.select([mask1, mask2], ['Cond1', 'Cond2'], default=pd.NA)
df = df.assign(C=c).convert_dtypes()

   A   B      C
0  0  10  Cond1
1  3   0   <NA>
2  4   2  Cond2

Then if we check the dtypes, we see that we are using the pandas arrays:

A     Int64
B     Int64
C    string
dtype: object
user_90812670

Finally found out why: numpy arrays cannot contain multiple data types. Since I am filling column C with strings, numpy will insert the default value as 'nan' instead of np.nan which is a float.

EDIT: a good workaround:

df = df.assign(C = pd.Series(np.select([mask1, mask2], ['Cond1', 'Cond2'], default=np.nan)).replace({'nan': np.nan}))
user_137658130
Related