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

cudf_polars/dsl/nodebase.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Base class for IR nodes, and utilities."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Hashable, Sequence
+
+    from typing_extensions import Self
+
+
+__all__: list[str] = ["Node"]
+
+T = TypeVar("T", bound="Node[Any]")
+
+
+class Node(Generic[T]):
+    """
+    An abstract node type.
+
+    Nodes are immutable!
+
+    This contains a (potentially empty) tuple of child nodes,
+    along with non-child data. For uniform reconstruction and
+    implementation of hashing and equality schemes, child classes need
+    to provide a certain amount of metadata when they are defined.
+    Specifically, the ``_non_child`` attribute must list, in-order,
+    the names of the slots that are passed to the constructor. The
+    constructor must take arguments in the order ``(*_non_child,
+    *children).``
+    """
+
+    __slots__ = ("_hash_value", "_repr_value", "children")
+    _hash_value: int
+    _repr_value: str
+    children: tuple[T, ...]
+    _non_child: ClassVar[tuple[str, ...]] = ()
+
+    def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any | T]:
+        return (*(getattr(self, attr) for attr in self._non_child), *children)
+
+    def reconstruct(self, children: Sequence[T]) -> Self:
+        """
+        Rebuild this node with new children.
+
+        Parameters
+        ----------
+        children
+            New children
+
+        Returns
+        -------
+        New node with new children. Non-child data is shared with the input.
+        """
+        return type(self)(*self._ctor_arguments(children))
+
+    def __reduce__(self) -> tuple[Any, ...]:
+        """Pickle a Node object."""
+        return (
+            type(self),
+            self._ctor_arguments(self.children),
+        )
+
+    def get_hashable(self) -> Hashable:
+        """
+        Return a hashable object for the node.
+
+        Returns
+        -------
+        Hashable object.
+
+        Notes
+        -----
+        This method is used by the :meth:`__hash__` implementation
+        (which does caching). If your node type needs special-case
+        handling for some of its attributes, override this method, not
+        :meth:`__hash__`.
+        """
+        return (type(self), self._ctor_arguments(self.children))
+
+    def __hash__(self) -> int:
+        """
+        Hash of an expression with caching.
+
+        See Also
+        --------
+        get_hashable
+        """
+        try:
+            return self._hash_value
+        except AttributeError:
+            self._hash_value = hash(self.get_hashable())
+            return self._hash_value
+
+    def is_equal(self, other: Self) -> bool:
+        """
+        Equality of two nodes of equal type.
+
+        Override this in subclasses, rather than :meth:`__eq__`.
+
+        Parameters
+        ----------
+        other
+            object of same type to compare to.
+
+        Notes
+        -----
+        Since nodes are immutable, this does common subexpression
+        elimination when two nodes are determined to be equal.
+
+        :meth:`__eq__` handles the case where the objects being
+        compared are not of the same type, so in this method, we only
+        need to implement equality of equal types.
+
+        Returns
+        -------
+        True if the two nodes are equal, false otherwise.
+        """
+        if self is other:
+            return True
+        result = self._ctor_arguments(self.children) == other._ctor_arguments(
+            other.children
+        )
+        # Eager CSE for nodes that match.
+        if result:
+            self.children = other.children
+        return result
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        Equality of expressions.
+
+        See Also
+        --------
+        is_equal
+        """
+        if type(self) is not type(other) or hash(self) != hash(other):
+            return False
+        else:
+            return self.is_equal(other)
+
+    def __ne__(self, other: Any) -> bool:
+        """Inequality of expressions."""
+        return not self.__eq__(other)
+
+    def __repr__(self) -> str:
+        """String representation of an expression with caching."""
+        try:
+            return self._repr_value
+        except AttributeError:
+            args = ", ".join(f"{arg!r}" for arg in self._ctor_arguments(self.children))
+            self._repr_value = f"{type(self).__name__}({args})"
+            return self._repr_value
+
+    def __rich_repr__(self) -> Generator[Any, None, None]:
+        """Formatting for rich.pretty.pprint."""
+        for attr in self._non_child:
+            yield attr, getattr(self, attr)
+
+        yield from self.children

cudf_polars/dsl/to_ast.py

@@ -0,0 +1,359 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Conversion of expression nodes to libcudf AST nodes."""
+
+from __future__ import annotations
+
+from functools import partial, reduce, singledispatch
+from typing import TYPE_CHECKING, TypeAlias, TypedDict
+
+import pylibcudf as plc
+from pylibcudf import expressions as plc_expr
+
+from cudf_polars.containers import DataType
+from cudf_polars.dsl import expr
+from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged
+from cudf_polars.typing import GenericTransformer
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+# Can't merge these op-mapping dictionaries because scoped enum values
+# are exposed by cython with equality/hash based one their underlying
+# representation type. So in a dict they are just treated as integers.
+BINOP_TO_ASTOP = {
+    plc.binaryop.BinaryOperator.EQUAL: plc_expr.ASTOperator.EQUAL,
+    plc.binaryop.BinaryOperator.NULL_EQUALS: plc_expr.ASTOperator.NULL_EQUAL,
+    plc.binaryop.BinaryOperator.NOT_EQUAL: plc_expr.ASTOperator.NOT_EQUAL,
+    plc.binaryop.BinaryOperator.LESS: plc_expr.ASTOperator.LESS,
+    plc.binaryop.BinaryOperator.LESS_EQUAL: plc_expr.ASTOperator.LESS_EQUAL,
+    plc.binaryop.BinaryOperator.GREATER: plc_expr.ASTOperator.GREATER,
+    plc.binaryop.BinaryOperator.GREATER_EQUAL: plc_expr.ASTOperator.GREATER_EQUAL,
+    plc.binaryop.BinaryOperator.ADD: plc_expr.ASTOperator.ADD,
+    plc.binaryop.BinaryOperator.SUB: plc_expr.ASTOperator.SUB,
+    plc.binaryop.BinaryOperator.MUL: plc_expr.ASTOperator.MUL,
+    plc.binaryop.BinaryOperator.DIV: plc_expr.ASTOperator.DIV,
+    plc.binaryop.BinaryOperator.TRUE_DIV: plc_expr.ASTOperator.TRUE_DIV,
+    plc.binaryop.BinaryOperator.FLOOR_DIV: plc_expr.ASTOperator.FLOOR_DIV,
+    plc.binaryop.BinaryOperator.PYMOD: plc_expr.ASTOperator.PYMOD,
+    plc.binaryop.BinaryOperator.BITWISE_AND: plc_expr.ASTOperator.BITWISE_AND,
+    plc.binaryop.BinaryOperator.BITWISE_OR: plc_expr.ASTOperator.BITWISE_OR,
+    plc.binaryop.BinaryOperator.BITWISE_XOR: plc_expr.ASTOperator.BITWISE_XOR,
+    plc.binaryop.BinaryOperator.LOGICAL_AND: plc_expr.ASTOperator.LOGICAL_AND,
+    plc.binaryop.BinaryOperator.LOGICAL_OR: plc_expr.ASTOperator.LOGICAL_OR,
+    plc.binaryop.BinaryOperator.NULL_LOGICAL_AND: plc_expr.ASTOperator.NULL_LOGICAL_AND,
+    plc.binaryop.BinaryOperator.NULL_LOGICAL_OR: plc_expr.ASTOperator.NULL_LOGICAL_OR,
+}
+
+UOP_TO_ASTOP = {
+    plc.unary.UnaryOperator.SIN: plc_expr.ASTOperator.SIN,
+    plc.unary.UnaryOperator.COS: plc_expr.ASTOperator.COS,
+    plc.unary.UnaryOperator.TAN: plc_expr.ASTOperator.TAN,
+    plc.unary.UnaryOperator.ARCSIN: plc_expr.ASTOperator.ARCSIN,
+    plc.unary.UnaryOperator.ARCCOS: plc_expr.ASTOperator.ARCCOS,
+    plc.unary.UnaryOperator.ARCTAN: plc_expr.ASTOperator.ARCTAN,
+    plc.unary.UnaryOperator.SINH: plc_expr.ASTOperator.SINH,
+    plc.unary.UnaryOperator.COSH: plc_expr.ASTOperator.COSH,
+    plc.unary.UnaryOperator.TANH: plc_expr.ASTOperator.TANH,
+    plc.unary.UnaryOperator.ARCSINH: plc_expr.ASTOperator.ARCSINH,
+    plc.unary.UnaryOperator.ARCCOSH: plc_expr.ASTOperator.ARCCOSH,
+    plc.unary.UnaryOperator.ARCTANH: plc_expr.ASTOperator.ARCTANH,
+    plc.unary.UnaryOperator.EXP: plc_expr.ASTOperator.EXP,
+    plc.unary.UnaryOperator.LOG: plc_expr.ASTOperator.LOG,
+    plc.unary.UnaryOperator.SQRT: plc_expr.ASTOperator.SQRT,
+    plc.unary.UnaryOperator.CBRT: plc_expr.ASTOperator.CBRT,
+    plc.unary.UnaryOperator.CEIL: plc_expr.ASTOperator.CEIL,
+    plc.unary.UnaryOperator.FLOOR: plc_expr.ASTOperator.FLOOR,
+    plc.unary.UnaryOperator.ABS: plc_expr.ASTOperator.ABS,
+    plc.unary.UnaryOperator.RINT: plc_expr.ASTOperator.RINT,
+    plc.unary.UnaryOperator.BIT_INVERT: plc_expr.ASTOperator.BIT_INVERT,
+    plc.unary.UnaryOperator.NOT: plc_expr.ASTOperator.NOT,
+}
+
+SUPPORTED_STATISTICS_BINOPS = {
+    plc.binaryop.BinaryOperator.EQUAL,
+    plc.binaryop.BinaryOperator.NOT_EQUAL,
+    plc.binaryop.BinaryOperator.LESS,
+    plc.binaryop.BinaryOperator.LESS_EQUAL,
+    plc.binaryop.BinaryOperator.GREATER,
+    plc.binaryop.BinaryOperator.GREATER_EQUAL,
+}
+
+REVERSED_COMPARISON = {
+    plc.binaryop.BinaryOperator.EQUAL: plc.binaryop.BinaryOperator.EQUAL,
+    plc.binaryop.BinaryOperator.NOT_EQUAL: plc.binaryop.BinaryOperator.NOT_EQUAL,
+    plc.binaryop.BinaryOperator.LESS: plc.binaryop.BinaryOperator.GREATER,
+    plc.binaryop.BinaryOperator.LESS_EQUAL: plc.binaryop.BinaryOperator.GREATER_EQUAL,
+    plc.binaryop.BinaryOperator.GREATER: plc.binaryop.BinaryOperator.LESS,
+    plc.binaryop.BinaryOperator.GREATER_EQUAL: plc.binaryop.BinaryOperator.LESS_EQUAL,
+}
+
+
+class ASTState(TypedDict):
+    """
+    State for AST transformations.
+
+    Parameters
+    ----------
+    for_parquet
+        Indicator for whether this transformation should provide an expression
+        suitable for use in parquet filters.
+    """
+
+    for_parquet: bool
+
+
+class ExprTransformerState(TypedDict):
+    """
+    State used for AST transformation when inserting column references.
+
+    Parameters
+    ----------
+    name_to_index
+        Mapping from column names to column indices in the table
+        eventually used for evaluation.
+    table_ref
+        pylibcudf `TableReference` indicating whether column
+        references are coming from the left or right table.
+    """
+
+    name_to_index: Mapping[str, int]
+    table_ref: plc.expressions.TableReference
+
+
+Transformer: TypeAlias = GenericTransformer[expr.Expr, plc_expr.Expression, ASTState]
+ExprTransformer: TypeAlias = GenericTransformer[
+    expr.Expr, expr.Expr, ExprTransformerState
+]
+"""Protocol for transformation of Expr nodes."""
+
+
+@singledispatch
+def _to_ast(node: expr.Expr, self: Transformer) -> plc_expr.Expression:
+    """
+    Translate an expression to a pylibcudf Expression.
+
+    Parameters
+    ----------
+    node
+        Expression to translate.
+    self
+        Recursive transformer. The state dictionary is an instance of
+        :class:`ASTState`.
+
+    Returns
+    -------
+    pylibcudf Expression.
+
+    Raises
+    ------
+    NotImplementedError or KeyError if the expression cannot be translated.
+    """
+    raise NotImplementedError(f"Unhandled expression type {type(node)}")
+
+
+@_to_ast.register
+def _(node: expr.Col, self: Transformer) -> plc_expr.Expression:
+    if self.state["for_parquet"]:
+        return plc_expr.ColumnNameReference(node.name)
+    raise TypeError("Should always be wrapped in a ColRef node before translation")
+
+
+@_to_ast.register
+def _(node: expr.ColRef, self: Transformer) -> plc_expr.Expression:
+    if self.state["for_parquet"]:
+        raise TypeError("Not expecting ColRef node in parquet filter")
+    return plc_expr.ColumnReference(node.index, node.table_ref)
+
+
+@_to_ast.register
+def _(node: expr.Literal, self: Transformer) -> plc_expr.Expression:
+    return plc_expr.Literal(plc.Scalar.from_py(node.value, node.dtype.plc))
+
+
+@_to_ast.register
+def _(node: expr.BinOp, self: Transformer) -> plc_expr.Expression:
+    if node.op == plc.binaryop.BinaryOperator.NULL_NOT_EQUALS:
+        return plc_expr.Operation(
+            plc_expr.ASTOperator.NOT,
+            self(
+                # Reconstruct and apply, rather than directly
+                # constructing the right expression so we get the
+                # handling of parquet special cases for free.
+                expr.BinOp(
+                    node.dtype, plc.binaryop.BinaryOperator.NULL_EQUALS, *node.children
+                )
+            ),
+        )
+    if self.state["for_parquet"]:
+        op1_col, op2_col = (isinstance(op, expr.Col) for op in node.children)
+        if op1_col ^ op2_col:
+            op = node.op
+            if op not in SUPPORTED_STATISTICS_BINOPS:
+                raise NotImplementedError(
+                    f"Parquet filter binop with column doesn't support {node.op!r}"
+                )
+            op1, op2 = node.children
+            if op2_col:
+                (op1, op2) = (op2, op1)
+                op = REVERSED_COMPARISON[op]
+            if not isinstance(op2, expr.Literal):
+                raise NotImplementedError(
+                    "Parquet filter binops must have form 'col binop literal'"
+                )
+            return plc_expr.Operation(BINOP_TO_ASTOP[op], self(op1), self(op2))
+        elif op1_col and op2_col:
+            raise NotImplementedError(
+                "Parquet filter binops must have one column reference not two"
+            )
+    return plc_expr.Operation(BINOP_TO_ASTOP[node.op], *map(self, node.children))
+
+
+@_to_ast.register
+def _(node: expr.BooleanFunction, self: Transformer) -> plc_expr.Expression:
+    if node.name is expr.BooleanFunction.Name.IsIn:
+        needles, haystack = node.children
+        if isinstance(haystack, expr.LiteralColumn) and len(haystack.value) < 16:
+            # 16 is an arbitrary limit
+            needle_ref = self(needles)
+            if haystack.dtype.id() == plc.TypeId.LIST:
+                # Because we originally translated pl_expr.Literal with a list scalar
+                # to a expr.LiteralColumn, so the actual type is in the inner type
+                #
+                # the type-ignore is safe because the for plc.TypeID.LIST, we know
+                # we have a polars.List type, which has an inner attribute.
+                plc_dtype = DataType(haystack.dtype.polars.inner).plc  # type: ignore[attr-defined]
+            else:
+                plc_dtype = haystack.dtype.plc  # pragma: no cover
+            values = (
+                plc_expr.Literal(plc.Scalar.from_py(val, plc_dtype))
+                for val in haystack.value
+            )
+            return reduce(
+                partial(plc_expr.Operation, plc_expr.ASTOperator.LOGICAL_OR),
+                (
+                    plc_expr.Operation(plc_expr.ASTOperator.EQUAL, needle_ref, value)
+                    for value in values
+                ),
+            )
+    if self.state["for_parquet"] and isinstance(node.children[0], expr.Col):
+        raise NotImplementedError(
+            f"Parquet filters don't support {node.name} on columns"
+        )
+    if node.name is expr.BooleanFunction.Name.IsNull:
+        return plc_expr.Operation(plc_expr.ASTOperator.IS_NULL, self(node.children[0]))
+    elif node.name is expr.BooleanFunction.Name.IsNotNull:
+        return plc_expr.Operation(
+            plc_expr.ASTOperator.NOT,
+            plc_expr.Operation(plc_expr.ASTOperator.IS_NULL, self(node.children[0])),
+        )
+    elif node.name is expr.BooleanFunction.Name.Not:
+        return plc_expr.Operation(plc_expr.ASTOperator.NOT, self(node.children[0]))
+    raise NotImplementedError(f"AST conversion does not support {node.name}")
+
+
+@_to_ast.register
+def _(node: expr.UnaryFunction, self: Transformer) -> plc_expr.Expression:
+    if isinstance(node.children[0], expr.Col) and self.state["for_parquet"]:
+        raise NotImplementedError(
+            "Parquet filters don't support {node.name} on columns"
+        )
+    return plc_expr.Operation(
+        UOP_TO_ASTOP[node._OP_MAPPING[node.name]], self(node.children[0])
+    )
+
+
+def to_parquet_filter(node: expr.Expr) -> plc_expr.Expression | None:
+    """
+    Convert an expression to libcudf AST nodes suitable for parquet filtering.
+
+    Parameters
+    ----------
+    node
+        Expression to convert.
+
+    Returns
+    -------
+    pylibcudf Expression if conversion is possible, otherwise None.
+    """
+    mapper: Transformer = CachingVisitor(_to_ast, state={"for_parquet": True})
+    try:
+        return mapper(node)
+    except (KeyError, NotImplementedError):
+        return None
+
+
+def to_ast(node: expr.Expr) -> plc_expr.Expression | None:
+    """
+    Convert an expression to libcudf AST nodes suitable for compute_column.
+
+    Parameters
+    ----------
+    node
+        Expression to convert.
+
+    Notes
+    -----
+    `Col` nodes must always be wrapped in `TableRef` nodes when
+    converting to an ast expression so that their table reference and
+    index are provided.
+
+    Returns
+    -------
+    pylibcudf Expression if conversion is possible, otherwise None.
+    """
+    mapper: Transformer = CachingVisitor(_to_ast, state={"for_parquet": False})
+    try:
+        return mapper(node)
+    except (KeyError, NotImplementedError):
+        return None
+
+
+def _insert_colrefs(node: expr.Expr, rec: ExprTransformer) -> expr.Expr:
+    if isinstance(node, expr.Col):
+        return expr.ColRef(
+            node.dtype,
+            rec.state["name_to_index"][node.name],
+            rec.state["table_ref"],
+            node,
+        )
+    return reuse_if_unchanged(node, rec)
+
+
+def insert_colrefs(
+    node: expr.Expr,
+    *,
+    table_ref: plc.expressions.TableReference,
+    name_to_index: Mapping[str, int],
+) -> expr.Expr:
+    """
+    Insert column references into an expression before conversion to libcudf AST.
+
+    Parameters
+    ----------
+    node
+        Expression to insert references into.
+    table_ref
+        pylibcudf `TableReference` indicating whether column
+        references are coming from the left or right table.
+    name_to_index:
+        Mapping from column names to column indices in the table
+        eventually used for evaluation.
+
+    Notes
+    -----
+    All column references are wrapped in the same, singular, table
+    reference, so this function relies on the expression only
+    containing column references from a single table.
+
+    Returns
+    -------
+    New expression with column references inserted.
+    """
+    mapper: ExprTransformer = CachingVisitor(
+        _insert_colrefs,
+        state={"name_to_index": name_to_index, "table_ref": table_ref},
+    )
+    return mapper(node)

cudf_polars/dsl/tracing.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for tracing and monitoring IR execution."""
+
+from __future__ import annotations
+
+import functools
+
+import nvtx
+
+CUDF_POLARS_NVTX_DOMAIN = "cudf_polars"
+
+nvtx_annotate_cudf_polars = functools.partial(
+    nvtx.annotate, domain=CUDF_POLARS_NVTX_DOMAIN
+)