Repository: pandas-dev/pandas

Issue:
Handle ExtensionArrays in Series.unstack / DataFrame.stack
Here's a test for `Series.unstack`

```diff
diff --git a/pandas/tests/extension/base/reshaping.py b/pandas/tests/extension/base/reshaping.py
index 0340289e0..a040fba63 100644
--- a/pandas/tests/extension/base/reshaping.py
+++ b/pandas/tests/extension/base/reshaping.py
@@ -171,3 +171,16 @@ class BaseReshapingTests(BaseExtensionTests):
                  [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_unstack(self, data):
+        data = data[:4]
+        ser = pd.Series(
+            data,
+            index=pd.MultiIndex.from_product([["A", "B"],
+                                              ["a", "b"]]),
+        )
+        result = ser.unstack()
+        expected = pd.DataFrame({"a": data.take([0, 2]),
+                                 "b": data.take([1, 3])},
+                                index=['A', 'B'])
+        self.assert_frame_equal(result, expected)
```

We don't do so well right now. Just categorical passes. 
```
========================================================================= FAILURES =========================================================================
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:112: IndexError: tuple index out of range
/Users/taugspurger/sandbox/pandas/pandas/core/reshape/reshape.py:204: AttributeError: 'IntervalArray' object has no attribute 'reshape'
/Users/taugspurger/sandbox/pandas/pandas/tests/extension/decimal/array.py:55: TypeError: All values must be of type <class 'decimal.Decimal'>
/Users/taugspurger/sandbox/pandas/pandas/tests/extension/json/array.py:88: TypeError: list indices must be integers or slices, not NoneType
```

No test for DataFrame.stack. In https://github.com/pandas-dev/pandas/pull/22862/files there's a WIP for stack that's based around `ExtensionArray._concat_same_type` instead of `.reshape`


Discussion:
I'm working on this now.
I have a WIP for `Series[extension_array].unstack` at https://github.com/pandas-dev/pandas/compare/master...TomAugspurger:ea-unstack

We can write `unstack` as a composition of

- `loc` (take)
- `reindex`
- concat

It seems to work... but it's quite slow. If you're reshaping to `(n_rows, n_columns)`, you'll end up with `n_columns` locs and reindexes each.

We should be able to improve the (maybe common?) case of a "uniform" index (like what you get from `from_product`), without too much additional code. That's just a bunch of uniform `take`s with no reindexes.
@TomAugspurger is it slower compared to the current implementation for non-consolidatable blocks? Or how does it currently work for them?
DataFrame.stack() has a regression from 0.23.4

```
In [3]: df = pd.DataFrame({"A": pd.Categorical(['a', 'b']), "B": pd.Categorical(['a', 'b'])})

In [4]: df.stack()
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-c3a581ae119f> in <module>
----> 1 df.stack()

~/sandbox/pandas/pandas/core/frame.py in stack(self, level, dropna)
   5753             return stack_multiple(self, level, dropna=dropna)
   5754         else:
-> 5755             return stack(self, level, dropna=dropna)
   5756
   5757     def unstack(self, level=-1, fill_value=None):

~/sandbox/pandas/pandas/core/reshape/reshape.py in stack(frame, level, dropna)
    471             arr = dtype.construct_array_type()
    472             new_values = arr._concat_same_type([
--> 473                 col for _, col in frame.iteritems()
    474             ])
    475         else:

~/sandbox/pandas/pandas/core/arrays/categorical.py in _concat_same_type(self, to_concat)
   2286         from pandas.core.dtypes.concat import _concat_categorical
   2287
-> 2288         return _concat_categorical(to_concat)
   2289
   2290     def _formatting_values(self):

~/sandbox/pandas/pandas/core/dtypes/concat.py in _concat_categorical(to_concat, axis)
    239         # when all categories are identical
    240         first = to_concat[0]
--> 241         if all(first.is_dtype_equal(other) for other in to_concat[1:]):
    242             return union_categoricals(categoricals)
    243

~/sandbox/pandas/pandas/core/dtypes/concat.py in <genexpr>(.0)
    239         # when all categories are identical
    240         first = to_concat[0]
--> 241         if all(first.is_dtype_equal(other) for other in to_concat[1:]):
    242             return union_categoricals(categoricals)
    243

~/sandbox/pandas/pandas/core/generic.py in __getattr__(self, name)
   4636             if self._info_axis._can_hold_identifiers_and_holds_name(name):
   4637                 return self[name]
-> 4638             return object.__getattribute__(self, name)
   4639
   4640     def __setattr__(self, name, value):

AttributeError: 'Series' object has no attribute 'is_dtype_equal'

```


previously that returned an object-dtype Series. Ideally, it would be a categorical.
@jorisvandenbossche I have a new implementation that lowers the overhead for EAs. Will post the timings in the PR (once I've written them).

It'll basically be identical to previously, but

1. ExtensionArrays are no longer cast to ndarrays (good)
2. An extra `ExtensionArray.take` per column + one extra `concat` (bad)
 


Resolving patch:
diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py
--- a/pandas/core/internals/blocks.py
+++ b/pandas/core/internals/blocks.py
@@ -35,9 +35,9 @@
     is_numeric_v_string_like, is_extension_type,
     is_extension_array_dtype,
     is_list_like,
-    is_sparse,
     is_re,
     is_re_compilable,
+    is_sparse,
     pandas_dtype)
 from pandas.core.dtypes.cast import (
     maybe_downcast_to_dtype,
diff --git a/pandas/core/reshape/reshape.py b/pandas/core/reshape/reshape.py
--- a/pandas/core/reshape/reshape.py
+++ b/pandas/core/reshape/reshape.py
@@ -494,8 +494,9 @@ 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()
             ])
+            new_values = _reorder_for_extension_array_stack(new_values, N, K)
         else:
             # homogeneous, non-EA
             new_values = frame.values.ravel()
@@ -624,16 +625,32 @@ def _convert_level_number(level_num, columns):
             slice_len = loc.stop - loc.start
 
         if slice_len != levsize:
-            chunk = this.loc[:, this.columns[loc]]
+            chunk = this[this.columns[loc]]
             chunk.columns = level_vals.take(chunk.columns.labels[-1])
             value_slice = chunk.reindex(columns=level_vals_used).values
         else:
-            if frame._is_mixed_type:
-                value_slice = this.loc[:, this.columns[loc]].values
+            if (frame._is_homogeneous_type and
+                    is_extension_array_dtype(frame.dtypes.iloc[0])):
+                dtype = this[this.columns[loc]].dtypes.iloc[0]
+                subset = this[this.columns[loc]]
+
+                value_slice = dtype.construct_array_type()._concat_same_type(
+                    [x._values for _, x in subset.iteritems()]
+                )
+                N, K = this.shape
+                idx = np.arange(N * K).reshape(K, N).T.ravel()
+                value_slice = value_slice.take(idx)
+
+            elif frame._is_mixed_type:
+                value_slice = this[this.columns[loc]].values
             else:
                 value_slice = this.values[:, loc]
 
-        new_data[key] = value_slice.ravel()
+        if value_slice.ndim > 1:
+            # i.e. not extension
+            value_slice = value_slice.ravel()
+
+        new_data[key] = value_slice
 
     if len(drop_cols) > 0:
         new_columns = new_columns.difference(drop_cols)
@@ -971,3 +988,38 @@ def make_axis_dummies(frame, axis='minor', transform=None):
     values = values.take(labels, axis=0)
 
     return DataFrame(values, columns=items, index=frame.index)
+
+
+def _reorder_for_extension_array_stack(arr, n_rows, n_columns):
+    """
+    Re-orders the values when stacking multiple extension-arrays.
+
+    The indirect stacking method used for EAs requires a followup
+    take to get the order correct.
+
+    Parameters
+    ----------
+    arr : ExtensionArray
+    n_rows, n_columns : int
+        The number of rows and columns in the original DataFrame.
+
+    Returns
+    -------
+    taken : ExtensionArray
+        The original `arr` with elements re-ordered appropriately
+
+    Examples
+    --------
+    >>> arr = np.array(['a', 'b', 'c', 'd', 'e', 'f'])
+    >>> _reorder_for_extension_array_stack(arr, 2, 3)
+    array(['a', 'c', 'e', 'b', 'd', 'f'], dtype='<U1')
+
+    >>> _reorder_for_extension_array_stack(arr, 3, 2)
+    array(['a', 'd', 'b', 'e', 'c', 'f'], dtype='<U1')
+    """
+    # final take to get the order correct.
+    # idx is an indexer like
+    # [c0r0, c1r0, c2r0, ...,
+    #  c0r1, c1r1, c2r1, ...]
+    idx = np.arange(n_rows * n_columns).reshape(n_columns, n_rows).T.ravel()
+    return arr.take(idx)


Test patch:
diff --git a/pandas/tests/extension/base/reshaping.py b/pandas/tests/extension/base/reshaping.py
--- a/pandas/tests/extension/base/reshaping.py
+++ b/pandas/tests/extension/base/reshaping.py
@@ -173,6 +173,28 @@ def test_merge(self, data, na_value):
                  dtype=data.dtype)})
         self.assert_frame_equal(res, exp[['ext', 'int1', 'key', 'int2']])
 
+    @pytest.mark.parametrize("columns", [
+        ["A", "B"],
+        pd.MultiIndex.from_tuples([('A', 'a'), ('A', 'b')],
+                                  names=['outer', 'inner']),
+    ])
+    def test_stack(self, data, columns):
+        df = pd.DataFrame({"A": data[:5], "B": data[:5]})
+        df.columns = columns
+        result = df.stack()
+        expected = df.astype(object).stack()
+        # we need a second astype(object), in case the constructor inferred
+        # object -> specialized, as is done for period.
+        expected = expected.astype(object)
+
+        if isinstance(expected, pd.Series):
+            assert result.dtype == df.iloc[:, 0].dtype
+        else:
+            assert all(result.dtypes == df.iloc[:, 0].dtype)
+
+        result = result.astype(object)
+        self.assert_equal(result, expected)
+
     @pytest.mark.parametrize("index", [
         # Two levels, uniform.
         pd.MultiIndex.from_product(([['A', 'B'], ['a', 'b']]),
diff --git a/pandas/tests/extension/json/test_json.py b/pandas/tests/extension/json/test_json.py
--- a/pandas/tests/extension/json/test_json.py
+++ b/pandas/tests/extension/json/test_json.py
@@ -139,6 +139,15 @@ def test_from_dtype(self, data):
 
 
 class TestReshaping(BaseJSON, base.BaseReshapingTests):
+
+    @pytest.mark.skip(reason="Different definitions of NA")
+    def test_stack(self):
+        """
+        The test does .astype(object).stack(). If we happen to have
+        any missing values in `data`, then we'll end up with different
+        rows since we consider `{}` NA, but `.astype(object)` doesn't.
+        """
+
     @pytest.mark.xfail(reason="dict for NA", strict=True)
     def test_unstack(self, data, index):
         # The base test has NaN for the expected NA value.
diff --git a/pandas/tests/frame/test_reshape.py b/pandas/tests/frame/test_reshape.py
--- a/pandas/tests/frame/test_reshape.py
+++ b/pandas/tests/frame/test_reshape.py
@@ -874,6 +874,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, 1])
     def test_unstack_mixed_extension_types(self, level):
         index = pd.MultiIndex.from_tuples([('A', 0), ('A', 1), ('B', 1)],
diff --git a/pandas/tests/sparse/frame/test_frame.py b/pandas/tests/sparse/frame/test_frame.py
--- a/pandas/tests/sparse/frame/test_frame.py
+++ b/pandas/tests/sparse/frame/test_frame.py
@@ -736,6 +736,16 @@ def test_astype_bool(self):
         assert res['A'].dtype == SparseDtype(np.bool)
         assert res['B'].dtype == SparseDtype(np.bool)
 
+    def test_astype_object(self):
+        # This may change in GH-23125
+        df = pd.DataFrame({"A": SparseArray([0, 1]),
+                           "B": SparseArray([0, 1])})
+        result = df.astype(object)
+        dtype = SparseDtype(object, 0)
+        expected = pd.DataFrame({"A": SparseArray([0, 1], dtype=dtype),
+                                 "B": SparseArray([0, 1], dtype=dtype)})
+        tm.assert_frame_equal(result, expected)
+
     def test_fillna(self, float_frame_fill0, float_frame_fill0_dense):
         df = float_frame_fill0.reindex(lrange(5))
         dense = float_frame_fill0_dense.reindex(lrange(5))