polars-runtime-compat 1.34.0b2__cp39-abi3-macosx_11_0_arm64.whl

polars/_utils/slice.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Union
+
+import polars._reexport as pl
+
+if TYPE_CHECKING:
+    from polars import DataFrame, LazyFrame, Series
+
+    FrameOrSeries = Union["DataFrame", "Series"]
+
+
+class PolarsSlice:
+    """
+    Apply Python slice object to Polars DataFrame or Series.
+
+    Has full support for negative indexing and/or stride.
+    """
+
+    stop: int
+    start: int
+    stride: int
+    slice_length: int
+    is_unbounded: bool
+    obj: FrameOrSeries
+
+    def __init__(self, obj: FrameOrSeries) -> None:
+        self.obj = obj
+
+    @staticmethod
+    def _as_original(lazy: LazyFrame, original: FrameOrSeries) -> FrameOrSeries:
+        """Return lazy variant back to its original type."""
+        frame = lazy.collect()
+        return frame if isinstance(original, pl.DataFrame) else frame.to_series()
+
+    @staticmethod
+    def _lazify(obj: FrameOrSeries) -> LazyFrame:
+        """Make lazy to ensure efficient/consistent handling."""
+        return obj.to_frame().lazy() if isinstance(obj, pl.Series) else obj.lazy()
+
+    def _slice_positive(self, obj: LazyFrame) -> LazyFrame:
+        """Logic for slices with positive stride."""
+        # note: at this point stride is guaranteed to be > 1
+        return obj.slice(self.start, self.slice_length).gather_every(self.stride)
+
+    def _slice_negative(self, obj: LazyFrame) -> LazyFrame:
+        """Logic for slices with negative stride."""
+        stride = abs(self.stride)
+        lazyslice = obj.slice(self.stop + 1, self.slice_length).reverse()
+        return lazyslice.gather_every(stride) if (stride > 1) else lazyslice
+
+    def _slice_setup(self, s: slice) -> None:
+        """Normalise slice bounds, identify unbounded and/or zero-length slices."""
+        # can normalise slice indices as we know object size
+        obj_len = len(self.obj)
+        start, stop, stride = slice(s.start, s.stop, s.step).indices(obj_len)
+
+        # check if slice is actually unbounded
+        if stride >= 1:
+            self.is_unbounded = (start <= 0) and (stop >= obj_len)
+        else:
+            self.is_unbounded = (stop == -1) and (start >= obj_len - 1)
+
+        # determine slice length
+        if self.obj.is_empty():
+            self.slice_length = 0
+        elif self.is_unbounded:
+            self.slice_length = obj_len
+        else:
+            self.slice_length = (
+                0
+                if (
+                    (start == stop)
+                    or (stride > 0 and start > stop)
+                    or (stride < 0 and start < stop)
+                )
+                else abs(stop - start)
+            )
+        self.start, self.stop, self.stride = start, stop, stride
+
+    def apply(self, s: slice) -> FrameOrSeries:
+        """Apply a slice operation, taking advantage of any potential fast paths."""
+        # normalise slice
+        self._slice_setup(s)
+
+        # check for fast-paths / single-operation calls
+        if self.slice_length == 0:
+            return self.obj.clear()
+
+        elif self.is_unbounded and self.stride in (-1, 1):
+            return self.obj.reverse() if (self.stride < 0) else self.obj.clone()
+
+        elif self.start >= 0 and self.stop >= 0 and self.stride == 1:
+            return self.obj.slice(self.start, self.slice_length)
+
+        elif self.stride < 0 and self.slice_length == 1:
+            return self.obj.slice(self.stop + 1, 1)
+        else:
+            # multi-operation calls; make lazy
+            lazyobj = self._lazify(self.obj)
+            sliced = (
+                self._slice_positive(lazyobj)
+                if self.stride > 0
+                else self._slice_negative(lazyobj)
+            )
+            return self._as_original(sliced, self.obj)
+
+
+class LazyPolarsSlice:
+    """
+    Apply python slice object to Polars LazyFrame.
+
+    Only slices with efficient computation paths that map directly
+    to existing lazy methods are supported.
+    """
+
+    obj: LazyFrame
+
+    def __init__(self, obj: LazyFrame) -> None:
+        self.obj = obj
+
+    def apply(self, s: slice) -> LazyFrame:
+        """
+        Apply a slice operation.
+
+        Note that LazyFrame is designed primarily for efficient computation and does not
+        know its own length so, unlike DataFrame, certain slice patterns (such as those
+        requiring negative stop/step) may not be supported.
+        """
+        start = s.start or 0
+        step = s.step or 1
+
+        # fail on operations that require length to do efficiently
+        if s.stop and s.stop < 0:
+            msg = "negative stop is not supported for lazy slices"
+            raise ValueError(msg)
+        if step < 0 and (start > 0 or s.stop is not None) and (start != s.stop):
+            if not (start > 0 > step and s.stop is None):
+                msg = "negative stride is not supported in conjunction with start+stop"
+                raise ValueError(msg)
+
+        # ---------------------------------------
+        # empty slice patterns
+        # ---------------------------------------
+        # [:0]
+        # [i:<=i]
+        # [i:>=i:-k]
+        if (step > 0 and (s.stop is not None and start >= s.stop)) or (
+            step < 0
+            and (s.start is not None and s.stop is not None and s.stop >= s.start >= 0)
+        ):
+            return self.obj.clear()
+
+        # ---------------------------------------
+        # straight-through mappings for "reverse"
+        # and/or "gather_every"
+        # ---------------------------------------
+        # [:]    => clone()
+        # [::k]  => gather_every(k),
+        # [::-1] => reverse(),
+        # [::-k] => reverse().gather_every(abs(k))
+        elif s.start is None and s.stop is None:
+            if step == 1:
+                return self.obj.clone()
+            elif step > 1:
+                return self.obj.gather_every(step)
+            elif step == -1:
+                return self.obj.reverse()
+            elif step < -1:
+                return self.obj.reverse().gather_every(abs(step))
+
+        # ---------------------------------------
+        # 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))
+
+        # ---------------------------------------
+        # straight-through mappings for "head"
+        # ---------------------------------------
+        # [:j]    => head(j)
+        # [:j:k]  => head(j).gather_every(k)
+        elif start == 0 and (s.stop or 0) >= 1:
+            obj = self.obj.head(s.stop)
+            return obj if (step == 1) else obj.gather_every(step)
+
+        # ---------------------------------------
+        # straight-through mappings for "tail"
+        # ---------------------------------------
+        # [-i:]    => tail(abs(i))
+        # [-i::k]  => tail(abs(i)).gather_every(k)
+        elif start < 0 and s.stop is None and step > 0:
+            obj = self.obj.tail(abs(start))
+            return obj if (step == 1) else obj.gather_every(step)
+
+        # ---------------------------------------
+        # straight-through mappings for "slice"
+        # ---------------------------------------
+        # [i:]     => slice(i)
+        # [i:j]    => slice(i,j-i)
+        # [i:j:k]  => slice(i,j-i).gather_every(k)
+        elif start > 0 and (s.stop is None or s.stop >= 0):
+            slice_length = None if (s.stop is None) else (s.stop - start)
+            obj = self.obj.slice(start, slice_length)
+            return obj if (step == 1) else obj.gather_every(step)
+
+        msg = (
+            f"the given slice {s!r} is not supported by lazy computation"
+            "\n\nConsider a more efficient approach, or construct explicitly with other methods."
+        )
+        raise ValueError(msg)