DF.__setitem__ creates extension column when given extension scalar (#34875)

* Bugfix to make DF.__setitem__ create extension column instead of object column when given an extension scalar

* removed bad whitespace

* Apply suggestions from code review

Checking if extension dtype via built in function instead of manually

Co-authored-by: Tom Augspurger <[email protected]>

* added missing :

* modified cast_extension_scalar_to_array test to include an Interval type

* added user-facing test for extension type bug

* fixed pep8 issues

* added note about bug in setting series to scalar extension type

* corrected order of imports

* corrected order of imports

* fixed black formatting errors

* removed extra comma

* updated cast_scalar_to_arr to support tuple shape for extension dtype

* removed unneeded code

* added coverage for datetime with timezone in extension_array test

* added TODO

* correct line that was too long

* fixed dtype issue with tz test

* creating distinct arrays for each column

* resolving mypy error

* added docstring info and test

* removed unneeded import

* flattened else case in init

* refactored extension type column fix

* reverted docstring changes

* reverted docstring changes

* removed unneeded imports

* reverted test changes

* fixed construct_1d_arraylike bug

* reorganized if statements

* moved what's new statement to correct file

* created new test for period df construction

* added assert_frame_equal to period_data test

* Using pandas array instead of df constructor for better test

Co-authored-by: Joris Van den Bossche <[email protected]>

* changed wording

* pylint fixes

* parameterized test and added comment

* removed extra comma

* parameterized test

* renamed test

Co-authored-by: Tom Augspurger <[email protected]>
Co-authored-by: Joris Van den Bossche <[email protected]>

Author: 3 people

doc/source/whatsnew/v1.1.0.rst

@@ -1146,6 +1146,7 @@ ExtensionArray
 - Fixed bug where :meth:`StringArray.memory_usage` was not implemented (:issue:`33963`)
 - Fixed bug where :meth:`DataFrameGroupBy` would ignore the ``min_count`` argument for aggregations on nullable boolean dtypes (:issue:`34051`)
 - Fixed bug that `DataFrame(columns=.., dtype='string')` would fail (:issue:`27953`, :issue:`33623`)
+- Bug where :class:`DataFrame` column set to scalar extension type was considered an object type rather than the extension type (:issue:`34832`)
 - Fixed bug in ``IntegerArray.astype`` to correctly copy the mask as well (:issue:`34931`).
 
 Other

pandas/core/frame.py

@@ -520,25 +520,43 @@ def __init__(
                     mgr = init_ndarray(data, index, columns, dtype=dtype, copy=copy)
             else:
                 mgr = init_dict({}, index, columns, dtype=dtype)
+        # For data is scalar
         else:
-            try:
-                arr = np.array(data, dtype=dtype, copy=copy)
-            except (ValueError, TypeError) as err:
-                exc = TypeError(
-                    "DataFrame constructor called with "
-                    f"incompatible data and dtype: {err}"
-                )
-                raise exc from err
+            if index is None or columns is None:
+                raise ValueError("DataFrame constructor not properly called!")
+
+            if not dtype:
+                dtype, _ = infer_dtype_from_scalar(data, pandas_dtype=True)
+
+            # For data is a scalar extension dtype
+            if is_extension_array_dtype(dtype):
+
+                values = [
+                    construct_1d_arraylike_from_scalar(data, len(index), dtype)
+                    for _ in range(len(columns))
+                ]
+                mgr = arrays_to_mgr(values, columns, index, columns, dtype=None)
+            else:
+                # Attempt to coerce to a numpy array
+                try:
+                    arr = np.array(data, dtype=dtype, copy=copy)
+                except (ValueError, TypeError) as err:
+                    exc = TypeError(
+                        "DataFrame constructor called with "
+                        f"incompatible data and dtype: {err}"
+                    )
+                    raise exc from err
+
+                if arr.ndim != 0:
+                    raise ValueError("DataFrame constructor not properly called!")
 
-            if arr.ndim == 0 and index is not None and columns is not None:
                 values = cast_scalar_to_array(
                     (len(index), len(columns)), data, dtype=dtype
                 )
+
                 mgr = init_ndarray(
                     values, index, columns, dtype=values.dtype, copy=False
                 )
-            else:
-                raise ValueError("DataFrame constructor not properly called!")
 
         NDFrame.__init__(self, mgr)
 
@@ -3740,7 +3758,13 @@ def reindexer(value):
             infer_dtype, _ = infer_dtype_from_scalar(value, pandas_dtype=True)
 
             # upcast
-            value = cast_scalar_to_array(len(self.index), value)
+            if is_extension_array_dtype(infer_dtype):
+                value = construct_1d_arraylike_from_scalar(
+                    value, len(self.index), infer_dtype
+                )
+            else:
+                value = cast_scalar_to_array(len(self.index), value)
+
             value = maybe_cast_to_datetime(value, infer_dtype)
 
         # return internal types directly

pandas/tests/frame/indexing/test_setitem.py

@@ -1,7 +1,18 @@
 import numpy as np
 import pytest
 
-from pandas import Categorical, DataFrame, Index, Series, Timestamp, date_range
+from pandas.core.dtypes.dtypes import DatetimeTZDtype, IntervalDtype, PeriodDtype
+
+from pandas import (
+    Categorical,
+    DataFrame,
+    Index,
+    Interval,
+    Period,
+    Series,
+    Timestamp,
+    date_range,
+)
 import pandas._testing as tm
 from pandas.core.arrays import SparseArray
 
@@ -150,3 +161,23 @@ def test_setitem_dict_preserves_dtypes(self):
                 "c": float(b),
             }
         tm.assert_frame_equal(df, expected)
+
+    @pytest.mark.parametrize(
+        "obj,dtype",
+        [
+            (Period("2020-01"), PeriodDtype("M")),
+            (Interval(left=0, right=5), IntervalDtype("int64")),
+            (
+                Timestamp("2011-01-01", tz="US/Eastern"),
+                DatetimeTZDtype(tz="US/Eastern"),
+            ),
+        ],
+    )
+    def test_setitem_extension_types(self, obj, dtype):
+        # GH: 34832
+        expected = DataFrame({"idx": [1, 2, 3], "obj": Series([obj] * 3, dtype=dtype)})
+
+        df = DataFrame({"idx": [1, 2, 3]})
+        df["obj"] = obj
+
+        tm.assert_frame_equal(df, expected)

pandas/tests/frame/methods/test_combine_first.py

@@ -199,12 +199,14 @@ def test_combine_first_timezone(self):
             columns=["UTCdatetime", "abc"],
             data=data1,
             index=pd.date_range("20140627", periods=1),
+            dtype="object",
         )
         data2 = pd.to_datetime("20121212 12:12").tz_localize("UTC")
         df2 = pd.DataFrame(
             columns=["UTCdatetime", "xyz"],
             data=data2,
             index=pd.date_range("20140628", periods=1),
+            dtype="object",
         )
         res = df2[["UTCdatetime"]].combine_first(df1)
         exp = pd.DataFrame(
@@ -217,10 +219,14 @@ def test_combine_first_timezone(self):
             },
             columns=["UTCdatetime", "abc"],
             index=pd.date_range("20140627", periods=2, freq="D"),
+            dtype="object",
         )
-        tm.assert_frame_equal(res, exp)
         assert res["UTCdatetime"].dtype == "datetime64[ns, UTC]"
         assert res["abc"].dtype == "datetime64[ns, UTC]"
+        # Need to cast all to "obejct" because combine_first does not retain dtypes:
+        # GH Issue 7509
+        res = res.astype("object")
+        tm.assert_frame_equal(res, exp)
 
         # see gh-10567
         dts1 = pd.date_range("2015-01-01", "2015-01-05", tz="UTC")

pandas/tests/frame/test_constructors.py

@@ -14,13 +14,16 @@
 from pandas.compat.numpy import _np_version_under1p19
 
 from pandas.core.dtypes.common import is_integer_dtype
+from pandas.core.dtypes.dtypes import DatetimeTZDtype, IntervalDtype, PeriodDtype
 
 import pandas as pd
 from pandas import (
     Categorical,
     DataFrame,
     Index,
+    Interval,
     MultiIndex,
+    Period,
     RangeIndex,
     Series,
     Timedelta,
@@ -700,7 +703,7 @@ def create_data(constructor):
         tm.assert_frame_equal(result_timedelta, expected)
         tm.assert_frame_equal(result_Timedelta, expected)
 
-    def test_constructor_period(self):
+    def test_constructor_period_dict(self):
         # PeriodIndex
         a = pd.PeriodIndex(["2012-01", "NaT", "2012-04"], freq="M")
         b = pd.PeriodIndex(["2012-02-01", "2012-03-01", "NaT"], freq="D")
@@ -713,6 +716,29 @@ def test_constructor_period(self):
         assert df["a"].dtype == a.dtype
         assert df["b"].dtype == b.dtype
 
+    @pytest.mark.parametrize(
+        "data,dtype",
+        [
+            (Period("2020-01"), PeriodDtype("M")),
+            (Interval(left=0, right=5), IntervalDtype("int64")),
+            (
+                Timestamp("2011-01-01", tz="US/Eastern"),
+                DatetimeTZDtype(tz="US/Eastern"),
+            ),
+        ],
+    )
+    def test_constructor_extension_scalar_data(self, data, dtype):
+        # GH 34832
+        df = DataFrame(index=[0, 1], columns=["a", "b"], data=data)
+
+        assert df["a"].dtype == dtype
+        assert df["b"].dtype == dtype
+
+        arr = pd.array([data] * 2, dtype=dtype)
+        expected = DataFrame({"a": arr, "b": arr})
+
+        tm.assert_frame_equal(df, expected)
+
     def test_nested_dict_frame_constructor(self):
         rng = pd.period_range("1/1/2000", periods=5)
         df = DataFrame(np.random.randn(10, 5), columns=rng)