Added arrow sum method (#662)

narwhals-dev · Jul 28, 2024 · 8d2cbb0 · 8d2cbb0
1 parent 435dbea
commit 8d2cbb0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
@@ -188,6 +188,12 @@ def is_between(
             closed,
         )
 
+    def sum(self) -> Self:
+        return self._from_call(
+            lambda _input: _input.sum(),
+            "sum",
+        )
+
     @property
     def str(self: Self) -> DaskExprStringNamespace:
         return DaskExprStringNamespace(self)

diff --git a/tests/dask_test.py b/tests/dask_test.py
@@ -83,6 +83,16 @@ def test_cum_sum() -> None:
     compare_dicts(result, expected)
 
 
+def test_sum() -> None:
+    import dask.dataframe as dd
+
+    dfdd = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    df = nw.from_native(dfdd)
+    result = df.with_columns((nw.col("a") + nw.col("b").sum()).alias("c"))
+    expected = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [16, 17, 18]}
+    compare_dicts(result, expected)
+
+
 @pytest.mark.parametrize(
     ("closed", "expected"),
     [