Prevent pylibcudf serialization in cudf-polars

python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py

@@ -31,7 +31,7 @@
 
 
 class Agg(Expr):
-    __slots__ = ("name", "op", "options", "request")
+    __slots__ = ("name", "options", "request")
     _non_child = ("dtype", "name", "options")
 
     def __init__(
@@ -46,58 +46,11 @@ def __init__(
             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":
+        if name == "quantile":
             _, quantile = self.children
             if not isinstance(quantile, Literal):
                 raise NotImplementedError("Only support literal quantile values")
-            req = plc.aggregation.quantile(
-                quantiles=[quantile.value.as_py()], 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 in {"count", "sum", "first", "last"}:
-            pass
-        else:
-            raise NotImplementedError(
-                f"Unreachable, supported agg {name=} has no implementation"
-            )  # pragma: no cover
-        self.op = op
+        self.request = None
 
     _SUPPORTED: ClassVar[frozenset[str]] = frozenset(
         [
@@ -124,6 +77,46 @@ def __init__(
         "linear": plc.types.Interpolation.LINEAR,
     }
 
+    def _fill_request(self):
+        if self.request is None:
+            # TODO: nan handling in groupby case
+            if self.name == "min":
+                req = plc.aggregation.min()
+            elif self.name == "max":
+                req = plc.aggregation.max()
+            elif self.name == "median":
+                req = plc.aggregation.median()
+            elif self.name == "n_unique":
+                # TODO: datatype of result
+                req = plc.aggregation.nunique(
+                    null_handling=plc.types.NullPolicy.INCLUDE
+                )
+            elif self.name == "first" or self.name == "last":
+                req = None
+            elif self.name == "mean":
+                req = plc.aggregation.mean()
+            elif self.name == "sum":
+                req = plc.aggregation.sum()
+            elif self.name == "std":
+                # TODO: handle nans
+                req = plc.aggregation.std(ddof=self.options)
+            elif self.name == "var":
+                # TODO: handle nans
+                req = plc.aggregation.variance(ddof=self.options)
+            elif self.name == "count":
+                req = plc.aggregation.count(null_handling=plc.types.NullPolicy.EXCLUDE)
+            elif self.name == "quantile":
+                _, quantile = self.children
+                req = plc.aggregation.quantile(
+                    quantiles=[quantile.value.as_py()],
+                    interp=Agg.interp_mapping[self.options],
+                )
+            else:
+                raise NotImplementedError(
+                    f"Unreachable, {self.name=} is incorrectly listed in _SUPPORTED"
+                )  # pragma: no cover
+            self.request = req
+
     def collect_agg(self, *, depth: int) -> AggInfo:
         """Collect information about aggregations in groupbys."""
         if depth >= 1:
@@ -134,6 +127,7 @@ def collect_agg(self, *, depth: int) -> AggInfo:
             raise NotImplementedError("Nan propagation in groupby for min/max")
         (child,) = self.children
         ((expr, _, _),) = child.collect_agg(depth=depth + 1).requests
+        self._fill_request()
         request = self.request
         # These are handled specially here because we don't set up the
         # request for the whole-frame agg because we can avoid a
@@ -240,7 +234,21 @@ def do_evaluate(
                 f"Agg in context {context}"
             )  # pragma: no cover; unreachable
 
+        self._fill_request()
+
+        op = getattr(self, f"_{self.name}", None)
+        if op is None:
+            op = partial(self._reduce, request=self.request)
+        elif self.name in {"min", "max"}:
+            op = partial(op, propagate_nans=self.options)
+        elif self.name in {"count", "sum", "first", "last"}:
+            pass
+        else:
+            raise NotImplementedError(
+                f"Unreachable, supported agg {self.name=} has no implementation"
+            )  # pragma: no cover
+
         # 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, mapping=mapping))
+        return op(child.evaluate(df, context=context, mapping=mapping))

python/cudf_polars/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
@@ -138,7 +138,7 @@ def _validate_input(self):
                     )
                 pattern = self.children[1].value.as_py()
                 try:
-                    self._regex_program = plc.strings.regex_program.RegexProgram.create(
+                    plc.strings.regex_program.RegexProgram.create(
                         pattern,
                         flags=plc.strings.regex_flags.RegexFlags.DEFAULT,
                     )
@@ -218,6 +218,12 @@ def do_evaluate(
                 )
                 return Column(plc.strings.find.contains(column.obj, pattern))
             else:
+                assert isinstance(arg, Literal)
+                pattern = arg.value.as_py()
+                self._regex_program = plc.strings.regex_program.RegexProgram.create(
+                    pattern,
+                    flags=plc.strings.regex_flags.RegexFlags.DEFAULT,
+                )
                 return Column(
                     plc.strings.contains.contains_re(column.obj, self._regex_program)
                 )

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -883,13 +883,11 @@ def __init__(
             raise NotImplementedError("dynamic group by")
         if any(GroupBy.check_agg(a.value) > 1 for a in self.agg_requests):
             raise NotImplementedError("Nested aggregations in groupby")
-        self.agg_infos = [req.collect_agg(depth=0) for req in self.agg_requests]
         self._non_child_args = (
             self.keys,
             self.agg_requests,
             maintain_order,
             options,
-            self.agg_infos,
         )
 
     @staticmethod
@@ -927,7 +925,6 @@ def do_evaluate(
         agg_requests: Sequence[expr.NamedExpr],
         maintain_order: bool,  # noqa: FBT001
         options: Any,
-        agg_infos: Sequence[expr.AggInfo],
         df: DataFrame,
     ):
         """Evaluate and return a dataframe."""
@@ -947,6 +944,7 @@ def do_evaluate(
         # TODO: uniquify
         requests = []
         replacements: list[expr.Expr] = []
+        agg_infos = [req.collect_agg(depth=0) for req in agg_requests]
         for info in agg_infos:
             for pre_eval, req, rep in info.requests:
                 if pre_eval is None: