Skip to content

Commit 7aa514c

Browse files
authored
fix(python): Fix DataFrame.__getitem__ for empty list input - df[[]] (#16520)
1 parent cc2c905 commit 7aa514c

File tree

2 files changed

+12
-11
lines changed

2 files changed

+12
-11
lines changed

py-polars/polars/_utils/getitem.py

+5-8
Original file line numberDiff line numberDiff line change
@@ -145,20 +145,17 @@ def get_df_item_by_key(
145145
else:
146146
return _select_rows(selection, row_key)
147147

148-
# Single input, e.g. df[1]
149-
elif isinstance(key, str):
148+
# Single string input, e.g. df["a"]
149+
if isinstance(key, str):
150150
# This case is required because empty strings are otherwise treated
151151
# as an empty Sequence in `_select_rows`
152152
return df.get_column(key)
153-
elif isinstance(key, Sequence) and len(key) == 0:
154-
# df[[]]
155-
# TODO: This removes all columns, but it should remove all rows.
156-
# https://github.com/pola-rs/polars/issues/4924
157-
return df.__class__()
153+
154+
# Single input - df[1] - or multiple inputs - df["a", "b", "c"]
158155
try:
159156
return _select_rows(df, key) # type: ignore[arg-type]
160157
except TypeError:
161-
return _select_columns(df, key) # type: ignore[arg-type]
158+
return _select_columns(df, key)
162159

163160

164161
# `str` overlaps with `Sequence[str]`

py-polars/tests/unit/dataframe/test_getitem.py

+7-3
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,13 @@ def test_df_getitem_row_range_single_input() -> None:
234234
assert_frame_equal(result, expected)
235235

236236

237+
def test_df_getitem_row_empty_list_single_input() -> None:
238+
df = pl.DataFrame({"a": [1, 2], "b": [5.0, 6.0]})
239+
result = df[[]]
240+
expected = df.clear()
241+
assert_frame_equal(result, expected)
242+
243+
237244
def test_df_getitem() -> None:
238245
"""Test all the methods to use [] on a dataframe."""
239246
df = pl.DataFrame({"a": [1.0, 2.0, 3.0, 4.0], "b": [3, 4, 5, 6]})
@@ -287,9 +294,6 @@ def test_df_getitem() -> None:
287294
# empty list with column selector drops rows but keeps columns
288295
assert_frame_equal(df[empty, :], df[:0])
289296

290-
# empty list without column select return empty frame
291-
assert_frame_equal(df[empty], pl.DataFrame({}))
292-
293297
# numpy array: assumed to be row indices if integers, or columns if strings
294298

295299
# numpy array: positive idxs and empty idx

0 commit comments

Comments
 (0)