Finish implementation

Author: phofl

doc/source/whatsnew/v1.6.0.rst

@@ -32,6 +32,7 @@ Other enhancements
 - :meth:`.DataFrameGroupBy.quantile` and :meth:`.SeriesGroupBy.quantile` now preserve nullable dtypes instead of casting to numpy dtypes (:issue:`37493`)
 - :meth:`Series.add_suffix`, :meth:`DataFrame.add_suffix`, :meth:`Series.add_prefix` and :meth:`DataFrame.add_prefix` support an ``axis`` argument. If ``axis`` is set, the default behaviour of which axis to consider can be overwritten (:issue:`47819`)
 - :func:`assert_frame_equal` now shows the first element where the DataFrames differ, analogously to ``pytest``'s output (:issue:`47910`)
+- Added new argument ``use_nullable_dtypes`` to :func:`read_csv` to enable automatic conversion to nullable dtypes (:issue:`36712`)
 - Added ``index`` parameter to :meth:`DataFrame.to_dict` (:issue:`46398`)
 - Added metadata propagation for binary operators on :class:`DataFrame` (:issue:`28283`)
 - :class:`.CategoricalConversionWarning`, :class:`.InvalidComparison`, :class:`.InvalidVersion`, :class:`.LossySetitemError`, and :class:`.NoBufferPresent` are now exposed in ``pandas.errors`` (:issue:`27656`)

pandas/_libs/parsers.pyx

@@ -936,6 +936,7 @@ cdef class TextReader:
             bint na_filter = 0
             int64_t num_cols
             dict result
+            bint use_nullable_dtypes
 
         start = self.parser_start
 
@@ -1056,8 +1057,15 @@ cdef class TextReader:
                     self._free_na_set(na_hashset)
 
             # don't try to upcast EAs
-            if na_count > 0 and not is_extension_array_dtype(col_dtype) or self.use_nullable_dtypes:
-                col_res = _maybe_upcast(col_res, use_nullable_dtypes=self.use_nullable_dtypes)
+            print(col_dtype)
+            if (
+                na_count > 0 and not is_extension_array_dtype(col_dtype)
+                or self.use_nullable_dtypes
+            ):
+                use_nullable_dtypes = self.use_nullable_dtypes and col_dtype is None
+                col_res = _maybe_upcast(
+                    col_res, use_nullable_dtypes=use_nullable_dtypes
+                )
 
             if col_res is None:
                 raise ParserError(f'Unable to parse column {i}')

pandas/io/parsers/base_parser.py

@@ -117,7 +117,7 @@ def __init__(self, kwds) -> None:
 
         self.dtype = copy(kwds.get("dtype", None))
         self.converters = kwds.get("converters")
-        self.use_nullable_dtypes = kwds.get("use_nullable_dtypes")
+        self.use_nullable_dtypes = kwds.get("use_nullable_dtypes", False)
 
         self.true_values = kwds.get("true_values")
         self.false_values = kwds.get("false_values")
@@ -516,7 +516,7 @@ def _agg_index(self, index, try_parse_dates: bool = True) -> Index:
             )
 
             arr, _ = self._infer_types(
-                arr, col_na_values | col_na_fvalues, try_num_bool
+                arr, col_na_values | col_na_fvalues, cast_type, try_num_bool
             )
             arrays.append(arr)
 
@@ -582,7 +582,10 @@ def _convert_to_ndarrays(
                     values = lib.map_infer_mask(values, conv_f, mask)
 
                 cvals, na_count = self._infer_types(
-                    values, set(col_na_values) | col_na_fvalues, try_num_bool=False
+                    values,
+                    set(col_na_values) | col_na_fvalues,
+                    cast_type,
+                    try_num_bool=False,
                 )
             else:
                 is_ea = is_extension_array_dtype(cast_type)
@@ -593,7 +596,7 @@ def _convert_to_ndarrays(
 
                 # general type inference and conversion
                 cvals, na_count = self._infer_types(
-                    values, set(col_na_values) | col_na_fvalues, try_num_bool
+                    values, set(col_na_values) | col_na_fvalues, cast_type, try_num_bool
                 )
 
                 # type specified in dtype param or cast_type is an EA
@@ -684,14 +687,15 @@ def _set(x) -> int:
 
         return noconvert_columns
 
-    def _infer_types(self, values, na_values, try_num_bool: bool = True):
+    def _infer_types(self, values, na_values, cast_type, try_num_bool: bool = True):
         """
         Infer types of values, possibly casting
 
         Parameters
         ----------
         values : ndarray
         na_values : set
+        cast_type: Specifies if we want to cast explicitly
         try_num_bool : bool, default try
            try to cast values to numeric (first preference) or boolean
 
@@ -712,26 +716,24 @@ def _infer_types(self, values, na_values, try_num_bool: bool = True):
                 np.putmask(values, mask, np.nan)
             return values, na_count
 
+        use_nullable_dtypes = self.use_nullable_dtypes and cast_type is None
+
         if try_num_bool and is_object_dtype(values.dtype):
             # exclude e.g DatetimeIndex here
             try:
                 result, result_mask = lib.maybe_convert_numeric(
                     values,
                     na_values,
                     False,
-                    convert_to_masked_nullable=self.use_nullable_dtypes,
+                    convert_to_masked_nullable=use_nullable_dtypes,
                 )
             except (ValueError, TypeError):
                 # e.g. encountering datetime string gets ValueError
                 #  TypeError can be raised in floatify
                 na_count = parsers.sanitize_objects(values, na_values)
-
-                if self.use_nullable_dtypes:
-                    result = StringDtype().construct_array_type()._from_sequence(values)
-                else:
-                    result = values
+                result = values
             else:
-                if self.use_nullable_dtypes:
+                if use_nullable_dtypes:
                     if result_mask is None:
                         result_mask = np.zeros(result.shape, dtype="bool")
 
@@ -751,11 +753,18 @@ def _infer_types(self, values, na_values, try_num_bool: bool = True):
                 na_count = parsers.sanitize_objects(values, na_values)
 
         if result.dtype == np.object_ and try_num_bool:
-            result, _ = libops.maybe_convert_bool(
+            result, mask = libops.maybe_convert_bool(
                 np.asarray(values),
                 true_values=self.true_values,
                 false_values=self.false_values,
+                convert_to_masked_nullable=use_nullable_dtypes,
             )
+            if result.dtype == np.bool_ and use_nullable_dtypes:
+                if mask is None:
+                    mask = np.zeros(result.shape, dtype=np.bool_)
+                result = BooleanArray(result, mask)
+            elif result.dtype == np.object_ and use_nullable_dtypes:
+                result = StringDtype().construct_array_type()._from_sequence(values)
 
         return result, na_count
 

pandas/tests/io/parser/dtypes/test_dtypes_basic.py

@@ -385,3 +385,45 @@ def test_dtypes_defaultdict_invalid(all_parsers):
     parser = all_parsers
     with pytest.raises(TypeError, match="not understood"):
         parser.read_csv(StringIO(data), dtype=dtype)
+
+
+def test_use_nullabla_dtypes(all_parsers):
+    # GH#36712
+
+    parser = all_parsers
+
+    data = """a,b,c,d,e,f,g,h,i
+1,2.5,True,a,,,,,12-31-2019
+3,4.5,False,b,6,7.5,True,a,12-31-2019
+"""
+    result = parser.read_csv(
+        StringIO(data), use_nullable_dtypes=True, parse_dates=["i"]
+    )
+    expected = DataFrame(
+        {
+            "a": pd.Series([1, 3], dtype="Int64"),
+            "b": pd.Series([2.5, 4.5], dtype="Float64"),
+            "c": pd.Series([True, False], dtype="boolean"),
+            "d": pd.Series(["a", "b"], dtype="string"),
+            "e": pd.Series([pd.NA, 6], dtype="Int64"),
+            "f": pd.Series([pd.NA, 7.5], dtype="Float64"),
+            "g": pd.Series([pd.NA, True], dtype="boolean"),
+            "h": pd.Series([pd.NA, "a"], dtype="string"),
+            "i": pd.Series([Timestamp("2019-12-31")] * 2),
+        }
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+def test_use_nullabla_dtypes_and_dtype(all_parsers):
+    # GH#36712
+
+    parser = all_parsers
+
+    data = """a,b
+1,2.5
+,
+"""
+    result = parser.read_csv(StringIO(data), use_nullable_dtypes=True, dtype="float64")
+    expected = DataFrame({"a": [1.0, np.nan], "b": [2.5, np.nan]})
+    tm.assert_frame_equal(result, expected)