BUG: Fix pyarrow categoricals not working for pivot and multiindex

doc/source/whatsnew/v3.0.0.rst

@@ -636,8 +636,8 @@ Bug fixes
 Categorical
 ^^^^^^^^^^^
 - Bug in :func:`Series.apply` where ``nan`` was ignored for :class:`CategoricalDtype` (:issue:`59938`)
+- Bug in :meth:`DataFrame.pivot` and :meth:`DataFrame.set_index` raising an ``ArrowNotImplementedError`` for columns with pyarrow dictionary dtype (:issue:`53051`)
 - Bug in :meth:`Series.convert_dtypes` with ``dtype_backend="pyarrow"`` where empty :class:`CategoricalDtype` :class:`Series` raised an error or got converted to ``null[pyarrow]`` (:issue:`59934`)
--
 
 Datetimelike
 ^^^^^^^^^^^^

pandas/core/arrays/categorical.py

@@ -452,7 +452,7 @@ def __init__(
                 if isinstance(values, Index):
                     arr = values._data._pa_array.combine_chunks()
                 else:
-                    arr = values._pa_array.combine_chunks()
+                    arr = extract_array(values)._pa_array.combine_chunks()
                 categories = arr.dictionary.to_pandas(types_mapper=ArrowDtype)
                 codes = arr.indices.to_numpy()
                 dtype = CategoricalDtype(categories, values.dtype.pyarrow_dtype.ordered)

pandas/tests/reshape/test_pivot.py

@@ -3,6 +3,7 @@
     datetime,
     timedelta,
 )
+import io
 from itertools import product
 import re
 
@@ -2851,3 +2852,43 @@ def test_pivot_margins_with_none_index(self):
             ),
         )
         tm.assert_frame_equal(result, expected)
+
+    # Ignore deprecation raised by old versions of pyarrow. Already fixed in
+    # newer versions
+    @pytest.mark.filterwarnings("ignore:Passing a BlockManager:DeprecationWarning")
+    def test_pivot_with_pyarrow_categorical(self):
+        # GH#53051
+
+        pytest.importorskip("pyarrow")
+
+        # Create dataframe with categorical column
+        df = (
+            DataFrame(
+                [("A", 1), ("B", 2), ("C", 3)],
+                columns=["string_column", "number_column"],
+            )
+            .astype({"string_column": "string", "number_column": "float32"})
+            .astype({"string_column": "category", "number_column": "float32"})
+        )
+
+        # Convert dataframe to pyarrow backend
+        with io.BytesIO() as buffer:
+            df.to_parquet(buffer)
+            buffer.seek(0)  # Reset buffer position
+            df = pd.read_parquet(buffer, dtype_backend="pyarrow")
+
+        # Check that pivot works
+        df = df.pivot(columns=["string_column"], values=["number_column"])
+
+        # Assert that values of result are correct to prevent silent failure
+        multi_index = MultiIndex.from_arrays(
+            [["number_column", "number_column", "number_column"], ["A", "B", "C"]],
+            names=(None, "string_column"),
+        )
+        df_expected = DataFrame(
+            [[1.0, np.nan, np.nan], [np.nan, 2.0, np.nan], [np.nan, np.nan, 3.0]],
+            columns=multi_index,
+        )
+        tm.assert_frame_equal(
+            df, df_expected, check_dtype=False, check_column_type=False
+        )

pandas/tests/test_multilevel.py

@@ -1,4 +1,5 @@
 import datetime
+import io
 
 import numpy as np
 import pytest
@@ -318,6 +319,33 @@ def test_multiindex_dt_with_nan(self):
         expected = Series(["a", "b", "c", "d"], name=("sub", np.nan))
         tm.assert_series_equal(result, expected)
 
+    # Ignore deprecation raised by old versions of pyarrow. Already fixed in
+    # newer versions
+    @pytest.mark.filterwarnings("ignore:Passing a BlockManager:DeprecationWarning")
+    def test_multiindex_with_pyarrow_categorical(self):
+        # GH#53051
+
+        pytest.importorskip("pyarrow")
+
+        # Create dataframe with categorical column
+        df = (
+            DataFrame(
+                [["A", 1], ["B", 2], ["C", 3]],
+                columns=["string_column", "number_column"],
+            )
+            .astype({"string_column": "string", "number_column": "float32"})
+            .astype({"string_column": "category", "number_column": "float32"})
+        )
+
+        # Convert dataframe to pyarrow backend
+        with io.BytesIO() as buffer:
+            df.to_parquet(buffer)
+            buffer.seek(0)  # Reset buffer position
+            df = pd.read_parquet(buffer, dtype_backend="pyarrow")
+
+        # Check that index can be set
+        df.set_index(["string_column", "number_column"])
+
 
 class TestSorted:
     """everything you wanted to test about sorting"""