From ec39570c799814bae758b7f7ab403cd49d47fcaf Mon Sep 17 00:00:00 2001 From: Garrett Wu Date: Wed, 3 Jan 2024 22:47:32 +0000 Subject: [PATCH 1/2] fix: DF.drop tuple input as multi-index --- bigframes/dataframe.py | 31 +++++++++++++++---- tests/system/small/test_multiindex.py | 2 ++ .../bigframes_vendored/pandas/core/frame.py | 2 +- 3 files changed, 28 insertions(+), 7 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 595670b0b6..7f72e6f959 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1062,12 +1062,31 @@ def drop( level_id = self._resolve_levels(level or 0)[0] if utils.is_list_like(index): - block, inverse_condition_id = block.apply_unary_op( - level_id, ops.IsInOp(index, match_nulls=True) - ) - block, condition_id = block.apply_unary_op( - inverse_condition_id, ops.invert_op - ) + # Only tuple is treated as multi-index value combinations + if isinstance(index, tuple): + if level is not None: + raise ValueError("Multi-index tuple can't specify level.") + condition_id = None + for i, idx in enumerate(index): + level_id = self._resolve_levels(i)[0] + block, condition_id_cur = block.apply_unary_op( + level_id, ops.partial_right(ops.ne_op, idx) + ) + if condition_id: + block, condition_id = block.apply_binary_op( + condition_id, condition_id_cur, ops.or_op + ) + else: + condition_id = condition_id_cur + + condition_id = typing.cast(str, condition_id) + else: + block, inverse_condition_id = block.apply_unary_op( + level_id, ops.IsInOp(index, match_nulls=True) + ) + block, condition_id = block.apply_unary_op( + inverse_condition_id, ops.invert_op + ) elif isinstance(index, indexes.Index): return self._drop_by_index(index) else: diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 2d4e1f0204..6af89e15c3 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -204,6 +204,7 @@ def test_series_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index, (1, 0), ([0, 1], 0), ([True, None], 1), + ((0, True), None), ], ) def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, level): @@ -213,6 +214,7 @@ def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, lev bf_result = bf_frame.drop(labels=labels, axis="index", level=level).to_pandas() pd_result = pd_frame.drop(labels=labels, axis="index", level=level) + print("over") pandas.testing.assert_frame_equal(bf_result, pd_result) diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index f2de8fcb6a..271c955220 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -1086,7 +1086,7 @@ def drop( Args: labels: - Index or column labels to drop. + Index or column labels to drop. A tuple will be used as a single label and not treated as a list-like. axis: Whether to drop labels from the index (0 or 'index') or columns (1 or 'columns'). From 452d4c0704837dce9b91ac5575a0ffe6b01eb22a Mon Sep 17 00:00:00 2001 From: Garrett Wu Date: Thu, 4 Jan 2024 19:21:35 +0000 Subject: [PATCH 2/2] fix test --- tests/system/small/test_multiindex.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/system/small/test_multiindex.py b/tests/system/small/test_multiindex.py index 6af89e15c3..4b1c3209bf 100644 --- a/tests/system/small/test_multiindex.py +++ b/tests/system/small/test_multiindex.py @@ -214,7 +214,6 @@ def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, lev bf_result = bf_frame.drop(labels=labels, axis="index", level=level).to_pandas() pd_result = pd_frame.drop(labels=labels, axis="index", level=level) - print("over") pandas.testing.assert_frame_equal(bf_result, pd_result)