simpler, top-level solution

Author: MarcoGorelli

narwhals/_arrow/namespace.py

@@ -247,23 +247,9 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
 
     def _if_then_else(
         self,
-        when: ChunkedArrayAny | list[ChunkedArrayAny],
-        then: ChunkedArrayAny | list[ChunkedArrayAny],
+        when: ChunkedArrayAny,
+        then: ChunkedArrayAny,
         otherwise: ChunkedArrayAny | None = None,
     ) -> ChunkedArrayAny:
-        if not isinstance(when, list):
-            otherwise = pa.nulls(len(when), then.type) if otherwise is None else otherwise
-            return pc.if_else(when, then, otherwise)
-
-        conditions = when
-        values = then
-
-        if otherwise is None:
-            result = pa.nulls(len(conditions[-1]), values[-1].type)
-        else:
-            result = otherwise
-
-        for cond, val in reversed(list(zip(conditions, values))):
-            result = pc.if_else(cond, val, result)
-
-        return result
+        otherwise = pa.nulls(len(when), then.type) if otherwise is None else otherwise
+        return pc.if_else(when, then, otherwise)

narwhals/_compliant/namespace.py

@@ -201,69 +201,40 @@ def _dataframe(self) -> type[EagerDataFrameT]: ...
     def _series(self) -> type[EagerSeriesT_co]: ...
     def _if_then_else(
         self,
-        when: NativeSeriesT | list[NativeSeriesT],
-        then: NativeSeriesT | list[NativeSeriesT],
+        when: NativeSeriesT,
+        then: NativeSeriesT,
         otherwise: NativeSeriesT | None = None,
     ) -> NativeSeriesT: ...
-    def _when_then_simple(
-        self, df: EagerDataFrameT, predicate: EagerExprT, then: EagerExprT
-    ) -> Sequence[EagerSeriesT_co]:
-        predicate_s = df._evaluate_single_output_expr(predicate)
-        then_s = df._evaluate_single_output_expr(then)
-        predicate_s, then_s = predicate_s._align_full_broadcast(predicate_s, then_s)
-        result = self._if_then_else(predicate_s.native, then_s.native)
-        return [then_s._with_native(result)]
-
-    def _when_then_chained(
-        self, df: EagerDataFrameT, args: tuple[EagerExprT, ...]
-    ) -> Sequence[EagerSeriesT_co]:
-        evaluated = [df._evaluate_single_output_expr(arg) for arg in args]
-        *pairs_list, otherwise_s = (
-            evaluated if len(evaluated) % 2 == 1 else (*evaluated, None)
-        )
-        conditions = pairs_list[::2]
-        values = pairs_list[1::2]
-
-        all_series = (
-            [*conditions, *values, otherwise_s]
-            if otherwise_s is not None
-            else [*conditions, *values]
-        )
-        aligned = conditions[0]._align_full_broadcast(*all_series)
-
-        num_conditions = len(conditions)
-        aligned_conditions = aligned[:num_conditions]
-        aligned_values = aligned[num_conditions : num_conditions * 2]
-        aligned_otherwise = aligned[-1] if otherwise_s is not None else None
-
-        if len(conditions) == 1:
-            result = self._if_then_else(
-                aligned_conditions[0].native,
-                aligned_values[0].native,
-                aligned_otherwise.native if aligned_otherwise is not None else None,
-            )
-        else:
-            result = self._if_then_else(
-                [c.native for c in aligned_conditions],
-                [v.native for v in aligned_values],
-                aligned_otherwise.native if aligned_otherwise is not None else None,
-            )
-
-        return [values[0]._with_native(result)]
-
-    def when_then(self, *args: EagerExprT) -> EagerExprT:
+    def when_then(
+        self, predicate: EagerExprT, then: EagerExprT, otherwise: EagerExprT | None = None
+    ) -> EagerExprT:
         def func(df: EagerDataFrameT) -> Sequence[EagerSeriesT_co]:
-            if len(args) == 2:
-                return self._when_then_simple(df, args[0], args[1])
-            return self._when_then_chained(df, args)
+            predicate_s = df._evaluate_single_output_expr(predicate)
+            align = predicate_s._align_full_broadcast
+
+            then_s = df._evaluate_single_output_expr(then)
+            if otherwise is None:
+                predicate_s, then_s = align(predicate_s, then_s)
+                result = self._if_then_else(predicate_s.native, then_s.native)
+
+            if otherwise is None:
+                predicate_s, then_s = align(predicate_s, then_s)
+                result = self._if_then_else(predicate_s.native, then_s.native)
+            else:
+                otherwise_s = df._evaluate_single_output_expr(otherwise)
+                predicate_s, then_s, otherwise_s = align(predicate_s, then_s, otherwise_s)
+                result = self._if_then_else(
+                    predicate_s.native, then_s.native, otherwise_s.native
+                )
+            return [then_s._with_native(result)]
 
         return self._expr._from_callable(
             func=func,
             evaluate_output_names=getattr(
-                args[1], "_evaluate_output_names", lambda _df: ["literal"]
+                then, "_evaluate_output_names", lambda _df: ["literal"]
             ),
-            alias_output_names=getattr(args[1], "_alias_output_names", None),
-            context=args[0],
+            alias_output_names=getattr(then, "_alias_output_names", None),
+            context=predicate,
         )
 
     def is_native(self, obj: Any, /) -> TypeIs[NativeFrameT | NativeSeriesT]:

narwhals/_dask/namespace.py

@@ -273,114 +273,50 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
             version=self._version,
         )
 
-    def _when_then_simple(
-        self,
-        df: DaskLazyFrame,
-        predicate: DaskExpr,
-        then: DaskExpr,
-        otherwise: DaskExpr | None,
-    ) -> list[dx.Series]:
-        then_value = df._evaluate_single_output_expr(then)
-        otherwise_value = (
-            df._evaluate_single_output_expr(otherwise)
-            if otherwise is not None
-            else otherwise
-        )
-
-        condition = df._evaluate_single_output_expr(predicate)
-        if all(
-            x._metadata.is_scalar_like
-            for x in (
-                (predicate, then) if otherwise is None else (predicate, then, otherwise)
-            )
-        ):
-            new_df = df._with_native(condition.to_frame())
-            condition = df._evaluate_single_output_expr(predicate.broadcast())
-            df = new_df
-
-        if otherwise is None:
-            (condition, then_series) = align_series_full_broadcast(
-                df, condition, then_value
-            )
-            validate_comparand(condition, then_series)
-            return [then_series.where(condition)]  # pyright: ignore[reportArgumentType]
-        (condition, then_series, otherwise_series) = align_series_full_broadcast(
-            df, condition, then_value, otherwise_value
-        )
-        validate_comparand(condition, then_series)
-        validate_comparand(condition, otherwise_series)
-        return [then_series.where(condition, otherwise_series)]  # pyright: ignore[reportArgumentType]
-
-    def _when_then_chained(
-        self, df: DaskLazyFrame, args: tuple[DaskExpr, ...]
-    ) -> list[dx.Series]:
-        import numpy as np  # ignore-banned-import
-
-        evaluated = [df._evaluate_single_output_expr(arg) for arg in args]
-        *pairs_list, otherwise_value = (
-            evaluated if len(evaluated) % 2 == 1 else (*evaluated, None)
-        )
-        conditions = pairs_list[::2]
-        values = pairs_list[1::2]
-
-        all_series = (
-            [*conditions, *values, otherwise_value]
-            if otherwise_value is not None
-            else [*conditions, *values]
-        )
-        aligned = align_series_full_broadcast(df, *all_series)
-
-        num_conditions = len(conditions)
-        aligned_conditions = aligned[:num_conditions]
-        aligned_values = aligned[num_conditions : num_conditions * 2]
-        aligned_otherwise = aligned[-1] if otherwise_value is not None else None
-
-        for cond, val in zip_strict(aligned_conditions, aligned_values):
-            validate_comparand(cond, val)
-        if aligned_otherwise is not None:
-            validate_comparand(aligned_conditions[0], aligned_otherwise)
-
-        def apply_select(*partition_series: pd.Series) -> pd.Series:
-            num_conds = len(aligned_conditions)
-            cond_parts = partition_series[:num_conds]
-            val_parts = partition_series[num_conds : num_conds * 2]
-            otherwise_part = (
-                partition_series[-1] if aligned_otherwise is not None else None
-            )
-            result = np.select(
-                list(cond_parts),
-                list(val_parts),
-                default=otherwise_part if otherwise_part is not None else np.nan,
-            )
-            return pd.Series(result, index=cond_parts[0].index)
-
-        map_args = [*aligned_conditions, *aligned_values]
-        if aligned_otherwise is not None:
-            map_args.append(aligned_otherwise)
-
-        return [
-            dd.map_partitions(
-                apply_select,
-                *map_args,
-                meta=(aligned_values[0].name, aligned_values[0].dtype),
+    def when_then(
+        self, predicate: DaskExpr, then: DaskExpr, otherwise: DaskExpr | None = None
+    ) -> DaskExpr:
+        def func(df: DaskLazyFrame) -> list[dx.Series]:
+            then_value = df._evaluate_single_output_expr(then)
+            otherwise_value = (
+                df._evaluate_single_output_expr(otherwise)
+                if otherwise is not None
+                else otherwise
             )
-        ]
 
-    def when_then(self, *args: DaskExpr) -> DaskExpr:
-        def func(df: DaskLazyFrame) -> list[dx.Series]:
-            if len(args) <= 3:
-                return self._when_then_simple(
-                    df, args[0], args[1], args[2] if len(args) == 3 else None
+            condition = df._evaluate_single_output_expr(predicate)
+            # re-evaluate DataFrame if the condition aggregates to force
+            # then/otherwise to be evaluated against the aggregated frame
+            if all(
+                x._metadata.is_scalar_like
+                for x in (
+                    (predicate, then)
+                    if otherwise is None
+                    else (predicate, then, otherwise)
                 )
-            return self._when_then_chained(df, args)
-
-        then_arg = args[1]
+            ):
+                new_df = df._with_native(condition.to_frame())
+                condition = df._evaluate_single_output_expr(predicate.broadcast())
+                df = new_df
+
+            if otherwise is None:
+                (condition, then_series) = align_series_full_broadcast(
+                    df, condition, then_value
+                )
+                validate_comparand(condition, then_series)
+                return [then_series.where(condition)]  # pyright: ignore[reportArgumentType]
+            (condition, then_series, otherwise_series) = align_series_full_broadcast(
+                df, condition, then_value, otherwise_value
+            )
+            validate_comparand(condition, then_series)
+            validate_comparand(condition, otherwise_series)
+            return [then_series.where(condition, otherwise_series)]  # pyright: ignore[reportArgumentType]
 
         return self._expr(
             call=func,
             evaluate_output_names=getattr(
-                then_arg, "_evaluate_output_names", lambda _df: ["literal"]
+                then, "_evaluate_output_names", lambda _df: ["literal"]
             ),
-            alias_output_names=getattr(then_arg, "_alias_output_names", None),
+            alias_output_names=getattr(then, "_alias_output_names", None),
             version=self._version,
         )

narwhals/_duckdb/namespace.py

@@ -67,34 +67,14 @@ def _lit(self, value: Any) -> Expression:
         return lit(value)
 
     def _when(
-        self, condition: Expression, value: Expression, *args: Expression
+        self,
+        condition: Expression,
+        value: Expression,
+        otherwise: Expression | None = None,
     ) -> Expression:
-        if not args:
+        if otherwise is None:
             return when(condition, value)
-
-        if len(args) == 1:
-            otherwise = args[0]
-            return when(condition, value).otherwise(otherwise)
-
-        all_exprs = [condition, value, *args]
-
-        has_otherwise = len(all_exprs) % 2 == 1
-
-        if has_otherwise:
-            *pairs, otherwise_expr = all_exprs
-        else:
-            pairs = all_exprs
-            otherwise_expr = None
-
-        result = when(pairs[0], pairs[1])
-
-        for cond, val in zip(pairs[2::2], pairs[3::2]):
-            result = result.when(cond, val)
-
-        if has_otherwise:
-            result = result.otherwise(otherwise_expr)
-
-        return result
+        return when(condition, value).otherwise(otherwise)
 
     def _coalesce(self, *exprs: Expression) -> Expression:
         return CoalesceOperator(*exprs)

narwhals/_ibis/namespace.py

@@ -54,29 +54,12 @@ def _function(self, name: str, *args: ir.Value | PythonLiteral) -> ir.Value:
     def _lit(self, value: Any) -> ir.Value:
         return lit(value)
 
-    def _when(self, condition: ir.Value, value: ir.Value, *args: ir.Value) -> ir.Value:
-        if not args:
+    def _when(
+        self, condition: ir.Value, value: ir.Value, otherwise: ir.Expr | None = None
+    ) -> ir.Value:
+        if otherwise is None:
             return ibis.cases((condition, value))
-
-        if len(args) == 1:
-            otherwise = args[0]
-            return ibis.cases((condition, value), else_=otherwise)  # pragma: no cover
-
-        all_exprs = [condition, value, *args]
-
-        has_otherwise = len(all_exprs) % 2 == 1
-
-        if has_otherwise:
-            *pairs, otherwise_expr = all_exprs
-        else:
-            pairs = all_exprs
-            otherwise_expr = None
-
-        tuples = list(zip(pairs[::2], pairs[1::2]))
-
-        if has_otherwise:
-            return ibis.cases(*tuples, else_=otherwise_expr)
-        return ibis.cases(*tuples)
+        return ibis.cases((condition, value), else_=otherwise)  # pragma: no cover
 
     def _coalesce(self, *exprs: ir.Value) -> ir.Value:
         return ibis.coalesce(*exprs)

narwhals/_pandas_like/namespace.py

@@ -366,24 +366,12 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
 
     def _if_then_else(
         self,
-        when: NativeSeriesT | list[NativeSeriesT],
-        then: NativeSeriesT | list[NativeSeriesT],
+        when: NativeSeriesT,
+        then: NativeSeriesT,
         otherwise: NativeSeriesT | None = None,
     ) -> NativeSeriesT:
-        if not isinstance(when, list):
-            where: Incomplete = then.where
-            return where(when) if otherwise is None else where(when, otherwise)
-
-        import numpy as np
-
-        condlist = when
-        choicelist = then
-
-        default = otherwise if otherwise is not None else np.nan
-
-        result = np.select(condlist, choicelist, default=default)
-
-        return then[0].__class__(result, index=then[0].index)
+        where: Incomplete = then.where
+        return where(when) if otherwise is None else where(when, otherwise)
 
 
 class _NativeConcat(Protocol[NativeDataFrameT, NativeSeriesT]):