Fix cudf.polars sum of empty not equalling zero

rapidsai · Jan 6, 2025 · 19a10ba · 19a10ba
1 parent c4f2e8e
commit 19a10ba
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py
@@ -91,7 +91,7 @@ def __init__(
             op = partial(self._reduce, request=req)
         elif name in {"min", "max"}:
             op = partial(op, propagate_nans=options)
-        elif name in {"count", "first", "last"}:
+        elif name in {"count", "sum", "first", "last"}:
             pass
         else:
             raise NotImplementedError(
@@ -180,6 +180,18 @@ def _count(self, column: Column) -> Column:
             )
         )
 
+    def _sum(self, column: Column) -> Column:
+        if column.obj.size() == 0:
+            return Column(
+                plc.Column.from_scalar(
+                    plc.interop.from_arrow(
+                        pa.scalar(0, type=plc.interop.to_arrow(self.dtype))
+                    ),
+                    1,
+                )
+            )
+        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(

diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.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
 from __future__ import annotations
 
@@ -148,3 +148,9 @@ def test_agg_singleton(op):
     q = df.select(op(pl.col("a")))
 
     assert_gpu_result_equal(q)
+
+
+def test_sum_empty_zero():
+    df = pl.LazyFrame({"a": pl.Series(values=[], dtype=pl.Int32())})
+    q = df.select(pl.col("a").sum())
+    assert_gpu_result_equal(q)