perf: Avoid evaluating expr twice in mean_horizontal (#2028)

Author: MarcoGorelli

narwhals/_arrow/namespace.py

@@ -237,14 +237,9 @@ def mean_horizontal(self: Self, *exprs: ArrowExpr) -> IntoArrowExpr:
         dtypes = import_dtypes_module(self._version)
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            series = (
-                s.fill_null(0, strategy=None, limit=None)
-                for _expr in exprs
-                for s in _expr(df)
-            )
-            non_na = (
-                1 - s.is_null().cast(dtypes.Int64()) for _expr in exprs for s in _expr(df)
-            )
+            expr_results = [s for _expr in exprs for s in _expr(df)]
+            series = (s.fill_null(0, strategy=None, limit=None) for s in expr_results)
+            non_na = (1 - s.is_null().cast(dtypes.Int64()) for s in expr_results)
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
         return self._create_expr_from_callable(

narwhals/_dask/namespace.py

@@ -237,8 +237,9 @@ def concat(
 
     def mean_horizontal(self: Self, *exprs: DaskExpr) -> DaskExpr:
         def func(df: DaskLazyFrame) -> list[dx.Series]:
-            series = (s.fillna(0) for _expr in exprs for s in _expr(df))
-            non_na = (1 - s.isna() for _expr in exprs for s in _expr(df))
+            expr_results = [s for _expr in exprs for s in _expr(df)]
+            series = (s.fillna(0) for s in expr_results)
+            non_na = (1 - s.isna() for s in expr_results)
             return [
                 name_preserving_div(
                     reduce(name_preserving_sum, series),

narwhals/_pandas_like/namespace.py

@@ -240,12 +240,9 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
 
     def mean_horizontal(self: Self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            series = (
-                s.fill_null(0, strategy=None, limit=None)
-                for _expr in exprs
-                for s in _expr(df)
-            )
-            non_na = (1 - s.is_null() for _expr in exprs for s in _expr(df))
+            expr_results = [s for _expr in exprs for s in _expr(df)]
+            series = (s.fill_null(0, strategy=None, limit=None) for s in expr_results)
+            non_na = (1 - s.is_null() for s in expr_results)
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
         return self._create_expr_from_callable(