Repository: narwhals-dev/narwhals
Pull request #429 (2024-07-07 10:51:34)

feat: implement DataFrame.with_columns for arrow backend
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 14a08050..c0869542 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -7,6 +7,7 @@ from typing import Sequence
 from typing import overload
 
 from narwhals._arrow.utils import translate_dtype
+from narwhals._arrow.utils import validate_dataframe_comparand
 from narwhals._pandas_like.utils import evaluate_into_exprs
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pyarrow
@@ -127,6 +128,34 @@ class ArrowDataFrame:
         df = pa.Table.from_arrays([s._native_series for s in new_series], names=names)
         return self._from_native_dataframe(df)
 
+    def with_columns(
+        self,
+        *exprs: IntoArrowExpr,
+        **named_exprs: IntoArrowExpr,
+    ) -> Self:
+        new_columns = evaluate_into_exprs(self, *exprs, **named_exprs)  # type: ignore[arg-type]
+        new_column_name_to_new_column_map = {s.name: s for s in new_columns}
+        to_concat = []
+        output_names = []
+        # Make sure to preserve column order
+        for name in self.columns:
+            if name in new_column_name_to_new_column_map:
+                to_concat.append(
+                    validate_dataframe_comparand(
+                        new_column_name_to_new_column_map.pop(name)
+                    )
+                )
+            else:
+                to_concat.append(self._native_dataframe[name])
+            output_names.append(name)
+        for s in new_column_name_to_new_column_map:
+            to_concat.append(
+                validate_dataframe_comparand(new_column_name_to_new_column_map[s])
+            )
+            output_names.append(s)
+        df = self._native_dataframe.__class__.from_arrays(to_concat, names=output_names)
+        return self._from_native_dataframe(df)
+
     def drop(self, *columns: str | Iterable[str]) -> Self:
         return self._from_native_dataframe(
             self._native_dataframe.drop(list(flatten(columns)))
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 1196ec62..5c74daca 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from typing import Any
 from typing import Callable
 
 from narwhals._pandas_like.utils import reuse_series_implementation
@@ -74,12 +75,30 @@ class ArrowExpr:
 
         return ArrowNamespace(backend_version=self._backend_version)
 
+    def __add__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__add__", other)  # type: ignore[type-var]
+
+    def __sub__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__sub__", other)  # type: ignore[type-var]
+
+    def __mul__(self, other: ArrowExpr | Any) -> Self:
+        return reuse_series_implementation(self, "__mul__", other)  # type: ignore[type-var]
+
+    def mean(self) -> Self:
+        return reuse_series_implementation(self, "mean", returns_scalar=True)  # type: ignore[type-var]
+
+    def std(self, ddof: int = 1) -> Self:
+        return reuse_series_implementation(self, "std", ddof=ddof, returns_scalar=True)  # type: ignore[type-var]
+
     def cast(self, dtype: DType) -> Self:
         return reuse_series_implementation(self, "cast", dtype)  # type: ignore[type-var]
 
     def abs(self) -> Self:
         return reuse_series_implementation(self, "abs")  # type: ignore[type-var]
 
+    def diff(self) -> Self:
+        return reuse_series_implementation(self, "diff")  # type: ignore[type-var]
+
     def cum_sum(self) -> Self:
         return reuse_series_implementation(self, "cum_sum")  # type: ignore[type-var]
 
@@ -89,6 +108,18 @@ class ArrowExpr:
     def all(self) -> Self:
         return reuse_series_implementation(self, "all", returns_scalar=True)  # type: ignore[type-var]
 
+    def alias(self, name: str) -> Self:
+        # Define this one manually, so that we can
+        # override `output_names` and not increase depth
+        return self.__class__(
+            lambda df: [series.alias(name) for series in self._call(df)],
+            depth=self._depth,
+            function_name=self._function_name,
+            root_names=self._root_names,
+            output_names=[name],
+            backend_version=self._backend_version,
+        )
+
     @property
     def dt(self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index fca21483..8359e684 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -5,8 +5,10 @@ from typing import Any
 from typing import Iterable
 
 from narwhals._arrow.namespace import ArrowNamespace
+from narwhals._arrow.utils import item
 from narwhals._arrow.utils import reverse_translate_dtype
 from narwhals._arrow.utils import translate_dtype
+from narwhals._arrow.utils import validate_column_comparand
 from narwhals._pandas_like.utils import native_series_from_iterable
 from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
@@ -55,6 +57,29 @@ class ArrowSeries:
     def __len__(self) -> int:
         return len(self._native_series)
 
+    def __add__(self, other: Any) -> Self:
+        pc = get_pyarrow_compute()
+        other = validate_column_comparand(other)
+        return self._from_native_series(pc.add(self._native_series, other))
+
+    def __sub__(self, other: Any) -> Self:
+        pc = get_pyarrow_compute()
+        other = validate_column_comparand(other)
+        return self._from_native_series(pc.subtract(self._native_series, other))
+
+    def __mul__(self, other: Any) -> Self:
+        pc = get_pyarrow_compute()
+        other = validate_column_comparand(other)
+        return self._from_native_series(pc.multiply(self._native_series, other))
+
+    def mean(self) -> int:
+        pc = get_pyarrow_compute()
+        return item(self._backend_version, pc.mean(self._native_series))  # type: ignore[no-any-return]
+
+    def std(self, ddof: int = 1) -> int:
+        pc = get_pyarrow_compute()
+        return item(self._backend_version, pc.stddev(self._native_series, ddof=ddof))  # type: ignore[no-any-return]
+
     def __narwhals_namespace__(self) -> ArrowNamespace:
         return ArrowNamespace(backend_version=self._backend_version)
 
@@ -96,6 +121,20 @@ class ArrowSeries:
         pc = get_pyarrow_compute()
         return self._from_native_series(pc.cumulative_sum(self._native_series))
 
+    def diff(self) -> Self:
+        pc = get_pyarrow_compute()
+        pa = get_pyarrow()
+        # todo: is all this rechunking really necessary?
+        ca = self._native_series.combine_chunks()
+        result = pa.chunked_array(
+            [
+                pa.chunked_array(
+                    [pa.array([None], type=ca.type), pc.subtract(ca[1:], ca[:-1])]
+                ).combine_chunks()
+            ]
+        )
+        return self._from_native_series(result)
+
     def any(self) -> bool:
         pc = get_pyarrow_compute()
         return pc.any(self._native_series)  # type: ignore[no-any-return]
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index fe60e2fd..55e3c415 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -93,3 +93,57 @@ def reverse_translate_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any:
         return pa.date32()
     msg = f"Unknown dtype: {dtype}"  # pragma: no cover
     raise AssertionError(msg)
+
+
+def validate_column_comparand(other: Any) -> Any:
+    """Validate RHS of binary operation.
+
+    If the comparison isn't supported, return `NotImplemented` so that the
+    "right-hand-side" operation (e.g. `__radd__`) can be tried.
+
+    If RHS is length 1, return the scalar value, so that the underlying
+    library can broadcast it.
+    """
+    from narwhals._arrow.dataframe import ArrowDataFrame
+    from narwhals._arrow.series import ArrowSeries
+
+    if isinstance(other, list):
+        if len(other) > 1:
+            # e.g. `plx.all() + plx.all()`
+            msg = "Multi-output expressions are not supported in this context"
+            raise ValueError(msg)
+        other = other[0]
+    if isinstance(other, ArrowDataFrame):
+        return NotImplemented
+    if isinstance(other, ArrowSeries):
+        if len(other) == 1:
+            # broadcast
+            return other[0]
+        return other._native_series
+    return other
+
+
+def validate_dataframe_comparand(other: Any) -> Any:
+    """Validate RHS of binary operation.
+
+    If the comparison isn't supported, return `NotImplemented` so that the
+    "right-hand-side" operation (e.g. `__radd__`) can be tried.
+    """
+    from narwhals._arrow.dataframe import ArrowDataFrame
+    from narwhals._arrow.series import ArrowSeries
+
+    if isinstance(other, ArrowDataFrame):
+        return NotImplemented
+    if isinstance(other, ArrowSeries):
+        if len(other) == 1:
+            # broadcast
+            msg = "not implemented yet"  # pragma: no cover
+            raise NotImplementedError(msg)
+        return other._native_series
+    raise AssertionError("Please report a bug")
+
+
+def item(backend_version: tuple[int, ...], obj: Any) -> Any:
+    if backend_version > (13,):
+        return obj
+    return obj.as_py()  # pragma: no cover
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index e2e52f27..21fc954f 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -92,10 +92,11 @@ def validate_dataframe_comparand(index: Any, other: Any) -> Any:
 
 def maybe_evaluate_expr(df: PandasDataFrame, expr: Any) -> Any:
     """Evaluate `expr` if it's an expression, otherwise return it as is."""
+    from narwhals._arrow.expr import ArrowExpr
     from narwhals._pandas_like.expr import PandasExpr
 
-    if isinstance(expr, PandasExpr):
-        return expr._call(df)
+    if isinstance(expr, (PandasExpr, ArrowExpr)):
+        return expr._call(df)  # type: ignore[arg-type]
     return expr
 
 
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index 18d2e6d9..6a9efad6 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -28,10 +28,8 @@ MISSING = [
     "DataFrame.to_dict",
     "DataFrame.to_numpy",
     "DataFrame.unique",
-    "DataFrame.with_columns",
     "DataFrame.with_row_index",
     "DataFrame.write_parquet",
-    "Series.diff",
     "Series.drop_nulls",
     "Series.fill_null",
     "Series.filter",
@@ -48,7 +46,6 @@ MISSING = [
     "Series.item",
     "Series.len",
     "Series.max",
-    "Series.mean",
     "Series.min",
     "Series.n_unique",
     "Series.null_count",
@@ -57,7 +54,6 @@ MISSING = [
     "Series.sample",
     "Series.shift",
     "Series.sort",
-    "Series.std",
     "Series.str",
     "Series.sum",
     "Series.tail",
@@ -105,6 +101,7 @@ class MockSeries:
 
 if __name__ == "__main__":
     missing = []
+    no_longer_missing = []
 
     df_pa = ArrowDataFrame(MockDataFrame({"a": [1, 2, 3]}), backend_version=(13, 0))
     df_pd = nw.DataFrame(
@@ -113,12 +110,14 @@ if __name__ == "__main__":
     pa_methods = [f"DataFrame.{x}" for x in df_pa.__dir__() if not x.startswith("_")]
     pd_methods = [f"DataFrame.{x}" for x in df_pd.__dir__() if not x.startswith("_")]
     missing.extend([x for x in pd_methods if x not in pa_methods and x not in MISSING])
+    no_longer_missing.extend([x for x in MISSING if x in pa_methods and x in pd_methods])
 
     ser_pa = df_pa["a"]
     ser_pd = df_pd["a"]
     pa_methods = [f"Series.{x}" for x in ser_pa.__dir__() if not x.startswith("_")]
     pd_methods = [f"Series.{x}" for x in ser_pd.__dir__() if not x.startswith("_")]
     missing.extend([x for x in pd_methods if x not in pa_methods and x not in MISSING])
+    no_longer_missing.extend([x for x in MISSING if x in pa_methods and x in pd_methods])
 
     if missing:
         print(
@@ -127,7 +126,6 @@ if __name__ == "__main__":
         )
         sys.exit(1)
 
-    no_longer_missing = [x for x in MISSING if x in pa_methods and x in pd_methods]
     if no_longer_missing:
         print(
             "Please remove the following from MISSING in utils/check_backend_completeness.py: ",


Test patch:
diff --git a/tests/expr/diff_test.py b/tests/expr/diff_test.py
index c4189dd5..db0784ce 100644
--- a/tests/expr/diff_test.py
+++ b/tests/expr/diff_test.py
@@ -10,9 +10,9 @@ data = {
 }
 
 
-def test_diff(constructor: Any) -> None:
-    df = nw.from_native(constructor(data), eager_only=True)
-    result = df.with_columns(c_diff=nw.col("c").diff()).filter(nw.col("i") > 0)
+def test_diff(constructor_with_pyarrow: Any) -> None:
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = df.with_columns(c_diff=nw.col("c").diff())[1:]
     expected = {
         "i": [1, 2, 3, 4],
         "b": [2, 3, 5, 3],
@@ -20,5 +20,5 @@ def test_diff(constructor: Any) -> None:
         "c_diff": [-1, -1, -1, -1],
     }
     compare_dicts(result, expected)
-    result = df.with_columns(c_diff=df["c"].diff()).filter(nw.col("i") > 0)
+    result = df.with_columns(c_diff=df["c"].diff())[1:]
     compare_dicts(result, expected)
diff --git a/tests/expr/double_test.py b/tests/expr/double_test.py
index 1495195c..d83535fd 100644
--- a/tests/expr/double_test.py
+++ b/tests/expr/double_test.py
@@ -4,14 +4,17 @@ import narwhals as nw
 from tests.utils import compare_dicts
 
 
-def test_double(constructor: Any) -> None:
+def test_double(constructor_with_pyarrow: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df = nw.from_native(constructor(data), eager_only=True)
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
     result = df.with_columns(nw.all() * 2)
-    result_native = nw.to_native(result)
     expected = {"a": [2, 6, 4], "b": [8, 8, 12], "z": [14.0, 16.0, 18.0]}
-    compare_dicts(result_native, expected)
+    compare_dicts(result, expected)
+
+
+def test_double_alias(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
     result = df.with_columns(nw.col("a").alias("o"), nw.all() * 2)
-    result_native = nw.to_native(result)
     expected = {"o": [1, 3, 2], "a": [2, 6, 4], "b": [8, 8, 12], "z": [14.0, 16.0, 18.0]}
-    compare_dicts(result_native, expected)
+    compare_dicts(result, expected)
diff --git a/tests/frame/add_test.py b/tests/frame/add_test.py
new file mode 100644
index 00000000..405dbacb
--- /dev/null
+++ b/tests/frame/add_test.py
@@ -0,0 +1,23 @@
+from typing import Any
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+
+def test_add(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+    df = nw.from_native(constructor_with_pyarrow(data))
+    result = df.with_columns(
+        c=nw.col("a") + nw.col("b"),
+        d=nw.col("a") - nw.col("a").mean(),
+        e=nw.col("a") - nw.col("a").std(),
+    )
+    expected = {
+        "a": [1, 3, 2],
+        "b": [4, 4, 6],
+        "z": [7.0, 8.0, 9.0],
+        "c": [5, 7, 8],
+        "d": [-1.0, 1.0, 0.0],
+        "e": [0.0, 2.0, 1.0],
+    }
+    compare_dicts(result, expected)
diff --git a/tests/frame/test_common.py b/tests/frame/test_common.py
index a352812c..d9f3578f 100644
--- a/tests/frame/test_common.py
+++ b/tests/frame/test_common.py
@@ -68,29 +68,6 @@ def test_empty_select(constructor: Any) -> None:
     assert result.shape == (0, 0)
 
 
-@pytest.mark.parametrize(
-    "df_raw",
-    [df_pandas, df_lazy, df_pandas_nullable, df_pandas_pyarrow],
-)
-def test_add(df_raw: Any) -> None:
-    df = nw.from_native(df_raw)
-    result = df.with_columns(
-        c=nw.col("a") + nw.col("b"),
-        d=nw.col("a") - nw.col("a").mean(),
-        e=nw.col("a") - nw.col("a").std(),
-    )
-    result_native = nw.to_native(result)
-    expected = {
-        "a": [1, 3, 2],
-        "b": [4, 4, 6],
-        "z": [7.0, 8.0, 9.0],
-        "c": [5, 7, 8],
-        "d": [-1.0, 1.0, 0.0],
-        "e": [0.0, 2.0, 1.0],
-    }
-    compare_dicts(result_native, expected)
-
-
 @pytest.mark.parametrize(
     "df_raw",
     [df_pandas, df_lazy, df_pandas_nullable, df_pandas_pyarrow],
@@ -542,6 +519,9 @@ def test_lazy(df_raw: Any) -> None:
 
 
 def test_invalid() -> None:
+    df = nw.from_native(pa.table({"a": [1, 2], "b": [3, 4]}))
+    with pytest.raises(ValueError, match="Multi-output"):
+        df.select(nw.all() + nw.all())
     df = nw.from_native(df_pandas)
     with pytest.raises(ValueError, match="Multi-output"):
         df.select(nw.all() + nw.all())