Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 971d091

Browse files
milkshakeiiiHenry J Solberg
and
Henry J Solberg
authoredOct 3, 2023
fix: fix df/series.iloc by list with multiindex (#79)
* fix: fix df/series.iloc by list with multiindex * fix git error --------- Co-authored-by: Henry J Solberg <henryjsolberg@google.com>
1 parent 4a84714 commit 971d091

File tree

2 files changed

+26
-7
lines changed

2 files changed

+26
-7
lines changed
 

‎bigframes/core/indexers.py

+8-7
Original file line numberDiff line numberDiff line change
@@ -332,8 +332,6 @@ def _iloc_getitem_series_or_dataframe(
332332
elif isinstance(key, slice):
333333
return series_or_dataframe._slice(key.start, key.stop, key.step)
334334
elif pd.api.types.is_list_like(key):
335-
# TODO(henryjsolberg): support MultiIndex
336-
337335
if len(key) == 0:
338336
return typing.cast(
339337
typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
@@ -346,15 +344,18 @@ def _iloc_getitem_series_or_dataframe(
346344
original_series_name if original_series_name is not None else "0"
347345
)
348346
df = series_or_dataframe.to_frame()
349-
original_index_name = df.index.name
350-
temporary_index_name = guid.generate_guid(prefix="temp_iloc_index_")
351-
df = df.rename_axis(temporary_index_name)
347+
original_index_names = df.index.names
348+
temporary_index_names = [
349+
guid.generate_guid(prefix="temp_iloc_index_")
350+
for _ in range(len(df.index.names))
351+
]
352+
df = df.rename_axis(temporary_index_names)
352353

353354
# set to offset index and use regular loc, then restore index
354355
df = df.reset_index(drop=False)
355356
result = df.loc[key]
356-
result = result.set_index(temporary_index_name)
357-
result = result.rename_axis(original_index_name)
357+
result = result.set_index(temporary_index_names)
358+
result = result.rename_axis(original_index_names)
358359

359360
if isinstance(series_or_dataframe, bigframes.series.Series):
360361
result = result[series_name]

‎tests/system/small/test_dataframe.py

+18
Original file line numberDiff line numberDiff line change
@@ -2525,6 +2525,24 @@ def test_iloc_list(scalars_df_index, scalars_pandas_df_index):
25252525
)
25262526

25272527

2528+
def test_iloc_list_multiindex(scalars_dfs):
2529+
scalars_df, scalars_pandas_df = scalars_dfs
2530+
scalars_df = scalars_df.copy()
2531+
scalars_pandas_df = scalars_pandas_df.copy()
2532+
scalars_df = scalars_df.set_index(["bytes_col", "numeric_col"])
2533+
scalars_pandas_df = scalars_pandas_df.set_index(["bytes_col", "numeric_col"])
2534+
2535+
index_list = [0, 0, 0, 5, 4, 7]
2536+
2537+
bf_result = scalars_df.iloc[index_list]
2538+
pd_result = scalars_pandas_df.iloc[index_list]
2539+
2540+
pd.testing.assert_frame_equal(
2541+
bf_result.to_pandas(),
2542+
pd_result,
2543+
)
2544+
2545+
25282546
def test_iloc_empty_list(scalars_df_index, scalars_pandas_df_index):
25292547
index_list = []
25302548

0 commit comments

Comments
 (0)
Failed to load comments.