From 7692044890d5e51cffe835970c4194f15a429b2f Mon Sep 17 00:00:00 2001 From: Henry J Solberg Date: Tue, 7 Nov 2023 20:00:24 +0000 Subject: [PATCH] fix: for reset_index on unnamed multiindex, always use level_[n] label --- bigframes/core/blocks.py | 2 +- tests/system/small/test_dataframe.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 635e7db865..e831b42752 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -282,7 +282,7 @@ def reset_index(self, drop: bool = True) -> Block: column_labels_modified = self.column_labels for level, label in enumerate(index_labels): if label is None: - if "index" not in self.column_labels: + if "index" not in self.column_labels and len(index_labels) <= 1: label = "index" else: label = f"level_{level}" diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index bd5930e508..6caf1226a2 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1207,6 +1207,28 @@ def test_reset_index_with_unnamed_index( pandas.testing.assert_frame_equal(bf_result, pd_result) +def test_reset_index_with_unnamed_multiindex( + scalars_df_index, + scalars_pandas_df_index, +): + bf_df = dataframe.DataFrame( + ([1, 2, 3], [2, 5, 7]), + index=pd.MultiIndex.from_tuples([("a", "aa"), ("a", "aa")]), + ) + pd_df = pd.DataFrame( + ([1, 2, 3], [2, 5, 7]), + index=pd.MultiIndex.from_tuples([("a", "aa"), ("a", "aa")]), + ) + + bf_df = bf_df.reset_index() + pd_df = pd_df.reset_index() + + assert pd_df.columns[0] == "level_0" + assert bf_df.columns[0] == "level_0" + assert pd_df.columns[1] == "level_1" + assert bf_df.columns[1] == "level_1" + + def test_reset_index_with_unnamed_index_and_index_column( scalars_df_index, scalars_pandas_df_index,