Support pl.Expr.product in cudf_polars

Author: mroeschke

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

@@ -118,6 +118,8 @@ def __init__(
             req = plc.aggregation.mean()
         elif name == "sum":
             req = plc.aggregation.sum()
+        elif name == "product":
+            req = plc.aggregation.product()
         elif name == "std":
             # TODO: handle nans
             req = plc.aggregation.std(ddof=options)
@@ -172,7 +174,7 @@ def __init__(
             op = partial(op, propagate_nans=options)
         elif name == "count":
             op = partial(op, include_nulls=options)
-        elif name in {"sum", "first", "last", "item", "first_non_null"}:
+        elif name in {"sum", "product", "first", "last", "item", "first_non_null"}:
             pass
         else:
             raise NotImplementedError(
@@ -194,6 +196,7 @@ def __init__(
             "m2",
             "merge_m2",
             "sum",
+            "product",
             "count",
             "std",
             "var",
@@ -280,6 +283,20 @@ def _sum(self, column: Column, stream: Stream) -> Column:
             )
         return self._reduce(column, request=plc.aggregation.sum(), stream=stream)
 
+    def _product(self, column: Column, stream: Stream) -> Column:
+        if column.size == 0 or column.null_count == column.size:
+            # The product of an empty or all-null column is 1 in polars.
+            return Column(
+                plc.Column.from_scalar(
+                    plc.Scalar.from_py(1, self.dtype.plc_type, stream=stream),
+                    1,
+                    stream=stream,
+                ),
+                name=column.name,
+                dtype=self.dtype,
+            )
+        return self._reduce(column, request=plc.aggregation.product(), stream=stream)
+
     def _min(self, column: Column, *, propagate_nans: bool, stream: Stream) -> Column:
         nan_count = column.nan_count(stream=stream)
         if propagate_nans and nan_count > 0:

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -1002,6 +1002,8 @@ def _(
             )
         elif name == "pow":
             return expr.BinOp(dtype, plc.binaryop.BinaryOperator.POW, *children)
+        elif name == "product":
+            return expr.Agg(dtype, "product", None, translator._expr_context, *children)
         elif not POLARS_VERSION_LT_141 and name == "quantile":
             # polars >= 1.41 emits quantile as a string-named Function
             # expression (function_data=("quantile", interpolation)) with the

python/cudf_polars/cudf_polars/dsl/utils/aggregations.py

@@ -321,7 +321,7 @@ def decompose_single_agg(
         new_children = [child] if not is_quantile else [child, agg.children[1]]
         named_expr = named_expr.reconstruct(agg.reconstruct(new_children))
 
-        if agg.name == "sum":
+        if agg.name in {"sum", "product"}:
             col = (
                 expr.Cast(
                     agg.dtype,
@@ -339,8 +339,10 @@ def decompose_single_agg(
             # - ROLLING: sum(all-null window) => null; sum(empty window) => 0 (fill only if empty)
             #
             # Must post-process because libcudf returns null for both empty and all-null windows/groups
+            # product uses an identity of 1 (empty/all-null product is 1 in polars).
+            identity = 1 if agg.name == "product" else 0
             return [(named_expr, True)], expr.NamedExpr(
-                name, replace_nulls(col, 0, is_top=is_top)
+                name, replace_nulls(col, identity, is_top=is_top)
             )
         elif agg.name in {"mean", "median", "quantile", "std", "var"}:
             post_agg_col: expr.Expr = expr.Col(

python/cudf_polars/tests/expressions/test_agg.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
@@ -112,6 +112,26 @@ def test_bool_agg(engine: pl.GPUEngine, agg):
     assert_gpu_result_equal(q, engine=engine, check_exact=False)
 
 
+@pytest.mark.parametrize(
+    "data,dtype",
+    [
+        ([1, 2, 3, 4], pl.Int32),
+        ([1, 2, None, 4], pl.Int32),
+        ([1, 0, 3], pl.Int32),
+        ([2, 3, 4], pl.Int8),
+        ([1.5, 2.0, 3.0], pl.Float64),
+        ([1.5, None, 3.0], pl.Float64),
+        ([True, False, True], pl.Boolean),
+        ([], pl.Int32),
+        ([None, None], pl.Int32),
+    ],
+)
+def test_product(engine: pl.GPUEngine, data, dtype):
+    df = pl.LazyFrame({"a": pl.Series(data, dtype=dtype)})
+    q = df.select(pl.col("a").product())
+    assert_gpu_result_equal(q, engine=engine, check_exact=False)
+
+
 @pytest.mark.parametrize("cum_agg", sorted(expr.UnaryFunction._supported_cum_aggs))
 def test_cum_agg_reverse_unsupported(engine: pl.GPUEngine, cum_agg):
     df = pl.LazyFrame({"a": [1, 2, 3]})