Skip to content

Commit 21391a9

Browse files
authored
fix: DF.drop tuple input as multi-index (#301)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes b/317131784
1 parent 9c21323 commit 21391a9

File tree

3 files changed

+27
-7
lines changed

3 files changed

+27
-7
lines changed

bigframes/dataframe.py

+25-6
Original file line numberDiff line numberDiff line change
@@ -1062,12 +1062,31 @@ def drop(
10621062
level_id = self._resolve_levels(level or 0)[0]
10631063

10641064
if utils.is_list_like(index):
1065-
block, inverse_condition_id = block.apply_unary_op(
1066-
level_id, ops.IsInOp(index, match_nulls=True)
1067-
)
1068-
block, condition_id = block.apply_unary_op(
1069-
inverse_condition_id, ops.invert_op
1070-
)
1065+
# Only tuple is treated as multi-index value combinations
1066+
if isinstance(index, tuple):
1067+
if level is not None:
1068+
raise ValueError("Multi-index tuple can't specify level.")
1069+
condition_id = None
1070+
for i, idx in enumerate(index):
1071+
level_id = self._resolve_levels(i)[0]
1072+
block, condition_id_cur = block.apply_unary_op(
1073+
level_id, ops.partial_right(ops.ne_op, idx)
1074+
)
1075+
if condition_id:
1076+
block, condition_id = block.apply_binary_op(
1077+
condition_id, condition_id_cur, ops.or_op
1078+
)
1079+
else:
1080+
condition_id = condition_id_cur
1081+
1082+
condition_id = typing.cast(str, condition_id)
1083+
else:
1084+
block, inverse_condition_id = block.apply_unary_op(
1085+
level_id, ops.IsInOp(index, match_nulls=True)
1086+
)
1087+
block, condition_id = block.apply_unary_op(
1088+
inverse_condition_id, ops.invert_op
1089+
)
10711090
elif isinstance(index, indexes.Index):
10721091
return self._drop_by_index(index)
10731092
else:

tests/system/small/test_multiindex.py

+1
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,7 @@ def test_series_multi_index_droplevel(scalars_df_index, scalars_pandas_df_index,
234234
(1, 0),
235235
([0, 1], 0),
236236
([True, None], 1),
237+
((0, True), None),
237238
],
238239
)
239240
def test_multi_index_drop(scalars_df_index, scalars_pandas_df_index, labels, level):

third_party/bigframes_vendored/pandas/core/frame.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1107,7 +1107,7 @@ def drop(
11071107
11081108
Args:
11091109
labels:
1110-
Index or column labels to drop.
1110+
Index or column labels to drop. A tuple will be used as a single label and not treated as a list-like.
11111111
axis:
11121112
Whether to drop labels from the index (0 or 'index') or
11131113
columns (1 or 'columns').

0 commit comments

Comments
 (0)