chore: Redo CompliantExpr.from_column_indices (#2561)

* chore: Redo `CompliantExpr.from_column_indices`

Resolves 1 task for (#2525)

* refactor: Get a bit DRY-er

Resolves (#2561 (comment))

Author: dangotbanned

narwhals/_arrow/expr.py

@@ -94,27 +94,20 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
         )
 
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self:
-        from narwhals._arrow.series import ArrowSeries
-
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            tbl = df.native
+            cols = df.columns
             return [
-                ArrowSeries(
-                    df.native[column_index],
-                    name=df.native.column_names[column_index],
-                    backend_version=df._backend_version,
-                    version=df._version,
-                )
-                for column_index in column_indices
+                ArrowSeries.from_native(tbl[i], name=cols[i], context=df)
+                for i in column_indices
             ]
 
         return cls(
             func,
             depth=0,
             function_name="nth",
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             backend_version=context._backend_version,
             version=context._version,

narwhals/_compliant/expr.py

@@ -109,9 +109,14 @@ def from_column_names(
         context: _FullContext,
     ) -> Self: ...
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self: ...
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self: ...
+    @staticmethod
+    def _eval_names_indices(indices: Sequence[int], /) -> EvalNames[CompliantFrameT]:
+        def fn(df: CompliantFrameT) -> Sequence[str]:
+            column_names = df.columns
+            return [column_names[i] for i in indices]
+
+        return fn
 
     def is_null(self) -> Self: ...
     def abs(self) -> Self: ...

narwhals/_dask/expr.py

@@ -135,19 +135,15 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
         )
 
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self:
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self:
         def func(df: DaskLazyFrame) -> list[dx.Series]:
-            return [
-                df._native_frame.iloc[:, column_index] for column_index in column_indices
-            ]
+            return [df.native.iloc[:, i] for i in column_indices]
 
         return cls(
             func,
             depth=0,
             function_name="nth",
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             backend_version=context._backend_version,
             version=context._version,

narwhals/_duckdb/expr.py

@@ -206,7 +206,7 @@ def func(df: DuckDBLazyFrame) -> list[duckdb.Expression]:
 
         return cls(
             func,
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             backend_version=context._backend_version,
             version=context._version,

narwhals/_ibis/expr.py

@@ -197,15 +197,13 @@ def func(df: IbisLazyFrame) -> list[ir.Column]:
         )
 
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self:
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self:
         def func(df: IbisLazyFrame) -> list[ir.Column]:
             return [df.native[i] for i in column_indices]
 
         return cls(
             func,
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             backend_version=context._backend_version,
             version=context._version,

narwhals/_pandas_like/expr.py

@@ -148,25 +148,19 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
         )
 
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self:
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+            native = df.native
             return [
-                PandasLikeSeries(
-                    df._native_frame.iloc[:, column_index],
-                    implementation=df._implementation,
-                    backend_version=df._backend_version,
-                    version=df._version,
-                )
-                for column_index in column_indices
+                PandasLikeSeries.from_native(native.iloc[:, i], context=df)
+                for i in column_indices
             ]
 
         return cls(
             func,
             depth=0,
             function_name="nth",
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             implementation=context._implementation,
             backend_version=context._backend_version,

narwhals/_polars/expr.py

@@ -263,6 +263,7 @@ def struct(self) -> PolarsExprStructNamespace:
     __call__: Any
     from_column_names: Any
     from_column_indices: Any
+    _eval_names_indices: Any
 
     # Polars
     abs: Method[Self]

narwhals/_spark_like/expr.py

@@ -299,16 +299,14 @@ def func(df: SparkLikeLazyFrame) -> list[Column]:
         )
 
     @classmethod
-    def from_column_indices(
-        cls: type[Self], *column_indices: int, context: _FullContext
-    ) -> Self:
+    def from_column_indices(cls, *column_indices: int, context: _FullContext) -> Self:
         def func(df: SparkLikeLazyFrame) -> list[Column]:
             columns = df.columns
             return [df._F.col(columns[i]) for i in column_indices]
 
         return cls(
             func,
-            evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
+            evaluate_output_names=cls._eval_names_indices(column_indices),
             alias_output_names=None,
             backend_version=context._backend_version,
             version=context._version,