Skip to content

Commit 563f0cb

Browse files
authored
feat: DataFrames.__getitem__ support for slice input (#1668)
1 parent 4fb54df commit 563f0cb

File tree

2 files changed

+18
-0
lines changed

2 files changed

+18
-0
lines changed

bigframes/dataframe.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -576,6 +576,9 @@ def __getitem__(
576576
if isinstance(key, bigframes.series.Series):
577577
return self._getitem_bool_series(key)
578578

579+
if isinstance(key, slice):
580+
return self.iloc[key]
581+
579582
if isinstance(key, typing.Hashable):
580583
return self._getitem_label(key)
581584
# Select a subset of columns or re-order columns.

tests/system/small/test_dataframe.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,21 @@ def test_get_column_nonstring(scalars_dfs):
221221
assert_series_equal(bf_result, pd_result)
222222

223223

224+
@pytest.mark.parametrize(
225+
"row_slice",
226+
[
227+
(slice(1, 7, 2)),
228+
(slice(1, 7, None)),
229+
(slice(None, -3, None)),
230+
],
231+
)
232+
def test_get_rows_with_slice(scalars_dfs, row_slice):
233+
scalars_df, scalars_pandas_df = scalars_dfs
234+
bf_result = scalars_df[row_slice].to_pandas()
235+
pd_result = scalars_pandas_df[row_slice]
236+
assert_pandas_df_equal(bf_result, pd_result)
237+
238+
224239
def test_hasattr(scalars_dfs):
225240
scalars_df, _ = scalars_dfs
226241
assert hasattr(scalars_df, "int64_col")

0 commit comments

Comments
 (0)