refactor(DRAFT): Possibly all ArrowExpr?

Main source of errors atm are expr namespaces - still needs migrating

Author: dangotbanned

narwhals/_arrow/dataframe.py

@@ -778,7 +778,7 @@ def unique(
 
         keep_idx = self.simple_select(*subset).is_unique()
         plx = self.__narwhals_namespace__()
-        return self.filter(plx._create_expr_from_series(keep_idx))
+        return self.filter(plx._expr._from_series(keep_idx))
 
     def gather_every(self: Self, n: int, offset: int) -> Self:
         return self._from_native_frame(

narwhals/_arrow/expr.py

@@ -29,6 +29,7 @@
     from narwhals._arrow.namespace import ArrowNamespace
     from narwhals.dtypes import DType
     from narwhals.utils import Version
+    from narwhals.utils import _FullContext
 
 
 class ArrowExpr(EagerExpr["ArrowDataFrame", ArrowSeries]):
@@ -45,6 +46,7 @@ def __init__(
         backend_version: tuple[int, ...],
         version: Version,
         call_kwargs: dict[str, Any] | None = None,
+        implementation: Implementation | None = None,
     ) -> None:
         self._call = call
         self._depth = depth
@@ -63,8 +65,7 @@ def from_column_names(
         /,
         *,
         function_name: str,
-        backend_version: tuple[int, ...],
-        version: Version,
+        context: _FullContext,
     ) -> Self:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             try:
@@ -91,16 +92,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             function_name=function_name,
             evaluate_output_names=evaluate_column_names,
             alias_output_names=None,
-            backend_version=backend_version,
-            version=version,
+            backend_version=context._backend_version,
+            version=context._version,
         )
 
     @classmethod
     def from_column_indices(
-        cls: type[Self],
-        *column_indices: int,
-        backend_version: tuple[int, ...],
-        version: Version,
+        cls: type[Self], *column_indices: int, context: _FullContext
     ) -> Self:
         from narwhals._arrow.series import ArrowSeries
 
@@ -121,8 +119,8 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             function_name="nth",
             evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
             alias_output_names=None,
-            backend_version=backend_version,
-            version=version,
+            backend_version=context._backend_version,
+            version=context._version,
         )
 
     def __narwhals_namespace__(self: Self) -> ArrowNamespace:

narwhals/_arrow/namespace.py

@@ -57,59 +57,7 @@ def _expr(self) -> type[ArrowExpr]:
     def _series(self) -> type[ArrowSeries]:
         return ArrowSeries
 
-    def _create_expr_from_callable(
-        self: Self,
-        func: Callable[[ArrowDataFrame], Sequence[ArrowSeries]],
-        *,
-        depth: int,
-        function_name: str,
-        evaluate_output_names: Callable[[ArrowDataFrame], Sequence[str]],
-        alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
-        call_kwargs: dict[str, Any] | None = None,
-    ) -> ArrowExpr:
-        from narwhals._arrow.expr import ArrowExpr
-
-        return ArrowExpr(
-            func,
-            depth=depth,
-            function_name=function_name,
-            evaluate_output_names=evaluate_output_names,
-            alias_output_names=alias_output_names,
-            backend_version=self._backend_version,
-            version=self._version,
-            call_kwargs=call_kwargs,
-        )
-
-    def _create_expr_from_series(self: Self, series: ArrowSeries) -> ArrowExpr:
-        from narwhals._arrow.expr import ArrowExpr
-
-        return ArrowExpr(
-            lambda _df: [series],
-            depth=0,
-            function_name="series",
-            evaluate_output_names=lambda _df: [series.name],
-            alias_output_names=None,
-            backend_version=self._backend_version,
-            version=self._version,
-        )
-
-    def _create_series_from_scalar(
-        self: Self, value: Any, *, reference_series: ArrowSeries
-    ) -> ArrowSeries:
-        from narwhals._arrow.series import ArrowSeries
-
-        if self._backend_version < (13,) and hasattr(value, "as_py"):
-            value = value.as_py()
-        return ArrowSeries._from_iterable(
-            [value],
-            name=reference_series.name,
-            backend_version=self._backend_version,
-            version=self._version,
-        )
-
     def _create_compliant_series(self: Self, value: Any) -> ArrowSeries:
-        from narwhals._arrow.series import ArrowSeries
-
         return ArrowSeries(
             native_series=pa.chunked_array([value]),
             name="",
@@ -127,39 +75,26 @@ def __init__(
 
     # --- selection ---
     def col(self: Self, *column_names: str) -> ArrowExpr:
-        from narwhals._arrow.expr import ArrowExpr
-
-        return ArrowExpr.from_column_names(
-            passthrough_column_names(column_names),
-            function_name="col",
-            backend_version=self._backend_version,
-            version=self._version,
+        return self._expr.from_column_names(
+            passthrough_column_names(column_names), function_name="col", context=self
         )
 
     def exclude(self: Self, excluded_names: Container[str]) -> ArrowExpr:
-        return ArrowExpr.from_column_names(
+        return self._expr.from_column_names(
             partial(exclude_column_names, names=excluded_names),
             function_name="exclude",
-            backend_version=self._backend_version,
-            version=self._version,
+            context=self,
         )
 
     def nth(self: Self, *column_indices: int) -> ArrowExpr:
-        from narwhals._arrow.expr import ArrowExpr
-
-        return ArrowExpr.from_column_indices(
-            *column_indices, backend_version=self._backend_version, version=self._version
-        )
+        return self._expr.from_column_indices(*column_indices, context=self)
 
     def len(self: Self) -> ArrowExpr:
         # coverage bug? this is definitely hit
-        return ArrowExpr(  # pragma: no cover
+        return self._expr(  # pragma: no cover
             lambda df: [
                 ArrowSeries._from_iterable(
-                    [len(df._native_frame)],
-                    name="len",
-                    backend_version=self._backend_version,
-                    version=self._version,
+                    [len(df._native_frame)], name="len", context=self
                 )
             ],
             depth=0,
@@ -171,26 +106,20 @@ def len(self: Self) -> ArrowExpr:
         )
 
     def all(self: Self) -> ArrowExpr:
-        return ArrowExpr.from_column_names(
-            get_column_names,
-            function_name="all",
-            backend_version=self._backend_version,
-            version=self._version,
+        return self._expr.from_column_names(
+            get_column_names, function_name="all", context=self
         )
 
     def lit(self: Self, value: Any, dtype: DType | None) -> ArrowExpr:
         def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
             arrow_series = ArrowSeries._from_iterable(
-                data=[value],
-                name="literal",
-                backend_version=self._backend_version,
-                version=self._version,
+                data=[value], name="literal", context=self
             )
             if dtype:
                 return arrow_series.cast(dtype)
             return arrow_series
 
-        return ArrowExpr(
+        return self._expr(
             lambda df: [_lit_arrow_series(df)],
             depth=0,
             function_name="lit",
@@ -200,30 +129,34 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
             version=self._version,
         )
 
-    def all_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
+    # NOTE: Needs to be resolved in `EagerNamespace`
+    # Probably, by adding an `EagerExprT` typevar
+    def all_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:  # type: ignore[override]
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = chain.from_iterable(expr(df) for expr in exprs)
             return [reduce(operator.and_, align_series_full_broadcast(*series))]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="all_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def any_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = chain.from_iterable(expr(df) for expr in exprs)
             return [reduce(operator.or_, align_series_full_broadcast(*series))]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="any_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def sum_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
@@ -232,12 +165,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = (s.fill_null(0, strategy=None, limit=None) for s in it)
             return [reduce(operator.add, align_series_full_broadcast(*series))]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="sum_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def mean_horizontal(self: Self, *exprs: ArrowExpr) -> IntoArrowExpr:
@@ -253,12 +187,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             )
             return [reduce(operator.add, series) / reduce(operator.add, non_na)]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="mean_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def min_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
@@ -281,12 +216,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
                 )
             ]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="min_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def max_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
@@ -310,12 +246,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
                 )
             ]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="max_horizontal",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
     def concat(
@@ -381,12 +318,13 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
                 )
             ]
 
-        return self._create_expr_from_callable(
+        return self._expr._from_callable(
             func=func,
             depth=max(x._depth for x in exprs) + 1,
             function_name="concat_str",
             evaluate_output_names=combine_evaluate_output_names(*exprs),
             alias_output_names=combine_alias_output_names(*exprs),
+            context=self,
         )
 
 
@@ -407,16 +345,13 @@ def __init__(
         self._version = version
 
     def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
-        plx = df.__narwhals_namespace__()
         condition = self._condition(df)[0]
         condition_native = condition._native_series
 
         if isinstance(self._then_value, ArrowExpr):
             value_series = self._then_value(df)[0]
         else:
-            value_series = plx._create_series_from_scalar(
-                self._then_value, reference_series=condition.alias("literal")
-            )
+            value_series = condition.alias("literal")._from_scalar(self._then_value)
             value_series._broadcast = True
         value_series_native = extract_dataframe_comparand(
             len(df), value_series, self._backend_version
@@ -474,6 +409,7 @@ def __init__(
         backend_version: tuple[int, ...],
         version: Version,
         call_kwargs: dict[str, Any] | None = None,
+        implementation: Implementation | None = None,
     ) -> None:
         self._backend_version = backend_version
         self._version = version

narwhals/_arrow/series.py

@@ -54,6 +54,7 @@
     from narwhals.typing import _1DArray
     from narwhals.typing import _2DArray
     from narwhals.utils import Version
+    from narwhals.utils import _FullContext
 
 
 # TODO @dangotbanned: move into `_arrow.utils`
@@ -140,16 +141,20 @@ def _from_iterable(
         data: Iterable[Any],
         name: str,
         *,
-        backend_version: tuple[int, ...],
-        version: Version,
+        context: _FullContext,
     ) -> Self:
         return cls(
             chunked_array([data]),
             name=name,
-            backend_version=backend_version,
-            version=version,
+            backend_version=context._backend_version,
+            version=context._version,
         )
 
+    def _from_scalar(self, value: Any) -> Self:
+        if self._backend_version < (13,) and hasattr(value, "as_py"):
+            value = value.as_py()
+        return super()._from_scalar(value)
+
     def __narwhals_namespace__(self: Self) -> ArrowNamespace:
         from narwhals._arrow.namespace import ArrowNamespace
 
@@ -570,12 +575,7 @@ def arg_true(self: Self) -> Self:
 
         ser = self._native_series
         res = np.flatnonzero(ser)
-        return self._from_iterable(
-            res,
-            name=self.name,
-            backend_version=self._backend_version,
-            version=self._version,
-        )
+        return self._from_iterable(res, name=self.name, context=self)
 
     def item(self: Self, index: int | None = None) -> Any:
         if index is None: