patch: fix when-then double lit case (#1810)

narwhals/_arrow/namespace.py

@@ -435,19 +435,18 @@ def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
 
         plx = df.__narwhals_namespace__()
         condition = parse_into_expr(self._condition, namespace=plx)(df)[0]
+
         try:
             value_series = parse_into_expr(self._then_value, namespace=plx)(df)[0]
         except TypeError:
-            # `self._otherwise_value` is a scalar and can't be converted to an expression
-            value_series = condition.__class__._from_iterable(
-                pa.repeat(pa.scalar(self._then_value), len(condition)),
-                name="literal",
-                backend_version=self._backend_version,
-                version=self._version,
+            # `self._then_value` is a scalar and can't be converted to an expression
+            value_series = plx._create_series_from_scalar(
+                self._then_value, reference_series=condition
             )
 
-        value_series_native = value_series._native_series
-        condition_native = condition._native_series
+        condition_native, value_series_native = broadcast_series(
+            [condition, value_series]
+        )
 
         if self._otherwise_value is None:
             otherwise_native = pa.repeat(
@@ -472,9 +471,7 @@ def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
             ]
         else:
             otherwise_series = otherwise_expr(df)[0]
-            condition_native, otherwise_native = broadcast_series(
-                [condition, otherwise_series]
-            )
+            _, otherwise_native = broadcast_series([condition, otherwise_series])
             return [
                 value_series._from_native_series(
                     pc.if_else(condition_native, value_series_native, otherwise_native)

narwhals/_dask/namespace.py

@@ -398,13 +398,24 @@ def __call__(self, df: DaskLazyFrame) -> Sequence[dx.Series]:
         plx = df.__narwhals_namespace__()
         condition = parse_into_expr(self._condition, namespace=plx)(df)[0]
         condition = cast("dx.Series", condition)
+
         try:
-            value_series = parse_into_expr(self._then_value, namespace=plx)(df)[0]
+            then_expr = parse_into_expr(self._then_value, namespace=plx)
         except TypeError:
-            # `self._otherwise_value` is a scalar and can't be converted to an expression
+            # `self._then_value` is a scalar and can't be converted to an expression
+            value_sequence: Sequence[Any] = [self._then_value]
+            is_scalar = True
+        else:
+            is_scalar = then_expr._returns_scalar  # type: ignore[attr-defined]
+            value_sequence = then_expr(df)[0]
+
+        if is_scalar:
             _df = condition.to_frame("a")
-            _df["tmp"] = self._then_value
+            _df["tmp"] = value_sequence[0]
             value_series = _df["tmp"]
+        else:
+            value_series = value_sequence
+
         value_series = cast("dx.Series", value_series)
         validate_comparand(condition, value_series)
 

narwhals/_pandas_like/namespace.py

@@ -461,20 +461,17 @@ def __call__(self, df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
 
         plx = df.__narwhals_namespace__()
         condition = parse_into_expr(self._condition, namespace=plx)(df)[0]
+
         try:
             value_series = parse_into_expr(self._then_value, namespace=plx)(df)[0]
         except TypeError:
-            # `self._otherwise_value` is a scalar and can't be converted to an expression
-            value_series = condition.__class__._from_iterable(
-                [self._then_value] * len(condition),
-                name="literal",
-                index=condition._native_series.index,
-                implementation=self._implementation,
-                backend_version=self._backend_version,
-                version=self._version,
+            # `self._then_value` is a scalar and can't be converted to an expression
+            value_series = plx._create_series_from_scalar(
+                self._then_value, reference_series=condition
             )
-        value_series_native, condition_native = broadcast_align_and_extract_native(
-            value_series, condition
+
+        condition_native, value_series_native = broadcast_align_and_extract_native(
+            condition, value_series
         )
 
         if self._otherwise_value is None:
@@ -494,7 +491,14 @@ def __call__(self, df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
             ]
         else:
             otherwise_series = otherwise_expr(df)[0]
-            return [value_series.zip_with(condition, otherwise_series)]
+            _, otherwise_native = broadcast_align_and_extract_native(
+                condition, otherwise_series
+            )
+            return [
+                value_series._from_native_series(
+                    value_series_native.where(condition_native, otherwise_native)
+                )
+            ]
 
     def then(self, value: PandasLikeExpr | PandasLikeSeries | Any) -> PandasThen:
         self._then_value = value

tests/expr_and_series/when_test.py

@@ -141,3 +141,13 @@ def test_when_then_otherwise_lit_str(
     result = df.select(nw.when(nw.col("a") > 1).then(nw.col("b")).otherwise(nw.lit("z")))
     expected = {"b": ["z", "b", "c"]}
     assert_equal_data(result, expected)
+
+
+def test_when_then_otherwise_both_lit(constructor: Constructor) -> None:
+    df = nw.from_native(constructor(data))
+    result = df.select(
+        x1=nw.when(nw.col("a") > 1).then(nw.lit(42)).otherwise(nw.lit(-1)),
+        x2=nw.when(nw.col("a") > 2).then(nw.lit(42)).otherwise(nw.lit(-1)),
+    )
+    expected = {"x1": [-1, 42, 42], "x2": [-1, -1, 42]}
+    assert_equal_data(result, expected)