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

cudf_polars/dsl/expressions/aggregation.py

@@ -0,0 +1,226 @@
+# 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 aggregations."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING, Any, ClassVar
+
+import pylibcudf as plc
+
+from cudf_polars.containers import Column
+from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
+from cudf_polars.dsl.expressions.literal import Literal
+
+if TYPE_CHECKING:
+    from cudf_polars.containers import DataFrame, DataType
+
+__all__ = ["Agg"]
+
+
+class Agg(Expr):
+    __slots__ = ("name", "op", "options", "request")
+    _non_child = ("dtype", "name", "options")
+
+    def __init__(
+        self, dtype: DataType, name: str, options: Any, *children: Expr
+    ) -> None:
+        self.dtype = dtype
+        self.name = name
+        self.options = options
+        self.is_pointwise = False
+        self.children = children
+        if name not in Agg._SUPPORTED:
+            raise NotImplementedError(
+                f"Unsupported aggregation {name=}"
+            )  # pragma: no cover; all valid aggs are supported
+        # TODO: nan handling in groupby case
+        if name == "min":
+            req = plc.aggregation.min()
+        elif name == "max":
+            req = plc.aggregation.max()
+        elif name == "median":
+            req = plc.aggregation.median()
+        elif name == "n_unique":
+            # TODO: datatype of result
+            req = plc.aggregation.nunique(null_handling=plc.types.NullPolicy.INCLUDE)
+        elif name == "first" or name == "last":
+            req = None
+        elif name == "mean":
+            req = plc.aggregation.mean()
+        elif name == "sum":
+            req = plc.aggregation.sum()
+        elif name == "std":
+            # TODO: handle nans
+            req = plc.aggregation.std(ddof=options)
+        elif name == "var":
+            # TODO: handle nans
+            req = plc.aggregation.variance(ddof=options)
+        elif name == "count":
+            req = plc.aggregation.count(
+                null_handling=plc.types.NullPolicy.EXCLUDE
+                if not options
+                else plc.types.NullPolicy.INCLUDE
+            )
+        elif name == "quantile":
+            child, quantile = self.children
+            if not isinstance(quantile, Literal):
+                raise NotImplementedError("Only support literal quantile values")
+            if options == "equiprobable":
+                raise NotImplementedError("Quantile with equiprobable interpolation")
+            if plc.traits.is_duration(child.dtype.plc):
+                raise NotImplementedError("Quantile with duration data type")
+            req = plc.aggregation.quantile(
+                quantiles=[quantile.value], interp=Agg.interp_mapping[options]
+            )
+        else:
+            raise NotImplementedError(
+                f"Unreachable, {name=} is incorrectly listed in _SUPPORTED"
+            )  # pragma: no cover
+        self.request = req
+        op = getattr(self, f"_{name}", None)
+        if op is None:
+            op = partial(self._reduce, request=req)
+        elif name in {"min", "max"}:
+            op = partial(op, propagate_nans=options)
+        elif name == "count":
+            op = partial(op, include_nulls=options)
+        elif name in {"sum", "first", "last"}:
+            pass
+        else:
+            raise NotImplementedError(
+                f"Unreachable, supported agg {name=} has no implementation"
+            )  # pragma: no cover
+        self.op = op
+
+    _SUPPORTED: ClassVar[frozenset[str]] = frozenset(
+        [
+            "min",
+            "max",
+            "median",
+            "n_unique",
+            "first",
+            "last",
+            "mean",
+            "sum",
+            "count",
+            "std",
+            "var",
+            "quantile",
+        ]
+    )
+
+    interp_mapping: ClassVar[dict[str, plc.types.Interpolation]] = {
+        "nearest": plc.types.Interpolation.NEAREST,
+        "higher": plc.types.Interpolation.HIGHER,
+        "lower": plc.types.Interpolation.LOWER,
+        "midpoint": plc.types.Interpolation.MIDPOINT,
+        "linear": plc.types.Interpolation.LINEAR,
+    }
+
+    @property
+    def agg_request(self) -> plc.aggregation.Aggregation:  # noqa: D102
+        if self.name == "first":
+            return plc.aggregation.nth_element(
+                0, null_handling=plc.types.NullPolicy.INCLUDE
+            )
+        elif self.name == "last":
+            return plc.aggregation.nth_element(
+                -1, null_handling=plc.types.NullPolicy.INCLUDE
+            )
+        else:
+            assert self.request is not None, "Init should have raised"
+            return self.request
+
+    def _reduce(
+        self, column: Column, *, request: plc.aggregation.Aggregation
+    ) -> Column:
+        return Column(
+            plc.Column.from_scalar(
+                plc.reduce.reduce(column.obj, request, self.dtype.plc),
+                1,
+            ),
+            name=column.name,
+            dtype=self.dtype,
+        )
+
+    def _count(self, column: Column, *, include_nulls: bool) -> Column:
+        null_count = column.null_count if not include_nulls else 0
+        return Column(
+            plc.Column.from_scalar(
+                plc.Scalar.from_py(column.size - null_count, self.dtype.plc),
+                1,
+            ),
+            name=column.name,
+            dtype=self.dtype,
+        )
+
+    def _sum(self, column: Column) -> Column:
+        if column.size == 0 or column.null_count == column.size:
+            return Column(
+                plc.Column.from_scalar(
+                    plc.Scalar.from_py(0, self.dtype.plc),
+                    1,
+                ),
+                name=column.name,
+                dtype=self.dtype,
+            )
+        return self._reduce(column, request=plc.aggregation.sum())
+
+    def _min(self, column: Column, *, propagate_nans: bool) -> Column:
+        if propagate_nans and column.nan_count > 0:
+            return Column(
+                plc.Column.from_scalar(
+                    plc.Scalar.from_py(float("nan"), self.dtype.plc),
+                    1,
+                ),
+                name=column.name,
+                dtype=self.dtype,
+            )
+        if column.nan_count > 0:
+            column = column.mask_nans()
+        return self._reduce(column, request=plc.aggregation.min())
+
+    def _max(self, column: Column, *, propagate_nans: bool) -> Column:
+        if propagate_nans and column.nan_count > 0:
+            return Column(
+                plc.Column.from_scalar(
+                    plc.Scalar.from_py(float("nan"), self.dtype.plc),
+                    1,
+                ),
+                name=column.name,
+                dtype=self.dtype,
+            )
+        if column.nan_count > 0:
+            column = column.mask_nans()
+        return self._reduce(column, request=plc.aggregation.max())
+
+    def _first(self, column: Column) -> Column:
+        return Column(
+            plc.copying.slice(column.obj, [0, 1])[0], name=column.name, dtype=self.dtype
+        )
+
+    def _last(self, column: Column) -> Column:
+        n = column.size
+        return Column(
+            plc.copying.slice(column.obj, [n - 1, n])[0],
+            name=column.name,
+            dtype=self.dtype,
+        )
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        if context is not ExecutionContext.FRAME:
+            raise NotImplementedError(
+                f"Agg in context {context}"
+            )  # pragma: no cover; unreachable
+
+        # Aggregations like quantiles may have additional children that were
+        # preprocessed into pylibcudf requests.
+        child = self.children[0]
+        return self.op(child.evaluate(df, context=context))

cudf_polars/dsl/expressions/base.py

@@ -0,0 +1,272 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# TODO: remove need for this
+# ruff: noqa: D101
+"""Base and common classes for expression DSL nodes."""
+
+from __future__ import annotations
+
+import enum
+from enum import IntEnum
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
+
+import pylibcudf as plc
+
+from cudf_polars.containers import Column
+from cudf_polars.dsl.nodebase import Node
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from cudf_polars.containers import Column, DataFrame, DataType
+
+__all__ = ["AggInfo", "Col", "ColRef", "ExecutionContext", "Expr", "NamedExpr"]
+
+
+class AggInfo(NamedTuple):
+    requests: list[tuple[Expr | None, plc.aggregation.Aggregation, Expr]]
+
+
+class ExecutionContext(IntEnum):
+    FRAME = enum.auto()
+    GROUPBY = enum.auto()
+    ROLLING = enum.auto()
+    # Follows GROUPBY semantics but useful
+    # to differentiate from GROUPBY so we can
+    # implement agg/per-row ops independently
+    WINDOW = enum.auto()
+
+
+class Expr(Node["Expr"]):
+    """An abstract expression object."""
+
+    __slots__ = ("dtype", "is_pointwise")
+    dtype: DataType
+    """Data type of the expression."""
+    is_pointwise: bool
+    """Whether this expression acts pointwise on its inputs."""
+    # This annotation is needed because of https://github.com/python/mypy/issues/17981
+    _non_child: ClassVar[tuple[str, ...]] = ("dtype",)
+    """Names of non-child data (not Exprs) for reconstruction."""
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """
+        Evaluate this expression given a dataframe for context.
+
+        Parameters
+        ----------
+        df
+            DataFrame that will provide columns.
+        context
+            What context are we performing this evaluation in?
+
+        Notes
+        -----
+        Do not call this function directly, but rather :meth:`evaluate`.
+
+        Returns
+        -------
+        Column representing the evaluation of the expression.
+
+        Raises
+        ------
+        NotImplementedError
+            If we couldn't evaluate the expression. Ideally all these
+            are returned during translation to the IR, but for now we
+            are not perfect.
+        """
+        raise NotImplementedError(
+            f"Evaluation of expression {type(self).__name__}"
+        )  # pragma: no cover; translation of unimplemented nodes trips first
+
+    def evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """
+        Evaluate this expression given a dataframe for context.
+
+        Parameters
+        ----------
+        df
+            DataFrame that will provide columns.
+        context
+            What context are we performing this evaluation in?
+
+        Notes
+        -----
+        Individual subclasses should implement :meth:`do_evaluate`,
+        this method provides logic to handle lookups in the
+        substitution mapping.
+
+        Returns
+        -------
+        Column representing the evaluation of the expression.
+
+        Raises
+        ------
+        NotImplementedError
+            If we couldn't evaluate the expression. Ideally all these
+            are returned during translation to the IR, but for now we
+            are not perfect.
+        """
+        return self.do_evaluate(df, context=context)
+
+    @property
+    def agg_request(self) -> plc.aggregation.Aggregation:
+        """
+        The aggregation for this expression in a grouped aggregation.
+
+        Returns
+        -------
+        Aggregation request. Default is to collect the expression.
+
+        Notes
+        -----
+        This presumes that the IR translation has decomposed groupby
+        reductions only into cases we can handle.
+
+        Raises
+        ------
+        NotImplementedError
+            If requesting an aggregation from an unexpected expression.
+        """
+        return plc.aggregation.collect_list()
+
+
+class ErrorExpr(Expr):
+    __slots__ = ("error",)
+    _non_child = ("dtype", "error")
+    error: str
+
+    def __init__(self, dtype: DataType, error: str) -> None:
+        self.dtype = dtype
+        self.error = error
+        self.children = ()
+        self.is_pointwise = False
+
+
+class NamedExpr:
+    # NamedExpr does not inherit from Expr since it does not appear
+    # when evaluating expressions themselves, only when constructing
+    # named return values in dataframe (IR) nodes.
+    __slots__ = ("name", "value")
+    value: Expr
+    name: str
+
+    def __init__(self, name: str, value: Expr) -> None:
+        self.name = name
+        self.value = value
+
+    def __hash__(self) -> int:
+        """Hash of the expression."""
+        return hash((type(self), self.name, self.value))
+
+    def __repr__(self) -> str:
+        """Repr of the expression."""
+        return f"NamedExpr({self.name}, {self.value})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Equality of two expressions."""
+        return (
+            type(self) is type(other)
+            and self.name == other.name
+            and self.value == other.value
+        )
+
+    def __ne__(self, other: Any) -> bool:
+        """Inequality of expressions."""
+        return not self.__eq__(other)
+
+    def evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """
+        Evaluate this expression given a dataframe for context.
+
+        Parameters
+        ----------
+        df
+            DataFrame providing context
+        context
+            Execution context
+
+        Returns
+        -------
+        Evaluated Column with name attached.
+
+        See Also
+        --------
+        :meth:`Expr.evaluate` for details, this function just adds the
+        name to a column produced from an expression.
+        """
+        return self.value.evaluate(df, context=context).rename(self.name)
+
+    def reconstruct(self, expr: Expr) -> Self:
+        """
+        Rebuild with a new `Expr` value.
+
+        Parameters
+        ----------
+        expr
+            New `Expr` value
+
+        Returns
+        -------
+        New `NamedExpr` with `expr` as the underlying expression.
+        The name of the original `NamedExpr` is preserved.
+        """
+        if expr is self.value:
+            return self
+        return type(self)(self.name, expr)
+
+
+class Col(Expr):
+    __slots__ = ("name",)
+    _non_child = ("dtype", "name")
+    name: str
+
+    def __init__(self, dtype: DataType, name: str) -> None:
+        self.dtype = dtype
+        self.name = name
+        self.is_pointwise = True
+        self.children = ()
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        # Deliberately remove the name here so that we guarantee
+        # evaluation of the IR produces names.
+        return df.column_map[self.name].rename(None)
+
+
+class ColRef(Expr):
+    __slots__ = ("index", "table_ref")
+    _non_child = ("dtype", "index", "table_ref")
+    index: int
+    table_ref: plc.expressions.TableReference
+
+    def __init__(
+        self,
+        dtype: DataType,
+        index: int,
+        table_ref: plc.expressions.TableReference,
+        column: Expr,
+    ) -> None:
+        if not isinstance(column, Col):
+            raise TypeError("Column reference should only apply to columns")
+        self.dtype = dtype
+        self.index = index
+        self.table_ref = table_ref
+        self.is_pointwise = True
+        self.children = (column,)
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        raise NotImplementedError(
+            "Only expect this node as part of an expression translated to libcudf AST."
+        )

cudf_polars/dsl/expressions/binaryop.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+# TODO: remove need for this
+# ruff: noqa: D101
+"""BinaryOp DSL nodes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+from polars.polars import _expr_nodes as pl_expr
+
+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__ = ["BinOp"]
+
+
+class BinOp(Expr):
+    __slots__ = ("op",)
+    _non_child = ("dtype", "op")
+
+    def __init__(
+        self,
+        dtype: DataType,
+        op: plc.binaryop.BinaryOperator,
+        left: Expr,
+        right: Expr,
+    ) -> None:
+        self.dtype = dtype
+        if plc.traits.is_boolean(self.dtype.plc):
+            # For boolean output types, bitand and bitor implement
+            # boolean logic, so translate. bitxor also does, but the
+            # default behaviour is correct.
+            op = BinOp._BOOL_KLEENE_MAPPING.get(op, op)
+        self.op = op
+        self.children = (left, right)
+        self.is_pointwise = True
+        if not plc.binaryop.is_supported_operation(
+            self.dtype.plc, left.dtype.plc, right.dtype.plc, op
+        ):
+            raise NotImplementedError(
+                f"Operation {op.name} not supported "
+                f"for types {left.dtype.id().name} and {right.dtype.id().name} "
+                f"with output type {self.dtype.id().name}"
+            )
+
+    _BOOL_KLEENE_MAPPING: ClassVar[
+        dict[plc.binaryop.BinaryOperator, plc.binaryop.BinaryOperator]
+    ] = {
+        plc.binaryop.BinaryOperator.BITWISE_AND: plc.binaryop.BinaryOperator.NULL_LOGICAL_AND,
+        plc.binaryop.BinaryOperator.BITWISE_OR: plc.binaryop.BinaryOperator.NULL_LOGICAL_OR,
+        plc.binaryop.BinaryOperator.LOGICAL_AND: plc.binaryop.BinaryOperator.NULL_LOGICAL_AND,
+        plc.binaryop.BinaryOperator.LOGICAL_OR: plc.binaryop.BinaryOperator.NULL_LOGICAL_OR,
+    }
+
+    _MAPPING: ClassVar[dict[pl_expr.Operator, plc.binaryop.BinaryOperator]] = {
+        pl_expr.Operator.Eq: plc.binaryop.BinaryOperator.EQUAL,
+        pl_expr.Operator.EqValidity: plc.binaryop.BinaryOperator.NULL_EQUALS,
+        pl_expr.Operator.NotEq: plc.binaryop.BinaryOperator.NOT_EQUAL,
+        pl_expr.Operator.NotEqValidity: plc.binaryop.BinaryOperator.NULL_NOT_EQUALS,
+        pl_expr.Operator.Lt: plc.binaryop.BinaryOperator.LESS,
+        pl_expr.Operator.LtEq: plc.binaryop.BinaryOperator.LESS_EQUAL,
+        pl_expr.Operator.Gt: plc.binaryop.BinaryOperator.GREATER,
+        pl_expr.Operator.GtEq: plc.binaryop.BinaryOperator.GREATER_EQUAL,
+        pl_expr.Operator.Plus: plc.binaryop.BinaryOperator.ADD,
+        pl_expr.Operator.Minus: plc.binaryop.BinaryOperator.SUB,
+        pl_expr.Operator.Multiply: plc.binaryop.BinaryOperator.MUL,
+        pl_expr.Operator.Divide: plc.binaryop.BinaryOperator.DIV,
+        pl_expr.Operator.TrueDivide: plc.binaryop.BinaryOperator.TRUE_DIV,
+        pl_expr.Operator.FloorDivide: plc.binaryop.BinaryOperator.FLOOR_DIV,
+        pl_expr.Operator.Modulus: plc.binaryop.BinaryOperator.PYMOD,
+        pl_expr.Operator.And: plc.binaryop.BinaryOperator.BITWISE_AND,
+        pl_expr.Operator.Or: plc.binaryop.BinaryOperator.BITWISE_OR,
+        pl_expr.Operator.Xor: plc.binaryop.BinaryOperator.BITWISE_XOR,
+        pl_expr.Operator.LogicalAnd: plc.binaryop.BinaryOperator.LOGICAL_AND,
+        pl_expr.Operator.LogicalOr: plc.binaryop.BinaryOperator.LOGICAL_OR,
+    }
+
+    def do_evaluate(
+        self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        left, right = (child.evaluate(df, context=context) for child in self.children)
+        lop = left.obj
+        rop = right.obj
+        if left.size != right.size:
+            if left.is_scalar:
+                lop = left.obj_scalar
+            elif right.is_scalar:
+                rop = right.obj_scalar
+        if plc.traits.is_integral_not_bool(self.dtype.plc) and self.op in {
+            plc.binaryop.BinaryOperator.FLOOR_DIV,
+            plc.binaryop.BinaryOperator.PYMOD,
+        }:
+            if right.obj.size() == 1 and right.obj.to_scalar().to_py() == 0:
+                return Column(
+                    plc.Column.all_null_like(left.obj, left.obj.size()),
+                    dtype=self.dtype,
+                )
+
+            if right.obj.size() > 1:
+                rop = plc.replace.find_and_replace_all(
+                    right.obj,
+                    plc.Column.from_scalar(
+                        plc.Scalar.from_py(0, dtype=self.dtype.plc), 1
+                    ),
+                    plc.Column.from_scalar(
+                        plc.Scalar.from_py(None, dtype=self.dtype.plc), 1
+                    ),
+                )
+        return Column(
+            plc.binaryop.binary_operation(lop, rop, self.op, self.dtype.plc),
+            dtype=self.dtype,
+        )