Replace a value with NA in an integer column while retaining the type closest to original

user_918959 0
Viewed 84

I've got a Numpy array with arbitrary flat structural data type with some columns as integers, some as datetime64, some floats etc...

>>> import numpy as np
>>> ar = np.array([(1, -1), (-1, 2)], dtype=[('a', 'i2'), ('b', 'i4')])

A Pandas dataframe is directly created from this numpy array precisely following its schema:

>>> import pandas as pd
>>> df = pd.DataFrame(ar)
>>> df
   a  b
0  1 -1
1 -1  2
>>> df.dtypes
a    int16
b    int32
dtype: object

Now I want to say replace all -1 values in this data frame to NA while retaining the original data types as closely as possible. Basically this:

>>> df_with_na = df.replace(-1, pd.NA)
>>> df_with_na
      a     b
0     1  <NA>
1  <NA>     2

except that this does not work because the column dtypes were converted to object:

>>> df_with_na.dtypes
a    object
b    object
dtype: object

which does not match the desired result of just mapping int16 to Int16:

>>> df_with_na.dtypes
a    Int16
b    Int32
dtype: object

Is there a way to convert the int16 etc whatever blindly to the masked nullable extension types, or preferably even making the replace with NA to change to the extension type if it wasn't originally, without having to write yet another set of lookup tables with all possible extensions?

pandas python-3.x python
1 Answers

I suggest going ahead and converting all your dtypes to the nullable ones before you do anything to the dataframe:

>>> df.dtypes
a    int16
b    int32
dtype: object
>>> df.convert_dtypes().dtypes
a    Int16
b    Int32
dtype: object

You might want to set all those kwargs other than convert_integer to False to be on the safe side.

user_50673110
Related