old polars and some more

Author: FBruzzesi

narwhals/_expression_parsing.py

@@ -30,17 +30,27 @@
     from narwhals._pandas_like.namespace import PandasLikeNamespace
     from narwhals._pandas_like.series import PandasLikeSeries
     from narwhals._pandas_like.typing import IntoPandasLikeExpr
+    from narwhals._polars.expr import PolarsExpr
+    from narwhals._polars.namespace import PolarsNamespace
+    from narwhals._polars.series import PolarsSeries
+    from narwhals._polars.typing import IntoPolarsExpr
 
-    CompliantNamespace = Union[PandasLikeNamespace, ArrowNamespace, DaskNamespace]
-    CompliantExpr = Union[PandasLikeExpr, ArrowExpr, DaskExpr]
-    IntoCompliantExpr = Union[IntoPandasLikeExpr, IntoArrowExpr, IntoDaskExpr]
+    CompliantNamespace = Union[
+        PandasLikeNamespace, ArrowNamespace, DaskNamespace, PolarsNamespace
+    ]
+    CompliantExpr = Union[PandasLikeExpr, ArrowExpr, DaskExpr, PolarsExpr]
+    IntoCompliantExpr = Union[
+        IntoPandasLikeExpr, IntoArrowExpr, IntoDaskExpr, IntoPolarsExpr
+    ]
     IntoCompliantExprT = TypeVar("IntoCompliantExprT", bound=IntoCompliantExpr)
     CompliantExprT = TypeVar("CompliantExprT", bound=CompliantExpr)
-    CompliantSeries = Union[PandasLikeSeries, ArrowSeries]
+    CompliantSeries = Union[PandasLikeSeries, ArrowSeries, PolarsSeries]
     ListOfCompliantSeries = Union[
-        list[PandasLikeSeries], list[ArrowSeries], list[DaskExpr]
+        list[PandasLikeSeries], list[ArrowSeries], list[DaskExpr], list[PolarsSeries]
+    ]
+    ListOfCompliantExpr = Union[
+        list[PandasLikeExpr], list[ArrowExpr], list[DaskExpr], list[PolarsExpr]
     ]
-    ListOfCompliantExpr = Union[list[PandasLikeExpr], list[ArrowExpr], list[DaskExpr]]
     CompliantDataFrame = Union[PandasLikeDataFrame, ArrowDataFrame, DaskLazyFrame]
 
     T = TypeVar("T")
@@ -133,14 +143,22 @@ def parse_into_exprs(
 ) -> list[DaskExpr]: ...
 
 
+@overload
+def parse_into_exprs(
+    *exprs: IntoPolarsExpr,
+    namespace: PolarsNamespace,
+    **named_exprs: IntoPolarsExpr,
+) -> list[PolarsExpr]: ...
+
+
 def parse_into_exprs(
     *exprs: IntoCompliantExpr,
     namespace: CompliantNamespace,
     **named_exprs: IntoCompliantExpr,
 ) -> ListOfCompliantExpr:
     """Parse each input as an expression (if it's not already one). See `parse_into_expr` for
     more details."""
-    return [  # type: ignore[return-value]
+    return [
         parse_into_expr(into_expr, namespace=namespace) for into_expr in flatten(exprs)
     ] + [
         parse_into_expr(expr, namespace=namespace).alias(name)

narwhals/_polars/namespace.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+from functools import reduce
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
 from typing import Sequence
 
 from narwhals import dtypes
+from narwhals._expression_parsing import parse_into_exprs
 from narwhals._polars.utils import extract_args_kwargs
 from narwhals._polars.utils import narwhals_to_native_dtype
 from narwhals.dependencies import get_polars
@@ -15,6 +17,7 @@
     from narwhals._polars.dataframe import PolarsDataFrame
     from narwhals._polars.dataframe import PolarsLazyFrame
     from narwhals._polars.expr import PolarsExpr
+    from narwhals._polars.typing import IntoPolarsExpr
 
 
 class PolarsNamespace:
@@ -85,11 +88,28 @@ def lit(self, value: Any, dtype: dtypes.DType | None = None) -> PolarsExpr:
             return PolarsExpr(pl.lit(value, dtype=narwhals_to_native_dtype(dtype)))
         return PolarsExpr(pl.lit(value))
 
-    def mean(self, *column_names: str) -> Any:
+    def mean(self, *column_names: str) -> PolarsExpr:
+        from narwhals._polars.expr import PolarsExpr
+
         pl = get_polars()
         if self._backend_version < (0, 20, 4):  # pragma: no cover
-            return pl.mean([*column_names])
-        return pl.mean(*column_names)
+            return PolarsExpr(pl.mean([*column_names]))
+        return PolarsExpr(pl.mean(*column_names))
+
+    def mean_horizontal(self, *exprs: IntoPolarsExpr) -> PolarsExpr:
+        from narwhals._polars.expr import PolarsExpr
+
+        pl = get_polars()
+        polars_exprs = parse_into_exprs(*exprs, namespace=self)
+
+        if self._backend_version < (0, 20, 8):  # pragma: no cover
+            total = reduce(lambda x, y: x + y, (e.fill_null(0.0) for e in polars_exprs))
+            n_non_zero = reduce(
+                lambda x, y: x + y, ((1 - e.is_null()) for e in polars_exprs)
+            )
+            return PolarsExpr(total._native_expr / n_non_zero._native_expr)
+
+        return PolarsExpr(pl.mean_horizontal([e._native_expr for e in polars_exprs]))
 
     @property
     def selectors(self) -> PolarsSelectors:

narwhals/_polars/typing.py

@@ -0,0 +1,17 @@
+from __future__ import annotations  # pragma: no cover
+
+from typing import TYPE_CHECKING  # pragma: no cover
+from typing import Union  # pragma: no cover
+
+if TYPE_CHECKING:
+    import sys
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+    from narwhals._polars.expr import PolarsExpr
+    from narwhals._polars.series import PolarsSeries
+
+    IntoPolarsExpr: TypeAlias = Union[PolarsExpr, str, PolarsSeries]

narwhals/stable/v1.py

@@ -1607,6 +1607,7 @@ def from_dict(
     "min",
     "max",
     "mean",
+    "mean_horizontal",
     "sum",
     "sum_horizontal",
     "DataFrame",

tests/expr_and_series/mean_horizontal_test.py

@@ -7,11 +7,9 @@
 
 
 @pytest.mark.parametrize("col_expr", [nw.col("a"), "a"])
-def test_sumh(constructor: Any, col_expr: Any) -> None:
-    data = {"a": [1, 3, None], "b": [4, None, 6]}
+def test_meanh(constructor: Any, col_expr: Any) -> None:
+    data = {"a": [1, 3, None, None], "b": [4, None, 6, None]}
     df = nw.from_native(constructor(data))
     result = df.select(horizontal_mean=nw.mean_horizontal(col_expr, nw.col("b")))
-    expected = {
-        "horizontal_mean": [2.5, 3.0, 6.0],
-    }
+    expected = {"horizontal_mean": [2.5, 3.0, 6.0, float("nan")]}
     compare_dicts(result, expected)