Preserve EA dtype in DataFrame.stack

Author: TomAugspurger

doc/source/whatsnew/v0.24.0.txt

@@ -724,6 +724,8 @@ update the ``ExtensionDtype._metadata`` tuple to match the signature of your
 - Updated the ``.type`` attribute for ``PeriodDtype``, ``DatetimeTZDtype``, and ``IntervalDtype`` to be instances of the dtype (``Period``, ``Timestamp``, and ``Interval`` respectively) (:issue:`22938`)
 - :func:`ExtensionArray.isna` is allowed to return an ``ExtensionArray`` (:issue:`22325`).
 - Support for reduction operations such as ``sum``, ``mean`` via opt-in base class method override (:issue:`22762`)
+- :meth:`DataFrame.stack` no longer converts to object dtype for DataFrames where each column has the same extension dtype. The output Series will have the same dtype as the columns (:issue:`23077`).
+
 
 .. _whatsnew_0240.api.incompatibilities:
 

pandas/core/reshape/reshape.py

@@ -470,8 +470,15 @@ def factorize(index):
         if is_extension_array_dtype(dtype):
             arr = dtype.construct_array_type()
             new_values = arr._concat_same_type([
-                col for _, col in frame.iteritems()
+                col._values for _, col in frame.iteritems()
             ])
+            # final take to get the order correct.
+            # idx is an indexer like
+            # [c0r0, c1r0, c2r0, ...,
+            #  c0r1, c1r1, c241, ...]
+            idx = np.arange(N * K).reshape(K, N).T.ravel()
+            new_values = new_values.take(idx)
+
         else:
             # homogeneous, non-EA
             new_values = frame.values.ravel()

pandas/tests/extension/base/reshaping.py

@@ -170,3 +170,11 @@ def test_merge(self, data, na_value):
                  [data[0], data[0], data[1], data[2], na_value],
                  dtype=data.dtype)})
         self.assert_frame_equal(res, exp[['ext', 'int1', 'key', 'int2']])
+
+    def test_stack(self, data):
+        df = pd.DataFrame({"A": data[:5], "B": data[:5]})
+        result = df.stack()
+        assert result.dtype == df.A.dtype
+        result = result.astype(object)
+        expected = df.astype(object).stack()
+        self.assert_series_equal(result, expected)

pandas/tests/frame/test_reshape.py

@@ -872,6 +872,17 @@ def test_stack_preserve_categorical_dtype(self, ordered, labels):
 
         tm.assert_series_equal(result, expected)
 
+    def test_stack_preserve_categorical_dtype_values(self):
+        # GH-23077
+        cat = pd.Categorical(['a', 'a', 'b', 'c'])
+        df = pd.DataFrame({"A": cat, "B": cat})
+        result = df.stack()
+        index = pd.MultiIndex.from_product([[0, 1, 2, 3], ['A', 'B']])
+        expected = pd.Series(pd.Categorical(['a', 'a', 'a', 'a',
+                                             'b', 'b', 'c', 'c']),
+                             index=index)
+        tm.assert_series_equal(result, expected)
+
     @pytest.mark.parametrize("level", [0, 'baz'])
     def test_unstack_swaplevel_sortlevel(self, level):
         # GH 20994