Skip to content

Commit

Permalink
fix(python): Fix DataFrame.__getitem__ for empty list input - `df[[…
Browse files Browse the repository at this point in the history
…]]` (pola-rs#16520)
  • Loading branch information
stinodego authored May 27, 2024
1 parent cc2c905 commit 7aa514c
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 11 deletions.
13 changes: 5 additions & 8 deletions py-polars/polars/_utils/getitem.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,20 +145,17 @@ def get_df_item_by_key(
else:
return _select_rows(selection, row_key)

# Single input, e.g. df[1]
elif isinstance(key, str):
# Single string input, e.g. df["a"]
if isinstance(key, str):
# This case is required because empty strings are otherwise treated
# as an empty Sequence in `_select_rows`
return df.get_column(key)
elif isinstance(key, Sequence) and len(key) == 0:
# df[[]]
# TODO: This removes all columns, but it should remove all rows.
# https://github.com/pola-rs/polars/issues/4924
return df.__class__()

# Single input - df[1] - or multiple inputs - df["a", "b", "c"]
try:
return _select_rows(df, key) # type: ignore[arg-type]
except TypeError:
return _select_columns(df, key) # type: ignore[arg-type]
return _select_columns(df, key)


# `str` overlaps with `Sequence[str]`
Expand Down
10 changes: 7 additions & 3 deletions py-polars/tests/unit/dataframe/test_getitem.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,13 @@ def test_df_getitem_row_range_single_input() -> None:
assert_frame_equal(result, expected)


def test_df_getitem_row_empty_list_single_input() -> None:
df = pl.DataFrame({"a": [1, 2], "b": [5.0, 6.0]})
result = df[[]]
expected = df.clear()
assert_frame_equal(result, expected)


def test_df_getitem() -> None:
"""Test all the methods to use [] on a dataframe."""
df = pl.DataFrame({"a": [1.0, 2.0, 3.0, 4.0], "b": [3, 4, 5, 6]})
Expand Down Expand Up @@ -287,9 +294,6 @@ def test_df_getitem() -> None:
# empty list with column selector drops rows but keeps columns
assert_frame_equal(df[empty, :], df[:0])

# empty list without column select return empty frame
assert_frame_equal(df[empty], pl.DataFrame({}))

# numpy array: assumed to be row indices if integers, or columns if strings

# numpy array: positive idxs and empty idx
Expand Down

0 comments on commit 7aa514c

Please sign in to comment.