cudf-polars-cu12 24.12.0__py3-none-any.whl → 25.2.1__py3-none-any.whl

cudf_polars/dsl/expressions/selection.py

@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 
     from cudf_polars.containers import DataFrame
 
-__all__ = ["Gather", "Filter"]
+__all__ = ["Filter", "Gather"]
 
 
 class Gather(Expr):
@@ -30,6 +30,7 @@ class Gather(Expr):
     def __init__(self, dtype: plc.DataType, values: Expr, indices: Expr) -> None:
         self.dtype = dtype
         self.children = (values, indices)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,
@@ -71,6 +72,7 @@ class Filter(Expr):
     def __init__(self, dtype: plc.DataType, values: Expr, indices: Expr):
         self.dtype = dtype
         self.children = (values, indices)
+        self.is_pointwise = True
 
     def do_evaluate(
         self,

cudf_polars/dsl/expressions/sorting.py

@@ -32,6 +32,7 @@ class Sort(Expr):
         self.dtype = dtype
         self.options = options
         self.children = (column,)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,
@@ -71,6 +72,7 @@ class SortBy(Expr):
         self.dtype = dtype
         self.options = options
         self.children = (column, *by)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,

cudf_polars/dsl/expressions/string.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 # TODO: remove need for this
 # ruff: noqa: D101
@@ -6,13 +6,13 @@
 
 from __future__ import annotations
 
+from enum import IntEnum, auto
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
 import pyarrow.compute as pc
 
 from polars.exceptions import InvalidOperationError
-from polars.polars import _expr_nodes as pl_expr
 
 import pylibcudf as plc
 
@@ -23,19 +23,83 @@ from cudf_polars.dsl.expressions.literal import Literal, LiteralColumn
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+    from typing_extensions import Self
+
+    from polars.polars import _expr_nodes as pl_expr
+
     from cudf_polars.containers import DataFrame
 
 __all__ = ["StringFunction"]
 
 
 class StringFunction(Expr):
-    __slots__ = ("name", "options", "_regex_program")
+    class Name(IntEnum):
+        """Internal and picklable representation of polars' `StringFunction`."""
+
+        Base64Decode = auto()
+        Base64Encode = auto()
+        ConcatHorizontal = auto()
+        ConcatVertical = auto()
+        Contains = auto()
+        ContainsAny = auto()
+        CountMatches = auto()
+        EndsWith = auto()
+        EscapeRegex = auto()
+        Extract = auto()
+        ExtractAll = auto()
+        ExtractGroups = auto()
+        Find = auto()
+        Head = auto()
+        HexDecode = auto()
+        HexEncode = auto()
+        JsonDecode = auto()
+        JsonPathMatch = auto()
+        LenBytes = auto()
+        LenChars = auto()
+        Lowercase = auto()
+        Normalize = auto()
+        PadEnd = auto()
+        PadStart = auto()
+        Replace = auto()
+        ReplaceMany = auto()
+        Reverse = auto()
+        Slice = auto()
+        Split = auto()
+        SplitExact = auto()
+        SplitN = auto()
+        StartsWith = auto()
+        StripChars = auto()
+        StripCharsEnd = auto()
+        StripCharsStart = auto()
+        StripPrefix = auto()
+        StripSuffix = auto()
+        Strptime = auto()
+        Tail = auto()
+        Titlecase = auto()
+        ToDecimal = auto()
+        ToInteger = auto()
+        Uppercase = auto()
+        ZFill = auto()
+
+        @classmethod
+        def from_polars(cls, obj: pl_expr.StringFunction) -> Self:
+            """Convert from polars' `StringFunction`."""
+            try:
+                function, name = str(obj).split(".", maxsplit=1)
+            except ValueError:
+                # Failed to unpack string
+                function = None
+            if function != "StringFunction":
+                raise ValueError("StringFunction required")
+            return getattr(cls, name)
+
+    __slots__ = ("_regex_program", "name", "options")
     _non_child = ("dtype", "name", "options")
 
     def __init__(
         self,
         dtype: plc.DataType,
-        name: pl_expr.StringFunction,
+        name: StringFunction.Name,
         options: tuple[Any, ...],
         *children: Expr,
     ) -> None:
@@ -43,25 +107,27 @@ class StringFunction(Expr):
         self.options = options
         self.name = name
         self.children = children
+        self.is_pointwise = True
         self._validate_input()
 
     def _validate_input(self):
         if self.name not in (
-            pl_expr.StringFunction.Contains,
-            pl_expr.StringFunction.EndsWith,
-            pl_expr.StringFunction.Lowercase,
-            pl_expr.StringFunction.Replace,
-            pl_expr.StringFunction.ReplaceMany,
-            pl_expr.StringFunction.Slice,
-            pl_expr.StringFunction.Strptime,
-            pl_expr.StringFunction.StartsWith,
-            pl_expr.StringFunction.StripChars,
-            pl_expr.StringFunction.StripCharsStart,
-            pl_expr.StringFunction.StripCharsEnd,
-            pl_expr.StringFunction.Uppercase,
+            StringFunction.Name.ConcatVertical,
+            StringFunction.Name.Contains,
+            StringFunction.Name.EndsWith,
+            StringFunction.Name.Lowercase,
+            StringFunction.Name.Replace,
+            StringFunction.Name.ReplaceMany,
+            StringFunction.Name.Slice,
+            StringFunction.Name.Strptime,
+            StringFunction.Name.StartsWith,
+            StringFunction.Name.StripChars,
+            StringFunction.Name.StripCharsStart,
+            StringFunction.Name.StripCharsEnd,
+            StringFunction.Name.Uppercase,
         ):
-            raise NotImplementedError(f"String function {self.name}")
-        if self.name == pl_expr.StringFunction.Contains:
+            raise NotImplementedError(f"String function {self.name!r}")
+        if self.name is StringFunction.Name.Contains:
             literal, strict = self.options
             if not literal:
                 if not strict:
@@ -82,7 +148,7 @@ class StringFunction(Expr):
                     raise NotImplementedError(
                         f"Unsupported regex {pattern} for GPU engine."
                     ) from e
-        elif self.name == pl_expr.StringFunction.Replace:
+        elif self.name is StringFunction.Name.Replace:
             _, literal = self.options
             if not literal:
                 raise NotImplementedError("literal=False is not supported for replace")
@@ -93,7 +159,7 @@ class StringFunction(Expr):
                 raise NotImplementedError(
                     "libcudf replace does not support empty strings"
                 )
-        elif self.name == pl_expr.StringFunction.ReplaceMany:
+        elif self.name is StringFunction.Name.ReplaceMany:
             (ascii_case_insensitive,) = self.options
             if ascii_case_insensitive:
                 raise NotImplementedError(
@@ -109,12 +175,12 @@ class StringFunction(Expr):
                     "libcudf replace_many is implemented differently from polars "
                     "for empty strings"
                 )
-        elif self.name == pl_expr.StringFunction.Slice:
+        elif self.name is StringFunction.Name.Slice:
             if not all(isinstance(child, Literal) for child in self.children[1:]):
                 raise NotImplementedError(
                     "Slice only supports literal start and stop values"
                 )
-        elif self.name == pl_expr.StringFunction.Strptime:
+        elif self.name is StringFunction.Name.Strptime:
             format, _, exact, cache = self.options
             if cache:
                 raise NotImplementedError("Strptime cache is a CPU feature")
@@ -123,9 +189,9 @@ class StringFunction(Expr):
             if not exact:
                 raise NotImplementedError("Strptime does not support exact=False")
         elif self.name in {
-            pl_expr.StringFunction.StripChars,
-            pl_expr.StringFunction.StripCharsStart,
-            pl_expr.StringFunction.StripCharsEnd,
+            StringFunction.Name.StripChars,
+            StringFunction.Name.StripCharsStart,
+            StringFunction.Name.StripCharsEnd,
         }:
             if not isinstance(self.children[1], Literal):
                 raise NotImplementedError(
@@ -140,7 +206,20 @@ class StringFunction(Expr):
         mapping: Mapping[Expr, Column] | None = None,
     ) -> Column:
         """Evaluate this expression given a dataframe for context."""
-        if self.name == pl_expr.StringFunction.Contains:
+        if self.name is StringFunction.Name.ConcatVertical:
+            (child,) = self.children
+            column = child.evaluate(df, context=context, mapping=mapping)
+            delimiter, ignore_nulls = self.options
+            if column.obj.null_count() > 0 and not ignore_nulls:
+                return Column(plc.Column.all_null_like(column.obj, 1))
+            return Column(
+                plc.strings.combine.join_strings(
+                    column.obj,
+                    plc.interop.from_arrow(pa.scalar(delimiter, type=pa.string())),
+                    plc.interop.from_arrow(pa.scalar(None, type=pa.string())),
+                )
+            )
+        elif self.name is StringFunction.Name.Contains:
             child, arg = self.children
             column = child.evaluate(df, context=context, mapping=mapping)
 
@@ -157,7 +236,7 @@ class StringFunction(Expr):
                 return Column(
                     plc.strings.contains.contains_re(column.obj, self._regex_program)
                 )
-        elif self.name == pl_expr.StringFunction.Slice:
+        elif self.name is StringFunction.Name.Slice:
             child, expr_offset, expr_length = self.children
             assert isinstance(expr_offset, Literal)
             assert isinstance(expr_length, Literal)
@@ -188,16 +267,16 @@ class StringFunction(Expr):
                 )
             )
         elif self.name in {
-            pl_expr.StringFunction.StripChars,
-            pl_expr.StringFunction.StripCharsStart,
-            pl_expr.StringFunction.StripCharsEnd,
+            StringFunction.Name.StripChars,
+            StringFunction.Name.StripCharsStart,
+            StringFunction.Name.StripCharsEnd,
         }:
             column, chars = (
                 c.evaluate(df, context=context, mapping=mapping) for c in self.children
             )
-            if self.name == pl_expr.StringFunction.StripCharsStart:
+            if self.name is StringFunction.Name.StripCharsStart:
                 side = plc.strings.SideType.LEFT
-            elif self.name == pl_expr.StringFunction.StripCharsEnd:
+            elif self.name is StringFunction.Name.StripCharsEnd:
                 side = plc.strings.SideType.RIGHT
             else:
                 side = plc.strings.SideType.BOTH
@@ -207,13 +286,13 @@ class StringFunction(Expr):
             child.evaluate(df, context=context, mapping=mapping)
             for child in self.children
         ]
-        if self.name == pl_expr.StringFunction.Lowercase:
+        if self.name is StringFunction.Name.Lowercase:
             (column,) = columns
             return Column(plc.strings.case.to_lower(column.obj))
-        elif self.name == pl_expr.StringFunction.Uppercase:
+        elif self.name is StringFunction.Name.Uppercase:
             (column,) = columns
             return Column(plc.strings.case.to_upper(column.obj))
-        elif self.name == pl_expr.StringFunction.EndsWith:
+        elif self.name is StringFunction.Name.EndsWith:
             column, suffix = columns
             return Column(
                 plc.strings.find.ends_with(
@@ -223,7 +302,7 @@ class StringFunction(Expr):
                     else suffix.obj,
                 )
             )
-        elif self.name == pl_expr.StringFunction.StartsWith:
+        elif self.name is StringFunction.Name.StartsWith:
             column, prefix = columns
             return Column(
                 plc.strings.find.starts_with(
@@ -233,7 +312,7 @@ class StringFunction(Expr):
                     else prefix.obj,
                 )
             )
-        elif self.name == pl_expr.StringFunction.Strptime:
+        elif self.name is StringFunction.Name.Strptime:
             # TODO: ignores ambiguous
             format, strict, exact, cache = self.options
             col = self.children[0].evaluate(df, context=context, mapping=mapping)
@@ -265,7 +344,7 @@ class StringFunction(Expr):
                         res.columns()[0], self.dtype, format
                     )
                 )
-        elif self.name == pl_expr.StringFunction.Replace:
+        elif self.name is StringFunction.Name.Replace:
             column, target, repl = columns
             n, _ = self.options
             return Column(
@@ -273,7 +352,7 @@ class StringFunction(Expr):
                     column.obj, target.obj_scalar, repl.obj_scalar, maxrepl=n
                 )
             )
-        elif self.name == pl_expr.StringFunction.ReplaceMany:
+        elif self.name is StringFunction.Name.ReplaceMany:
             column, target, repl = columns
             return Column(
                 plc.strings.replace.replace_multiple(column.obj, target.obj, repl.obj)

cudf_polars/dsl/expressions/ternary.py

@@ -34,6 +34,7 @@ class Ternary(Expr):
     ) -> None:
         self.dtype = dtype
         self.children = (when, then, otherwise)
+        self.is_pointwise = True
 
     def do_evaluate(
         self,

cudf_polars/dsl/expressions/unary.py

@@ -21,7 +21,7 @@ if TYPE_CHECKING:
 
     from cudf_polars.containers import DataFrame
 
-__all__ = ["Cast", "UnaryFunction", "Len"]
+__all__ = ["Cast", "Len", "UnaryFunction"]
 
 
 class Cast(Expr):
@@ -33,6 +33,7 @@ class Cast(Expr):
     def __init__(self, dtype: plc.DataType, value: Expr) -> None:
         self.dtype = dtype
         self.children = (value,)
+        self.is_pointwise = True
         if not dtypes.can_cast(value.dtype, self.dtype):
             raise NotImplementedError(
                 f"Can't cast {value.dtype.id().name} to {self.dtype.id().name}"
@@ -63,6 +64,7 @@ class Len(Expr):
     def __init__(self, dtype: plc.DataType) -> None:
         self.dtype = dtype
         self.children = ()
+        self.is_pointwise = False
 
     def do_evaluate(
         self,
@@ -147,6 +149,14 @@ class UnaryFunction(Expr):
         self.name = name
         self.options = options
         self.children = children
+        self.is_pointwise = self.name not in (
+            "cum_min",
+            "cum_max",
+            "cum_prod",
+            "cum_sum",
+            "drop_nulls",
+            "unique",
+        )
 
         if self.name not in UnaryFunction._supported_fns:
             raise NotImplementedError(f"Unary function {name=}")