perf: avoid full broadcast for otherwise_value in when/then/otherwise (#2098)

Author: MarcoGorelli

narwhals/_arrow/namespace.py

@@ -424,10 +424,10 @@ def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
         if isinstance(self._otherwise_value, ArrowExpr):
             otherwise_series = self._otherwise_value(df)[0]
         else:
-            otherwise_series = plx._create_series_from_scalar(
-                self._otherwise_value, reference_series=condition.alias("literal")
+            native_result = pc.if_else(
+                condition_native, value_series_native, self._otherwise_value
             )
-            otherwise_series._broadcast = True
+            return [value_series._from_native_series(native_result)]
 
         otherwise_series_native = extract_dataframe_comparand(
             len(df), otherwise_series, self._backend_version

narwhals/_dask/namespace.py

@@ -363,7 +363,7 @@ def __call__(self: Self, df: DaskLazyFrame) -> Sequence[dx.Series]:
         if isinstance(self._otherwise_value, DaskExpr):
             otherwise_value = self._otherwise_value(df)[0]
         else:
-            otherwise_value = self._otherwise_value
+            return [then_series.where(condition, self._otherwise_value)]
         (otherwise_series,) = align_series_full_broadcast(df, otherwise_value)
         validate_comparand(condition, otherwise_series)
         return [then_series.where(condition, otherwise_series)]  # pyright: ignore[reportArgumentType]

narwhals/_pandas_like/namespace.py

@@ -445,10 +445,10 @@ def __call__(self: Self, df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
         if isinstance(self._otherwise_value, PandasLikeExpr):
             otherwise_series = self._otherwise_value(df)[0]
         else:
-            otherwise_series = plx._create_series_from_scalar(
-                self._otherwise_value, reference_series=condition.alias("literal")
+            native_result = value_series_native.where(
+                condition_native, self._otherwise_value
             )
-            otherwise_series._broadcast = True
+            return [value_series._from_native_series(native_result)]
         otherwise_series_native = extract_dataframe_comparand(
             df._native_frame.index, otherwise_series
         )

narwhals/_spark_like/utils.py

@@ -19,6 +19,7 @@
     from narwhals.utils import Version
 
 
+# NOTE: don't lru_cache this as `ModuleType` isn't hashable
 def native_to_narwhals_dtype(
     dtype: pyspark_types.DataType,
     version: Version,