diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 101e7a1c5e..9cb388329e 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -576,6 +576,9 @@ def __getitem__( if isinstance(key, bigframes.series.Series): return self._getitem_bool_series(key) + if isinstance(key, slice): + return self.iloc[key] + if isinstance(key, typing.Hashable): return self._getitem_label(key) # Select a subset of columns or re-order columns. diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 452fa841e4..362d736aeb 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -221,6 +221,21 @@ def test_get_column_nonstring(scalars_dfs): assert_series_equal(bf_result, pd_result) +@pytest.mark.parametrize( + "row_slice", + [ + (slice(1, 7, 2)), + (slice(1, 7, None)), + (slice(None, -3, None)), + ], +) +def test_get_rows_with_slice(scalars_dfs, row_slice): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = scalars_df[row_slice].to_pandas() + pd_result = scalars_pandas_df[row_slice] + assert_pandas_df_equal(bf_result, pd_result) + + def test_hasattr(scalars_dfs): scalars_df, _ = scalars_dfs assert hasattr(scalars_df, "int64_col")