python - Getting boolean pandas column that supports NA/ is nullable

Question

How can I create a pandas dataframe column with dtype bool (or int for that matter) with support for Nan/missing values?

When I try like this:

d = {'one' : np.ma.MaskedArray([True, False, True, True], mask = [0,0,1,0]),
'two' : pd.Series([1., 2., 3., 4.], index=['a', 'b', 'c', 'd'])}
df = pd.DataFrame(d)
print (df.dtypes)
print (df)

column one is implicitly converted to object. Likewise similar for ints:

d = {'one' : np.ma.MaskedArray([1,3,2,1], mask = [0,0,1,0]),
'two' : pd.Series([1., 2., 3., 4.], index=['a', 'b', 'c', 'd'])}
df = pd.DataFrame(d)
print (df.dtypes)
print (df)

one is here implicitly converted to float64, and I'd prefer if I stayed in int domain and not handle floating point arithmetic with its idiosyncrasies (always have tolerance when comparing, rounding errors, etc.)

Answer 1

pandas >= 1.0

As of pandas 1.0.0 (January 2020), there is experimental support for nullable booleans directly:

In [183]: df.one.astype('boolean')
Out[183]:
a     True
b    False
c     <NA>
d     True
Name: one, dtype: object

In this version, pandas will also use pd.NA instead of np.nan in the integer case:

In [166]: df.astype('Int64')
Out[166]:
    one  two
a     1    1
b     3    2
c  <NA>    3
d     1    4

---

pandas >= 0.24

In the integer case, as of pandas 0.24 (January 2019), you can use nullable integers to achieve what you want:

In [165]: df
Out[165]:
   one  two
a  1.0  1.0
b  3.0  2.0
c  NaN  3.0
d  1.0  4.0

In [166]: df.astype('Int64')
Out[166]:
   one  two
a    1    1
b    3    2
c  NaN    3
d    1    4

This works by converting the backing array to an arrays.IntegerArray, and there is no equivalent thing for booleans, but some work in that direction is discussed in this GitHub issue and this PyData talk. You could write your own extension type to cover this case as well, but if you can live with your booleans being represented by the integers 0 and 1, one approach could be the following:

In [183]: df.one
Out[183]:
a     True
b    False
c      NaN
d     True
Name: one, dtype: object

In [184]: (df.one * 1).astype('Int64')
Out[184]:
a      1
b      0
c    NaN
d      1
Name: one, dtype: Int64