Skip to content

Commit

Permalink
fix(python): Handle special case correctly when slicing a LazyFrame (
Browse files Browse the repository at this point in the history
  • Loading branch information
petrosbar authored Mar 26, 2024
1 parent a3cc9e8 commit 2774bbb
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
10 changes: 8 additions & 2 deletions py-polars/polars/slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ def apply(self, s: slice) -> LazyFrame:
# [::k] => gather_every(k),
# [::-1] => reverse(),
# [::-k] => reverse().gather_every(abs(k))
elif start == 0 and s.stop is None:
elif s.start is None and s.stop is None:
if step == 1:
return self.obj.clone()
elif step > 1:
Expand All @@ -168,7 +168,13 @@ def apply(self, s: slice) -> LazyFrame:
elif step < -1:
return self.obj.reverse().gather_every(abs(step))

elif start > 0 > step and s.stop is None:
# ---------------------------------------
# straight-through mappings for "head",
# "reverse" and "gather_every"
# ---------------------------------------
# [i::-1] => head(i+1).reverse()
# [i::k], k<-1 => head(i+1).reverse().gather_every(abs(k))
elif start >= 0 > step and s.stop is None:
obj = self.obj.head(s.start + 1).reverse()
return obj if (abs(step) == 1) else obj.gather_every(abs(step))

Expand Down
3 changes: 3 additions & 0 deletions py-polars/tests/unit/operations/test_slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,10 +88,13 @@ def test_python_slicing_lazy_frame() -> None:
slice(None, 2, 2),
slice(3, None, -1),
slice(1, None, -2),
slice(0, None, -1),
):
# confirm frame slice matches python slice
assert ldf[py_slice].collect().rows() == ldf.collect().rows()[py_slice]

assert_frame_equal(ldf[0::-1], ldf.head(1))
assert_frame_equal(ldf[2::-1], ldf.head(3).reverse())
assert_frame_equal(ldf[::-1], ldf.reverse())
assert_frame_equal(ldf[::-2], ldf.reverse().gather_every(2))

Expand Down

0 comments on commit 2774bbb

Please sign in to comment.