BUG: preserve categorical & sparse types when grouping / pivot & preserve dtypes on ufuncs
State: closed

closes #18502
closes #23743

Labels:
name: Bug
name: Groupby
name: Sparse
description: Sparse Data Type
name: Categorical
description: Categorical Data Type
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.

Comments:
# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=h1) Report
> Merging [#26550](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/7629a182f34d0374cb329e07076fba16605853ee?src=pr&el=desc) will **decrease** coverage by `0.02%`.
> The diff coverage is `88.88%`.

[![Impacted file tree graph](https://codecov.io/gh/pandas-dev/pandas/pull/26550/graphs/tree.svg?width=650&token=eZ4WkYLtcO&height=150&src=pr)](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master   #26550      +/-   ##
==========================================
- Coverage   91.77%   91.74%   -0.03%     
==========================================
  Files         174      174              
  Lines       50642    50666      +24     
==========================================
+ Hits        46476    46484       +8     
- Misses       4166     4182      +16
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.28% <88.88%> (-0.02%)` | :arrow_down: |
| #single | `41.66% <44.44%> (-0.12%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/groupby/generic.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9nZW5lcmljLnB5) | `88.07% <100%> (-0.89%)` | :arrow_down: |
| [pandas/core/series.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc2VyaWVzLnB5) | `93.63% <100%> (+0.01%)` | :arrow_up: |
| [pandas/core/frame.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZnJhbWUucHk=) | `97.02% <100%> (-0.11%)` | :arrow_down: |
| [pandas/core/groupby/ops.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9vcHMucHk=) | `96% <100%> (ø)` | :arrow_up: |
| [pandas/core/nanops.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvbmFub3BzLnB5) | `94.11% <100%> (ø)` | :arrow_up: |
| [pandas/core/groupby/groupby.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9ncm91cGJ5LnB5) | `97.39% <100%> (+0.16%)` | :arrow_up: |
| [pandas/core/dtypes/cast.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Nhc3QucHk=) | `91.54% <100%> (ø)` | :arrow_up: |
| [pandas/core/generic.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ2VuZXJpYy5weQ==) | `93.42% <25%> (-0.15%)` | :arrow_down: |
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `93.94% <71.42%> (-0.19%)` | :arrow_down: |
| [pandas/io/gbq.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2lvL2dicS5weQ==) | `78.94% <0%> (-10.53%)` | :arrow_down: |
| ... and [6 more](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=footer). Last update [7629a18...f397908](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=h1) Report
> Merging [#26550](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/b9b081dc6b510c8290ded12fe751b1216843527e?src=pr&el=desc) will **decrease** coverage by `0.02%`.
> The diff coverage is `92.3%`.

[![Impacted file tree graph](https://codecov.io/gh/pandas-dev/pandas/pull/26550/graphs/tree.svg?width=650&token=eZ4WkYLtcO&height=150&src=pr)](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master   #26550      +/-   ##
==========================================
- Coverage   91.97%   91.95%   -0.03%     
==========================================
  Files         180      180              
  Lines       50756    50789      +33     
==========================================
+ Hits        46685    46705      +20     
- Misses       4071     4084      +13
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.55% <92.3%> (-0.02%)` | :arrow_down: |
| #single | `41.81% <41.53%> (-0.1%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/sparse/frame.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc3BhcnNlL2ZyYW1lLnB5) | `95.67% <100%> (+0.03%)` | :arrow_up: |
| [pandas/core/groupby/generic.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9nZW5lcmljLnB5) | `88.48% <100%> (-0.86%)` | :arrow_down: |
| [pandas/core/series.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc2VyaWVzLnB5) | `93.69% <100%> (+0.02%)` | :arrow_up: |
| [pandas/core/dtypes/cast.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Nhc3QucHk=) | `90.53% <100%> (ø)` | :arrow_up: |
| [pandas/core/nanops.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvbmFub3BzLnB5) | `94.76% <100%> (ø)` | :arrow_up: |
| [pandas/core/frame.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZnJhbWUucHk=) | `96.9% <100%> (-0.11%)` | :arrow_down: |
| [pandas/core/groupby/ops.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9vcHMucHk=) | `96% <100%> (ø)` | :arrow_up: |
| [pandas/core/groupby/groupby.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ3JvdXBieS9ncm91cGJ5LnB5) | `97.46% <100%> (+0.29%)` | :arrow_up: |
| [pandas/core/arrays/sparse.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3NwYXJzZS5weQ==) | `93.88% <100%> (+0.15%)` | :arrow_up: |
| [pandas/core/internals/construction.py](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2NvbnN0cnVjdGlvbi5weQ==) | `96.19% <100%> (+0.25%)` | :arrow_up: |
| ... and [9 more](https://codecov.io/gh/pandas-dev/pandas/pull/26550/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=footer). Last update [b9b081d...4bd486e](https://codecov.io/gh/pandas-dev/pandas/pull/26550?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

@jreback rough sketch of the idea in https://github.com/pandas-dev/pandas/pull/26550#discussion_r289198328


```diff
diff --git a/pandas/core/generic.py b/pandas/core/generic.py
index ce1cb37ab..898d03e6c 100644
--- a/pandas/core/generic.py
+++ b/pandas/core/generic.py
@@ -1945,6 +1945,10 @@ class NDFrame(PandasObject, SelectionMixin):
     # GH#23114 Ensure ndarray.__op__(DataFrame) returns NotImplemented
     __array_priority__ = 1000
 
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        result = self._data.map_blocks(ufunc, **kwargs)
+        return type(self)(result)
+
     def __array__(self, dtype=None):
         return com.values_from_object(self)
 
diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py
index 0b63588c9..9c7cc4bd5 100644
--- a/pandas/core/internals/managers.py
+++ b/pandas/core/internals/managers.py
@@ -117,6 +117,10 @@ class BlockManager(PandasObject):
 
         self._rebuild_blknos_and_blklocs()
 
+    def map_blocks(self, func, **kwargs):
+        newbs = [make_block(func(block.values, **kwargs), block.mgr_locs) for block in self.blocks]
+        return type(self)(newbs, self.axes)
+
     def make_empty(self, axes=None):
         """ return an empty BlockManager with the items axis of len 0 """
         if axes is None:
```

usage:

```python
In [1]: import numpy as np; import pandas as pd

In [2]: df = pd.DataFrame({"A": [1, 0]}, dtype=pd.SparseDtype('int'))

In [3]: np.sin(df)
Out[3]:
          A
0  0.841471
1  0.000000

In [4]: np.sin(df).dtypes
Out[4]:
A    Sparse[float64, 0.0]
dtype: object
```

The `__array_ufunc__` isn't really correct yet, but does make sense? I think that's much cleaner than doing the op and trying to cast back.


On this branch, the casting seems problematic:

```python
In [10]: df = pd.DataFrame({"A": [1, 2, 3]}, dtype=pd.SparseDtype('int'))

In [11]: np.sqrt(df)
Out[11]:
   A
0  1
1  1
2  1
```
yeah i can do via blocks 


array_ufunc is orthogonal here - i don’t think we have any implementation yet (should though)
fixed up the casting, a bit tricky.
push should make this pass, will address comments soon.
replacing with a PR just with the categorical grouping changes; the sparse ufuncs is conflicting with what @TomAugspurger is doing and is messy.