feat: Add coalesce for all backends (#2647)

camriddell · FBruzzesi · web-flow · commit 6bf173100b04 · 2025-07-07T08:06:52.000-07:00
* add coalesce for all backends

* add coalesce docstring, stable version, and protocol info

* add simple coalesce v1 test

* fix coalesce to appropriately cast literals

* coalesce mark backends that do not upcast as xfail

* coalesce to accept literals, not unwrapped scalars

* fix up coalesce v1 typing

* lazy coalesce backends use _with_elementwise

* remove test coalesce upcast

* enh coalesce tests explicit &amp; implicit lit

* enh coalesce push literal conversion to backends

* use self._series._align_full_broadcast

* replace _with_elementwise

* remove coalesce manual literal guessing

* Enh raise in coalesce non {str,lit,Series} inputs

- add coalesce test for series inputs
- add coalesce test for raising TypeError from inappropriate inputs
- remove coalesce tests for scalar inputs, only accept valid IntoExpr

* ref add in upstream elementwise changes

* remove dask backend_version from coalesce

* clean up _pandas_like namespace diff

* cleanup coalesce ibis impl

* enh coalesce docstring and error message

---------

Co-authored-by: Francesco Bruzzesi &lt;42817048+FBruzzesi@users.noreply.github.com&gt;
Co-authored-by: FBruzzesi &lt;francesco.bruzzesi.93@gmail.com&gt;
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
@@ -9,6 +9,7 @@ Here are the top-level functions available in Narwhals.
         - all
         - all_horizontal
         - any_horizontal
+        - coalesce
         - col
         - concat
         - concat_str
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
@@ -49,6 +49,7 @@
     all_ as all,
     all_horizontal,
     any_horizontal,
+    coalesce,
     col,
     concat,
     concat_str,
@@ -127,6 +128,7 @@
     "all",
     "all_horizontal",
     "any_horizontal",
+    "coalesce",
     "col",
     "concat",
     "concat_str",
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
@@ -262,6 +262,27 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             context=self,
         )
 
+    def coalesce(self, *exprs: ArrowExpr) -> ArrowExpr:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            align = self._series._align_full_broadcast
+            init_series, *series = align(*chain.from_iterable(expr(df) for expr in exprs))
+            return [
+                ArrowSeries(
+                    pc.coalesce(init_series.native, *(s.native for s in series)),
+                    name=init_series.name,
+                    version=self._version,
+                )
+            ]
+
+        return self._expr._from_callable(
+            func=func,
+            depth=max(x._depth for x in exprs) + 1,
+            function_name="coalesce",
+            evaluate_output_names=combine_evaluate_output_names(*exprs),
+            alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
+        )
+
 
 class ArrowWhen(EagerWhen[ArrowDataFrame, ArrowSeries, ArrowExpr, "ChunkedArrayAny"]):
     @property
diff --git a/narwhals/_compliant/namespace.py b/narwhals/_compliant/namespace.py
@@ -96,6 +96,7 @@ def concat_str(
     def selectors(self) -> CompliantSelectorNamespace[Any, Any]: ...
     @property
     def _expr(self) -> type[CompliantExprT]: ...
+    def coalesce(self, *exprs: CompliantExprT) -> CompliantExprT: ...
 
 
 class DepthTrackingNamespace(
diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py
@@ -286,6 +286,22 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
             version=self._version,
         )
 
+    def coalesce(self, *exprs: DaskExpr) -> DaskExpr:
+        def func(df: DaskLazyFrame) -> list[dx.Series]:
+            series = align_series_full_broadcast(
+                df, *(s for _expr in exprs for s in _expr(df))
+            )
+            return [reduce(lambda x, y: x.fillna(y), series)]
+
+        return self._expr(
+            call=func,
+            depth=max(x._depth for x in exprs) + 1,
+            function_name="coalesce",
+            evaluate_output_names=combine_evaluate_output_names(*exprs),
+            alias_output_names=combine_alias_output_names(*exprs),
+            version=self._version,
+        )
+
 
 class DaskWhen(CompliantWhen[DaskLazyFrame, "dx.Series", DaskExpr]):
     @property
diff --git a/narwhals/_duckdb/namespace.py b/narwhals/_duckdb/namespace.py
@@ -174,6 +174,12 @@ def func(_df: DuckDBLazyFrame) -> list[Expression]:
             version=self._version,
         )
 
+    def coalesce(self, *exprs: DuckDBExpr) -> DuckDBExpr:
+        def func(cols: Iterable[Expression]) -> Expression:
+            return CoalesceOperator(*cols)
+
+        return self._expr._from_elementwise_horizontal_op(func, *exprs)
+
 
 class DuckDBWhen(LazyWhen["DuckDBLazyFrame", Expression, DuckDBExpr]):
     @property
diff --git a/narwhals/_ibis/namespace.py b/narwhals/_ibis/namespace.py
@@ -160,6 +160,12 @@ def func(_df: IbisLazyFrame) -> list[ir.Value]:
             version=self._version,
         )
 
+    def coalesce(self, *exprs: IbisExpr) -> IbisExpr:
+        def func(cols: Iterable[ir.Value]) -> ir.Value:
+            return ibis.coalesce(*cols)
+
+        return self._expr._from_elementwise_horizontal_op(func, *exprs)
+
 
 class IbisWhen(LazyWhen["IbisLazyFrame", "ir.Value", IbisExpr]):
     lit = lit
diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py
@@ -68,6 +68,23 @@ def __init__(self, implementation: Implementation, version: Version) -> None:
         self._implementation = implementation
         self._version = version
 
+    def coalesce(self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
+        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+            align = self._series._align_full_broadcast
+            series = align(*(s for _expr in exprs for s in _expr(df)))
+            return [
+                reduce(lambda x, y: x.fill_null(y, strategy=None, limit=None), series)
+            ]
+
+        return self._expr._from_callable(
+            func=func,
+            depth=max(x._depth for x in exprs) + 1,
+            function_name="coalesce",
+            evaluate_output_names=combine_evaluate_output_names(*exprs),
+            alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
+        )
+
     def lit(self, value: NonNestedLiteral, dtype: IntoDType | None) -> PandasLikeExpr:
         def _lit_pandas_series(df: PandasLikeDataFrame) -> PandasLikeSeries:
             pandas_series = self._series.from_iterable(
diff --git a/narwhals/_polars/namespace.py b/narwhals/_polars/namespace.py
@@ -26,6 +26,7 @@
 
 class PolarsNamespace:
     all: Method[PolarsExpr]
+    coalesce: Method[PolarsExpr]
     col: Method[PolarsExpr]
     exclude: Method[PolarsExpr]
     sum_horizontal: Method[PolarsExpr]
diff --git a/narwhals/_spark_like/namespace.py b/narwhals/_spark_like/namespace.py
@@ -236,6 +236,12 @@ def func(df: SparkLikeLazyFrame) -> list[Column]:
     def when(self, predicate: SparkLikeExpr) -> SparkLikeWhen:
         return SparkLikeWhen.from_expr(predicate, context=self)
 
+    def coalesce(self, *exprs: SparkLikeExpr) -> SparkLikeExpr:
+        def func(cols: Iterable[Column]) -> Column:
+            return self._F.coalesce(*cols)
+
+        return self._expr._from_elementwise_horizontal_op(func, *exprs)
+
 
 class SparkLikeWhen(LazyWhen[SparkLikeLazyFrame, "Column", SparkLikeExpr]):
     @property
diff --git a/narwhals/functions.py b/narwhals/functions.py
@@ -36,6 +36,7 @@
 )
 from narwhals.exceptions import InvalidOperationError
 from narwhals.expr import Expr
+from narwhals.series import Series
 from narwhals.translate import from_native, to_native
 
 if TYPE_CHECKING:
@@ -48,7 +49,6 @@
     from narwhals.dataframe import DataFrame, LazyFrame
     from narwhals.dtypes import DType
     from narwhals.schema import Schema
-    from narwhals.series import Series
     from narwhals.typing import (
         ConcatMethod,
         FrameT,
@@ -1847,3 +1847,69 @@ def concat_str(
             *flat_exprs, str_as_lit=False, allow_multi_output=True, to_single_output=True
         ),
     )
+
+
+def coalesce(
+    exprs: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr | NonNestedLiteral
+) -> Expr:
+    """Folds the columns from left to right, keeping the first non-null value.
+
+    Arguments:
+        exprs: Columns to coalesce, must be a str, nw.Expr, or nw.Series
+            where strings are parsed as column names and both nw.Expr/nw.Series
+            are passed through as-is. Scalar values must be wrapped in `nw.lit`.
+
+        *more_exprs: Additional columns to coalesce, specified as positional arguments.
+
+    Raises:
+        TypeError: If any of the inputs are not a str, nw.Expr, or nw.Series.
+
+    Returns:
+        A new expression.
+
+    Examples:
+    >>> import polars as pl
+    >>> import narwhals as nw
+    >>> data = [
+    ...     (1, 5, None),
+    ...     (None, 6, None),
+    ...     (None, None, 9),
+    ...     (4, 8, 10),
+    ...     (None, None, None),
+    ... ]
+    >>> df = pl.DataFrame(data, schema=["a", "b", "c"], orient="row")
+    >>> nw.from_native(df).select(nw.coalesce("a", "b", "c", nw.lit(-1)))
+    ┌──────────────────┐
+    |Narwhals DataFrame|
+    |------------------|
+    |  shape: (5, 1)   |
+    |  ┌─────┐         |
+    |  │ a   │         |
+    |  │ --- │         |
+    |  │ i64 │         |
+    |  ╞═════╡         |
+    |  │ 1   │         |
+    |  │ 6   │         |
+    |  │ 9   │         |
+    |  │ 4   │         |
+    |  │ -1  │         |
+    |  └─────┘         |
+    └──────────────────┘
+    """
+    flat_exprs = flatten([*flatten([exprs]), *more_exprs])
+
+    non_exprs = [expr for expr in flat_exprs if not isinstance(expr, (str, Expr, Series))]
+    if non_exprs:
+        msg = (
+            f"All arguments to `coalesce` must be of type {str!r}, {Expr!r}, or {Series!r}."
+            "\nGot the following invalid arguments (type, value):"
+            f"\n    {', '.join(repr((type(e), e)) for e in non_exprs)}"
+        )
+        raise TypeError(msg)
+
+    return Expr(
+        lambda plx: apply_n_ary_operation(
+            plx, lambda *args: plx.coalesce(*args), *flat_exprs, str_as_lit=False
+        ),
+        ExprMetadata.from_horizontal_op(*flat_exprs),
+    )
diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py
@@ -1448,6 +1448,25 @@ def concat_str(
     )
 
 
+def coalesce(exprs: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr) -> Expr:
+    """Folds the columns from left to right, keeping the first non-null value.
+
+    Arguments:
+        exprs: Columns to coalesce, must be a str, nw.Expr, or nw.Series
+            where strings are parsed as column names and both nw.Expr/nw.Series
+            are passed through as-is. Scalar values must be wrapped in `nw.lit`.
+
+        *more_exprs: Additional columns to coalesce, specified as positional arguments.
+
+    Raises:
+        TypeError: If any of the inputs are not a str, nw.Expr, or nw.Series.
+
+    Returns:
+        A new expression.
+    """
+    return _stableify(nw.coalesce(exprs, *more_exprs))
+
+
 def get_level(
     obj: DataFrame[Any] | LazyFrame[Any] | Series[IntoSeriesT],
 ) -> Literal["full", "lazy", "interchange"]:
@@ -1874,6 +1893,7 @@ def scan_parquet(
     "all",
     "all_horizontal",
     "any_horizontal",
+    "coalesce",
     "col",
     "concat",
     "concat_str",
diff --git a/tests/expr_and_series/coalesce_test.py b/tests/expr_and_series/coalesce_test.py
@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals as nw
+from tests.utils import Constructor, ConstructorEager, assert_equal_data
+
+
+def test_coalesce_numeric(constructor: Constructor) -> None:
+    data = {
+        "a": [0, None, None, None, None],
+        "b": [1, None, None, 5, 3],
+        "c": [5, None, 3, 2, 1],
+    }
+    df = nw.from_native(constructor(data))
+
+    result = df.select(
+        no_lit=nw.coalesce("a", "b", "c"),
+        explicit_lit=nw.coalesce("a", "b", "c", nw.lit(-1)),
+    )
+    expected = {"no_lit": [0, None, 3, 5, 3], "explicit_lit": [0, -1, 3, 5, 3]}
+    assert_equal_data(result, expected)
+
+
+def test_coalesce_strings(constructor: Constructor) -> None:
+    data = {
+        "a": ["0", None, None, None, None],
+        "b": ["1", None, None, "5", "3"],
+        "c": ["5", None, "3", "2", "1"],
+    }
+    df = nw.from_native(constructor(data))
+
+    result = df.select(
+        no_lit=nw.coalesce("a", "b", "c"),
+        explicit_lit=nw.coalesce("a", "b", "c", nw.lit("xyz")),
+    )
+    expected = {
+        "no_lit": ["0", None, "3", "5", "3"],
+        "explicit_lit": ["0", "xyz", "3", "5", "3"],
+    }
+    assert_equal_data(result, expected)
+
+
+def test_coalesce_series(constructor_eager: ConstructorEager) -> None:
+    data = {
+        "a": ["0", None, None, None, None],
+        "b": ["1", None, None, "5", "3"],
+        "c": ["5", None, "3", "2", "1"],
+    }
+    df = nw.from_native(constructor_eager(data))
+
+    result = df.select(result=nw.coalesce("a", "b", df["c"]))
+    expected = {"result": ["0", None, "3", "5", "3"]}
+    assert_equal_data(result, expected)
+
+
+def test_coalesce_raises_non_expr(constructor: Constructor) -> None:
+    data = {
+        "a": ["0", None, None, None, None],
+        "b": ["1", None, None, "5", "3"],
+        "c": ["5", None, "3", "2", "1"],
+    }
+    df = nw.from_native(constructor(data))
+
+    with pytest.raises(TypeError, match="All arguments to `coalesce` must be of type"):
+        df.select(implicit_lit=nw.coalesce("a", "b", "c", 10))
diff --git a/tests/v1_test.py b/tests/v1_test.py
@@ -46,7 +46,9 @@
 def test_toplevel(constructor_eager: ConstructorEager) -> None:
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1,):
         pytest.skip()
-    df = nw_v1.from_native(constructor_eager({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    df = nw_v1.from_native(
+        constructor_eager({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, None, 9]})
+    )
     result = df.select(
         min=nw_v1.min("a"),
         max=nw_v1.max("a"),
@@ -62,7 +64,8 @@ def test_toplevel(constructor_eager: ConstructorEager) -> None:
         any_h=nw_v1.any_horizontal(nw_v1.lit(True), nw_v1.lit(True)),  # noqa: FBT003
         all_h=nw_v1.all_horizontal(nw_v1.lit(True), nw_v1.lit(True)),  # noqa: FBT003
         first=nw_v1.nth(0),
-        no_first=nw_v1.exclude("a"),
+        no_first=nw_v1.exclude("a", "c"),
+        coalesce=nw_v1.coalesce("c", "a"),
     )
     expected = {
         "min": [1, 1, 1],
@@ -80,6 +83,7 @@ def test_toplevel(constructor_eager: ConstructorEager) -> None:
         "all_h": [True, True, True],
         "first": [1, 2, 3],
         "no_first": [4, 5, 6],
+        "coalesce": [7, 2, 9],
     }
     assert_equal_data(result, expected)
     assert isinstance(result, nw_v1.DataFrame)
