CLN: Refactor code and clean up according to PR feedback

Author: robin-mader-bis

doc/source/whatsnew/v3.0.0.rst

@@ -639,6 +639,7 @@ 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/tests/reshape/test_pivot.py

@@ -3,7 +3,6 @@
     datetime,
     timedelta,
 )
-import io
 from itertools import product
 import re
 
@@ -16,6 +15,7 @@
 
 import pandas as pd
 from pandas import (
+    ArrowDtype,
     Categorical,
     DataFrame,
     Grouper,
@@ -2853,29 +2853,24 @@ 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")
+        pa = 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"})
-        )
+        df = DataFrame(
+            {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
+        ).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")
+        df = df.astype(
+            {
+                "string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
+                "number_column": "float[pyarrow]",
+            }
+        )
 
         # Check that pivot works
         df = df.pivot(columns=["string_column"], values=["number_column"])

pandas/tests/test_multilevel.py

@@ -1,11 +1,11 @@
 import datetime
-import io
 
 import numpy as np
 import pytest
 
 import pandas as pd
 from pandas import (
+    ArrowDtype,
     DataFrame,
     MultiIndex,
     Series,
@@ -325,23 +325,20 @@ def test_multiindex_dt_with_nan(self):
     def test_multiindex_with_pyarrow_categorical(self):
         # GH#53051
 
-        pytest.importorskip("pyarrow")
+        pa = 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"})
-        )
+        df = DataFrame(
+            {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
+        ).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")
+        df = df.astype(
+            {
+                "string_column": ArrowDtype(pa.dictionary(pa.int32(), pa.string())),
+                "number_column": "float[pyarrow]",
+            }
+        )
 
         # Check that index can be set
         df.set_index(["string_column", "number_column"])