diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py index 2e53e694..b81eaf47 100644 --- a/dask_expr/_collection.py +++ b/dask_expr/_collection.py @@ -154,6 +154,8 @@ def __reduce__(self): def __getitem__(self, other): if isinstance(other, FrameBase): return new_collection(self.expr.__getitem__(other.expr)) + elif isinstance(other, slice): + return self.loc[other] return new_collection(self.expr.__getitem__(other)) def persist(self, fuse=True, combine_similar=True, **kwargs): diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py index 58503e5f..bd62117d 100644 --- a/dask_expr/_expr.py +++ b/dask_expr/_expr.py @@ -1991,7 +1991,7 @@ def _meta(self): if is_dataframe_like(self.frame._meta): return super()._meta # if we are not a DataFrame and have a scalar, we reduce to a scalar - if not isinstance(self.operand("columns"), list) and not hasattr( + if not isinstance(self.operand("columns"), (list, slice)) and not hasattr( self.operand("columns"), "dtype" ): return meta_nonempty(self.frame._meta).iloc[0] diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py index 74eaa94e..3b905a2f 100644 --- a/dask_expr/tests/test_collection.py +++ b/dask_expr/tests/test_collection.py @@ -1052,6 +1052,11 @@ def test_dropna_simplify(pdf, subset): assert_eq(q, pdf.dropna(subset=subset)["y"]) +def test_series_slice_getitem(df, pdf): + with pytest.raises(NotImplementedError): + df.x[:4] + + def test_dir(df): assert all(c in dir(df) for c in df.columns) assert "sum" in dir(df)