From b6ab35ad9ecf0f426e8e456f62e923efcb1aec65 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Sun, 17 Dec 2023 00:24:46 +0100 Subject: [PATCH 1/2] BUG: pivot dropping wrong column level with numeric columns and ea dtype --- doc/source/whatsnew/v2.2.0.rst | 1 + pandas/core/reshape/reshape.py | 2 +- pandas/tests/reshape/test_pivot.py | 5 +++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/doc/source/whatsnew/v2.2.0.rst b/doc/source/whatsnew/v2.2.0.rst index 0c4fb6d3d1164..08ef9ead94efb 100644 --- a/doc/source/whatsnew/v2.2.0.rst +++ b/doc/source/whatsnew/v2.2.0.rst @@ -675,6 +675,7 @@ Reshaping - Bug in :meth:`DataFrame.melt` where an exception was raised if ``var_name`` was not a string (:issue:`55948`) - Bug in :meth:`DataFrame.melt` where it would not preserve the datetime (:issue:`55254`) - Bug in :meth:`DataFrame.pivot_table` where the row margin is incorrect when the columns have numeric names (:issue:`26568`) +- Bug in :meth:`DataFrame.pivot` with numeric columns and extension dtype for data (:issue:`56528`) Sparse ^^^^^^ diff --git a/pandas/core/reshape/reshape.py b/pandas/core/reshape/reshape.py index d6922ba58d2b9..8c822ec58e011 100644 --- a/pandas/core/reshape/reshape.py +++ b/pandas/core/reshape/reshape.py @@ -572,7 +572,7 @@ def _unstack_extension_series( # equiv: result.droplevel(level=0, axis=1) # but this avoids an extra copy - result.columns = result.columns.droplevel(0) + result.columns = result.columns._drop_level_numbers([0]) return result diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py index dab2b034d3fd4..1f84d525a0a89 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -2520,8 +2520,9 @@ def test_pivot_empty(self): expected = DataFrame(index=[], columns=[]) tm.assert_frame_equal(result, expected, check_names=False) - def test_pivot_integer_bug(self): - df = DataFrame(data=[("A", "1", "A1"), ("B", "2", "B2")]) + @pytest.mark.parametrize("dtype", [object, "string"]) + def test_pivot_integer_bug(self, dtype): + df = DataFrame(data=[("A", "1", "A1"), ("B", "2", "B2")], dtype=dtype) result = df.pivot(index=1, columns=0, values=2) tm.assert_index_equal(result.columns, Index(["A", "B"], name=0)) From d75358a965f936b4bae920157b999e1a06a9915f Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Sun, 17 Dec 2023 18:48:09 +0100 Subject: [PATCH 2/2] Fixup --- pandas/tests/reshape/test_pivot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py index 1f84d525a0a89..d9ae8ddfb930e 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -2525,7 +2525,7 @@ def test_pivot_integer_bug(self, dtype): df = DataFrame(data=[("A", "1", "A1"), ("B", "2", "B2")], dtype=dtype) result = df.pivot(index=1, columns=0, values=2) - tm.assert_index_equal(result.columns, Index(["A", "B"], name=0)) + tm.assert_index_equal(result.columns, Index(["A", "B"], name=0, dtype=dtype)) def test_pivot_index_none(self): # GH#3962