cudf-polars-cu13 25.10.0__py3-none-any.whl

cudf_polars/dsl/expressions/selection.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# TODO: remove need for this
+# ruff: noqa: D101
+"""DSL nodes for selection operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pylibcudf as plc
+
+from cudf_polars.containers import Column
+from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
+
+if TYPE_CHECKING:
+    from cudf_polars.containers import DataFrame, DataType
+
+__all__ = ["Filter", "Gather"]
+
+
+class Gather(Expr):
+    __slots__ = ()
+    _non_child = ("dtype",)
+
+    def __init__(self, dtype: DataType, values: Expr, indices: Expr) -> None:
+        self.dtype = dtype
+        self.children = (values, indices)
+        self.is_pointwise = False
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        values, indices = (
+            child.evaluate(df, context=context) for child in self.children
+        )
+        lo, hi = plc.reduce.minmax(indices.obj)
+        lo = lo.to_py()
+        hi = hi.to_py()
+        n = df.num_rows
+        if hi >= n or lo < -n:
+            raise ValueError("gather indices are out of bounds")
+        if indices.null_count:
+            bounds_policy = plc.copying.OutOfBoundsPolicy.NULLIFY
+            obj = plc.replace.replace_nulls(
+                indices.obj,
+                plc.Scalar.from_py(n, dtype=indices.obj.type()),
+            )
+        else:
+            bounds_policy = plc.copying.OutOfBoundsPolicy.DONT_CHECK
+            obj = indices.obj
+        table = plc.copying.gather(plc.Table([values.obj]), obj, bounds_policy)
+        return Column(table.columns()[0], dtype=self.dtype)
+
+
+class Filter(Expr):
+    __slots__ = ()
+    _non_child = ("dtype",)
+
+    def __init__(self, dtype: DataType, values: Expr, indices: Expr):
+        self.dtype = dtype
+        self.children = (values, indices)
+        self.is_pointwise = False
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        values, mask = (child.evaluate(df, context=context) for child in self.children)
+        table = plc.stream_compaction.apply_boolean_mask(
+            plc.Table([values.obj]), mask.obj
+        )
+        return Column(table.columns()[0], dtype=self.dtype).sorted_like(values)

cudf_polars/dsl/expressions/slicing.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# TODO: remove need for this
+# ruff: noqa: D101
+"""Slicing DSL nodes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl.expressions.base import (
+    ExecutionContext,
+    Expr,
+)
+
+if TYPE_CHECKING:
+    from cudf_polars.containers import Column, DataFrame, DataType
+
+
+__all__ = ["Slice"]
+
+
+class Slice(Expr):
+    __slots__ = ("length", "offset")
+    _non_child = ("dtype", "offset", "length")
+
+    def __init__(
+        self,
+        dtype: DataType,
+        offset: int,
+        length: int | None,
+        column: Expr,
+    ) -> None:
+        self.dtype = dtype
+        self.offset = offset
+        self.length = length
+        self.children = (column,)
+        self.is_pointwise = False
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        (child,) = self.children
+        column = child.evaluate(df, context=context)
+        return column.slice((self.offset, self.length))

cudf_polars/dsl/expressions/sorting.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# TODO: remove need for this
+# ruff: noqa: D101
+"""Sorting DSL nodes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pylibcudf as plc
+
+from cudf_polars.containers import Column
+from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
+from cudf_polars.utils import sorting
+
+if TYPE_CHECKING:
+    from cudf_polars.containers import DataFrame, DataType
+
+__all__ = ["Sort", "SortBy"]
+
+
+class Sort(Expr):
+    __slots__ = ("options",)
+    _non_child = ("dtype", "options")
+
+    def __init__(
+        self, dtype: DataType, options: tuple[bool, bool, bool], column: Expr
+    ) -> None:
+        self.dtype = dtype
+        self.options = options
+        self.children = (column,)
+        self.is_pointwise = False
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        (child,) = self.children
+        column = child.evaluate(df, context=context)
+        (stable, nulls_last, descending) = self.options
+        order, null_order = sorting.sort_order(
+            [descending], nulls_last=[nulls_last], num_keys=1
+        )
+        do_sort = plc.sorting.stable_sort if stable else plc.sorting.sort
+        table = do_sort(plc.Table([column.obj]), order, null_order)
+        return Column(
+            table.columns()[0],
+            is_sorted=plc.types.Sorted.YES,
+            order=order[0],
+            null_order=null_order[0],
+            dtype=self.dtype,
+        )
+
+
+class SortBy(Expr):
+    __slots__ = ("options",)
+    _non_child = ("dtype", "options")
+
+    def __init__(
+        self,
+        dtype: DataType,
+        options: tuple[bool, tuple[bool], tuple[bool]],
+        column: Expr,
+        *by: Expr,
+    ) -> None:
+        self.dtype = dtype
+        self.options = options
+        self.children = (column, *by)
+        self.is_pointwise = False
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        column, *by = (child.evaluate(df, context=context) for child in self.children)
+        (stable, nulls_last, descending) = self.options
+        order, null_order = sorting.sort_order(
+            descending, nulls_last=nulls_last, num_keys=len(by)
+        )
+        do_sort = plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key
+        table = do_sort(
+            plc.Table([column.obj]), plc.Table([c.obj for c in by]), order, null_order
+        )
+        return Column(table.columns()[0], dtype=self.dtype)