Use IntergerArray for integer arrays with null

Author: Jiang Yue

pandas/_libs/lib.pyx

@@ -944,6 +944,7 @@ cdef class Seen:
         bint int_             # seen_int
         bint bool_            # seen_bool
         bint null_            # seen_null
+        bint nan_             # seen_np.nan
         bint uint_            # seen_uint (unsigned integer)
         bint sint_            # seen_sint (signed integer)
         bint float_           # seen_float
@@ -967,6 +968,7 @@ cdef class Seen:
         self.int_ = 0
         self.bool_ = 0
         self.null_ = 0
+        self.nan_ = 0
         self.uint_ = 0
         self.sint_ = 0
         self.float_ = 0
@@ -1905,7 +1907,7 @@ def maybe_convert_numeric(ndarray[object] values, set na_values,
 @cython.wraparound(False)
 def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
                           bint safe=0, bint convert_datetime=0,
-                          bint convert_timedelta=0):
+                          bint convert_timedelta=0, to_integer_array=False):
     """
     Type inference function-- convert object array to proper dtype
     """
@@ -1929,6 +1931,7 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
     ints = np.empty(n, dtype='i8')
     uints = np.empty(n, dtype='u8')
     bools = np.empty(n, dtype=np.uint8)
+    mask = np.full(n, False)
 
     if convert_datetime:
         datetimes = np.empty(n, dtype='M8[ns]')
@@ -1946,6 +1949,7 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
         if val is None:
             seen.null_ = 1
             floats[i] = complexes[i] = fnan
+            mask[i] = True
         elif val is NaT:
             if convert_datetime:
                 idatetimes[i] = NPY_NAT
@@ -1956,6 +1960,10 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
             if not (convert_datetime or convert_timedelta):
                 seen.object_ = 1
                 break
+        elif val is np.nan:
+            seen.nan_ = 1
+            mask[i] = True
+            floats[i] = complexes[i] = val
         elif util.is_bool_object(val):
             seen.bool_ = 1
             bools[i] = val
@@ -2037,11 +2045,19 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
 
     if not seen.object_:
         if not safe:
-            if seen.null_:
+            if seen.null_ or seen.nan_:
                 if seen.is_float_or_complex:
                     if seen.complex_:
                         return complexes
-                    elif seen.float_ or seen.int_:
+                    elif seen.float_:
+                        return floats
+                    elif seen.int_:
+                        if to_integer_array:
+                            from pandas.core.arrays import IntegerArray
+                            return IntegerArray(ints, mask)
+                        else:
+                            return floats
+                    elif seen.nan_:
                         return floats
             else:
                 if not seen.bool_:
@@ -2071,7 +2087,7 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
                     if seen.complex_:
                         if not seen.int_:
                             return complexes
-                    elif seen.float_:
+                    elif seen.float_ or seen.nan_:
                         if not seen.int_:
                             return floats
             else:
@@ -2086,7 +2102,7 @@ def maybe_convert_objects(ndarray[object] objects, bint try_float=0,
                         if seen.complex_:
                             if not seen.int_:
                                 return complexes
-                        elif seen.float_:
+                        elif seen.float_ or seen.nan_:
                             if not seen.int_:
                                 return floats
                         elif seen.int_:

pandas/core/frame.py

@@ -389,7 +389,8 @@ def _constructor_expanddim(self):
     # ----------------------------------------------------------------------
     # Constructors
 
-    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False):
+    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False,
+                 to_integer_array=False):
         if data is None:
             data = {}
         if dtype is not None:
@@ -442,7 +443,8 @@ def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False):
                 if is_list_like(data[0]) and getattr(data[0], "ndim", 1) == 1:
                     if is_named_tuple(data[0]) and columns is None:
                         columns = data[0]._fields
-                    arrays, columns = to_arrays(data, columns, dtype=dtype)
+                    arrays, columns = to_arrays(data, columns, dtype=dtype,
+                                                to_integer_array=to_integer_array)
                     columns = ensure_index(columns)
 
                     # set the index

pandas/core/internals/construction.py

@@ -437,7 +437,8 @@ def _get_axes(N, K, index, columns):
 # Conversion of Inputs to Arrays
 
 
-def to_arrays(data, columns, coerce_float=False, dtype=None):
+def to_arrays(data, columns, coerce_float=False, dtype=None,
+              to_integer_array=False):
     """
     Return list of arrays, columns.
     """
@@ -464,7 +465,8 @@ def to_arrays(data, columns, coerce_float=False, dtype=None):
         return _list_to_arrays(data, columns, coerce_float=coerce_float, dtype=dtype)
     elif isinstance(data[0], abc.Mapping):
         return _list_of_dict_to_arrays(
-            data, columns, coerce_float=coerce_float, dtype=dtype
+            data, columns, coerce_float=coerce_float, dtype=dtype,
+            to_integer_array=to_integer_array
         )
     elif isinstance(data[0], ABCSeries):
         return _list_of_series_to_arrays(
@@ -535,7 +537,8 @@ def _list_of_series_to_arrays(data, columns, coerce_float=False, dtype=None):
         return values.T, columns
 
 
-def _list_of_dict_to_arrays(data, columns, coerce_float=False, dtype=None):
+def _list_of_dict_to_arrays(data, columns, coerce_float=False, dtype=None,
+                            to_integer_array=False):
     if columns is None:
         gen = (list(x.keys()) for x in data)
         sort = not any(isinstance(d, OrderedDict) for d in data)
@@ -547,11 +550,13 @@ def _list_of_dict_to_arrays(data, columns, coerce_float=False, dtype=None):
 
     content = list(lib.dicts_to_array(data, list(columns)).T)
     return _convert_object_array(
-        content, columns, dtype=dtype, coerce_float=coerce_float
+        content, columns, dtype=dtype, coerce_float=coerce_float,
+        to_integer_array=to_integer_array
     )
 
 
-def _convert_object_array(content, columns, coerce_float=False, dtype=None):
+def _convert_object_array(content, columns, coerce_float=False, dtype=None,
+                          to_integer_array=False):
     if columns is None:
         columns = ibase.default_index(len(content))
     else:
@@ -565,7 +570,8 @@ def _convert_object_array(content, columns, coerce_float=False, dtype=None):
     # provide soft conversion of object dtypes
     def convert(arr):
         if dtype != object and dtype != np.object:
-            arr = lib.maybe_convert_objects(arr, try_float=coerce_float)
+            arr = lib.maybe_convert_objects(arr, try_float=coerce_float,
+                                            to_integer_array=to_integer_array)
             arr = maybe_cast_to_datetime(arr, dtype)
         return arr
 

pandas/io/json/_normalize.py

@@ -120,6 +120,7 @@ def json_normalize(
     errors: Optional[str] = "raise",
     sep: str = ".",
     max_level: Optional[int] = None,
+    to_integer_array: Optional[bool] = False
 ):
     """
     Normalize semi-structured JSON data into a flat table.
@@ -264,7 +265,7 @@ def _pull_field(js, spec):
             # TODO: handle record value which are lists, at least error
             #       reasonably
             data = nested_to_record(data, sep=sep, max_level=max_level)
-        return DataFrame(data)
+        return DataFrame(data, to_integer_array=to_integer_array)
     elif not isinstance(record_path, list):
         record_path = [record_path]
 
@@ -326,7 +327,7 @@ def _recursive_extract(data, path, seen_meta, level=0):
 
     _recursive_extract(data, record_path, {}, level=0)
 
-    result = DataFrame(records)
+    result = DataFrame(records, to_integer_array=to_integer_array)
 
     if record_prefix is not None:
         result = result.rename(columns=lambda x: "{p}{c}".format(p=record_prefix, c=x))