refactor: Generic CompliantSelector

narwhals/_arrow/dataframe.py

@@ -71,7 +71,7 @@
 from narwhals.typing import CompliantLazyFrame
 
 
-class ArrowDataFrame(CompliantDataFrame, CompliantLazyFrame):
+class ArrowDataFrame(CompliantDataFrame["ArrowSeries"], CompliantLazyFrame):
     # --- not in the spec ---
     def __init__(
         self: Self,

narwhals/_arrow/namespace.py

@@ -26,6 +26,7 @@
 from narwhals._expression_parsing import combine_evaluate_output_names
 from narwhals.typing import CompliantNamespace
 from narwhals.utils import Implementation
+from narwhals.utils import get_column_names
 from narwhals.utils import import_dtypes_module
 from narwhals.utils import is_compliant_expr
 
@@ -159,7 +160,7 @@ def all(self: Self) -> ArrowExpr:
             ],
             depth=0,
             function_name="all",
-            evaluate_output_names=lambda df: df.columns,
+            evaluate_output_names=get_column_names,
             alias_output_names=None,
             backend_version=self._backend_version,
             version=self._version,

narwhals/_arrow/selectors.py

@@ -1,127 +1,54 @@
 from __future__ import annotations
 
-import re
 from typing import TYPE_CHECKING
-from typing import Any
-from typing import Callable
-from typing import Iterable
-from typing import Sequence
+from typing import Iterator
 
 from narwhals._arrow.expr import ArrowExpr
-from narwhals.utils import _parse_time_unit_and_time_zone
-from narwhals.utils import dtype_matches_time_unit_and_time_zone
-from narwhals.utils import import_dtypes_module
+from narwhals._selectors import CompliantSelector
+from narwhals._selectors import CompliantSelectorNamespace
 
 if TYPE_CHECKING:
-    from datetime import timezone
-
     from typing_extensions import Self
 
     from narwhals._arrow.dataframe import ArrowDataFrame
     from narwhals._arrow.series import ArrowSeries
-    from narwhals.dtypes import DType
-    from narwhals.typing import TimeUnit
-    from narwhals.utils import _LimitedContext
-
-
-class ArrowSelectorNamespace:
-    def __init__(self: Self, context: _LimitedContext, /) -> None:
-        self._backend_version = context._backend_version
-        self._version = context._version
-
-    def by_dtype(self: Self, dtypes: Iterable[DType | type[DType]]) -> ArrowSelector:
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            return [df[col] for col in df.columns if df.schema[col] in dtypes]
-
-        def evaluate_output_names(df: ArrowDataFrame) -> Sequence[str]:
-            return [col for col in df.columns if df.schema[col] in dtypes]
-
-        return selector(self, func, evaluate_output_names)
-
-    def matches(self: Self, pattern: str) -> ArrowSelector:
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            return [df[col] for col in df.columns if re.search(pattern, col)]
-
-        def evaluate_output_names(df: ArrowDataFrame) -> Sequence[str]:
-            return [col for col in df.columns if re.search(pattern, col)]
-
-        return selector(self, func, evaluate_output_names)
-
-    def numeric(self: Self) -> ArrowSelector:
-        dtypes = import_dtypes_module(self._version)
-        return self.by_dtype(
-            [
-                dtypes.Int128,
-                dtypes.Int64,
-                dtypes.Int32,
-                dtypes.Int16,
-                dtypes.Int8,
-                dtypes.UInt128,
-                dtypes.UInt64,
-                dtypes.UInt32,
-                dtypes.UInt16,
-                dtypes.UInt8,
-                dtypes.Float64,
-                dtypes.Float32,
-            ],
+    from narwhals._selectors import EvalNames
+    from narwhals._selectors import EvalSeries
+    from narwhals.utils import _FullContext
+
+
+class ArrowSelectorNamespace(CompliantSelectorNamespace["ArrowDataFrame", "ArrowSeries"]):
+    def _iter_columns(self, df: ArrowDataFrame) -> Iterator[ArrowSeries]:
+        from narwhals._arrow.series import ArrowSeries
+
+        for col, ser in zip(df.columns, df._native_frame.itercolumns()):
+            yield ArrowSeries(
+                ser, name=col, backend_version=df._backend_version, version=df._version
+            )
+
+    def _selector(
+        self,
+        context: _FullContext,
+        call: EvalSeries[ArrowDataFrame, ArrowSeries],
+        evaluate_output_names: EvalNames[ArrowDataFrame],
+        /,
+    ) -> CompliantSelector[ArrowDataFrame, ArrowSeries]:
+        return ArrowSelector(
+            call,
+            depth=0,
+            function_name="selector",
+            evaluate_output_names=evaluate_output_names,
+            alias_output_names=None,
+            backend_version=context._backend_version,
+            version=context._version,
         )
 
-    def categorical(self: Self) -> ArrowSelector:
-        dtypes = import_dtypes_module(self._version)
-        return self.by_dtype([dtypes.Categorical])
-
-    def string(self: Self) -> ArrowSelector:
-        dtypes = import_dtypes_module(self._version)
-        return self.by_dtype([dtypes.String])
-
-    def boolean(self: Self) -> ArrowSelector:
-        dtypes = import_dtypes_module(self._version)
-        return self.by_dtype([dtypes.Boolean])
-
-    def all(self: Self) -> ArrowSelector:
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            return [df[col] for col in df.columns]
-
-        return selector(self, func, lambda df: df.columns)
-
-    def datetime(
-        self: Self,
-        time_unit: TimeUnit | Iterable[TimeUnit] | None,
-        time_zone: str | timezone | Iterable[str | timezone | None] | None,
-    ) -> ArrowSelector:
-        dtypes = import_dtypes_module(version=self._version)
-        time_units, time_zones = _parse_time_unit_and_time_zone(
-            time_unit=time_unit, time_zone=time_zone
-        )
-
-        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            return [
-                df[col]
-                for col in df.columns
-                if dtype_matches_time_unit_and_time_zone(
-                    dtype=df.schema[col],
-                    dtypes=dtypes,
-                    time_units=time_units,
-                    time_zones=time_zones,
-                )
-            ]
-
-        def evaluate_output_names(df: ArrowDataFrame) -> Sequence[str]:
-            return [
-                col
-                for col in df.columns
-                if dtype_matches_time_unit_and_time_zone(
-                    dtype=df.schema[col],
-                    dtypes=dtypes,
-                    time_units=time_units,
-                    time_zones=time_zones,
-                )
-            ]
 
-        return selector(self, func, evaluate_output_names)
+class ArrowSelector(CompliantSelector["ArrowDataFrame", "ArrowSeries"], ArrowExpr):  # type: ignore[misc]
+    @property
+    def selectors(self) -> ArrowSelectorNamespace:
+        return ArrowSelectorNamespace(self)
 
-
-class ArrowSelector(ArrowExpr):
     def __repr__(self: Self) -> str:  # pragma: no cover
         return f"ArrowSelector(depth={self._depth}, function_name={self._function_name})"
 
@@ -135,82 +62,3 @@ def _to_expr(self: Self) -> ArrowExpr:
             backend_version=self._backend_version,
             version=self._version,
         )
-
-    def __sub__(self: Self, other: Self | Any) -> ArrowSelector | Any:
-        if isinstance(other, ArrowSelector):
-
-            def call(df: ArrowDataFrame) -> list[ArrowSeries]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                lhs = self._call(df)
-                return [x for x, name in zip(lhs, lhs_names) if name not in rhs_names]
-
-            def evaluate_output_names(df: ArrowDataFrame) -> list[str]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                return [x for x in lhs_names if x not in rhs_names]
-
-            return selector(self, call, evaluate_output_names)
-        else:
-            return self._to_expr() - other
-
-    def __or__(self: Self, other: Self | Any) -> ArrowSelector | Any:
-        if isinstance(other, ArrowSelector):
-
-            def call(df: ArrowDataFrame) -> list[ArrowSeries]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                lhs = self._call(df)
-                rhs = other._call(df)
-                return [
-                    *(x for x, name in zip(lhs, lhs_names) if name not in rhs_names),
-                    *rhs,
-                ]
-
-            def evaluate_output_names(df: ArrowDataFrame) -> list[str]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                return [*(x for x in lhs_names if x not in rhs_names), *rhs_names]
-
-            return selector(self, call, evaluate_output_names)
-        else:
-            return self._to_expr() | other
-
-    def __and__(self: Self, other: Self | Any) -> ArrowSelector | Any:
-        if isinstance(other, ArrowSelector):
-
-            def call(df: ArrowDataFrame) -> list[ArrowSeries]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                lhs = self._call(df)
-                return [x for x, name in zip(lhs, lhs_names) if name in rhs_names]
-
-            def evaluate_output_names(df: ArrowDataFrame) -> list[str]:
-                lhs_names = self._evaluate_output_names(df)
-                rhs_names = other._evaluate_output_names(df)
-                return [x for x in lhs_names if x in rhs_names]
-
-            return selector(self, call, evaluate_output_names)
-
-        else:
-            return self._to_expr() & other
-
-    def __invert__(self: Self) -> ArrowSelector:
-        return ArrowSelectorNamespace(self).all() - self
-
-
-def selector(
-    context: _LimitedContext,
-    call: Callable[[ArrowDataFrame], Sequence[ArrowSeries]],
-    evaluate_output_names: Callable[[ArrowDataFrame], Sequence[str]],
-    /,
-) -> ArrowSelector:
-    return ArrowSelector(
-        call,
-        depth=0,
-        function_name="selector",
-        evaluate_output_names=evaluate_output_names,
-        alias_output_names=None,
-        backend_version=context._backend_version,
-        version=context._version,
-    )

narwhals/_dask/dataframe.py

@@ -30,6 +30,7 @@
     from narwhals._dask.expr import DaskExpr
     from narwhals._dask.group_by import DaskLazyGroupBy
     from narwhals._dask.namespace import DaskNamespace
+    from narwhals._polars.dataframe import PolarsDataFrame
     from narwhals.dtypes import DType
     from narwhals.utils import Version
 
@@ -94,7 +95,7 @@ def collect(
         self: Self,
         backend: Implementation | None,
         **kwargs: Any,
-    ) -> CompliantDataFrame:
+    ) -> CompliantDataFrame[Any] | PolarsDataFrame:
         import pandas as pd
 
         result = self._native_frame.compute(**kwargs)

narwhals/_dask/namespace.py

@@ -24,6 +24,7 @@
 from narwhals._expression_parsing import combine_alias_output_names
 from narwhals._expression_parsing import combine_evaluate_output_names
 from narwhals.typing import CompliantNamespace
+from narwhals.utils import get_column_names
 from narwhals.utils import is_compliant_expr
 
 if TYPE_CHECKING:
@@ -57,7 +58,7 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
             func,
             depth=0,
             function_name="all",
-            evaluate_output_names=lambda df: df.columns,
+            evaluate_output_names=get_column_names,
             alias_output_names=None,
             backend_version=self._backend_version,
             version=self._version,

narwhals/_dask/selectors.py

@@ -10,6 +10,7 @@
 from narwhals._dask.expr import DaskExpr
 from narwhals.utils import _parse_time_unit_and_time_zone
 from narwhals.utils import dtype_matches_time_unit_and_time_zone
+from narwhals.utils import get_column_names
 from narwhals.utils import import_dtypes_module
 
 if TYPE_CHECKING:
@@ -95,7 +96,7 @@ def all(self: Self) -> DaskSelector:
         def func(df: DaskLazyFrame) -> list[dx.Series]:
             return [df._native_frame[col] for col in df.columns]
 
-        return selector(self, func, lambda df: df.columns)
+        return selector(self, func, get_column_names)
 
     def datetime(
         self: Self,

narwhals/_duckdb/dataframe.py

@@ -36,6 +36,7 @@
     from narwhals._duckdb.group_by import DuckDBGroupBy
     from narwhals._duckdb.namespace import DuckDBNamespace
     from narwhals._duckdb.series import DuckDBInterchangeSeries
+    from narwhals._polars.dataframe import PolarsDataFrame
     from narwhals.dtypes import DType
 
 from narwhals.typing import CompliantLazyFrame
@@ -90,7 +91,7 @@ def collect(
         self: Self,
         backend: ModuleType | Implementation | str | None,
         **kwargs: Any,
-    ) -> CompliantDataFrame:
+    ) -> CompliantDataFrame[Any] | PolarsDataFrame:
         if backend is None or backend is Implementation.PYARROW:
             import pyarrow as pa  # ignore-banned-import
 

narwhals/_duckdb/namespace.py

@@ -24,6 +24,7 @@
 from narwhals._expression_parsing import combine_alias_output_names
 from narwhals._expression_parsing import combine_evaluate_output_names
 from narwhals.typing import CompliantNamespace
+from narwhals.utils import get_column_names
 
 if TYPE_CHECKING:
     import duckdb
@@ -52,7 +53,7 @@ def _all(df: DuckDBLazyFrame) -> list[duckdb.Expression]:
         return DuckDBExpr(
             call=_all,
             function_name="all",
-            evaluate_output_names=lambda df: df.columns,
+            evaluate_output_names=get_column_names,
             alias_output_names=None,
             backend_version=self._backend_version,
             version=self._version,

narwhals/_duckdb/selectors.py

@@ -12,6 +12,7 @@
 from narwhals._duckdb.expr import DuckDBExpr
 from narwhals.utils import _parse_time_unit_and_time_zone
 from narwhals.utils import dtype_matches_time_unit_and_time_zone
+from narwhals.utils import get_column_names
 from narwhals.utils import import_dtypes_module
 
 if TYPE_CHECKING:
@@ -88,7 +89,7 @@ def all(self: Self) -> DuckDBSelector:
         def func(df: DuckDBLazyFrame) -> list[duckdb.Expression]:
             return [ColumnExpression(col) for col in df.columns]
 
-        return selector(self, func, lambda df: df.columns)
+        return selector(self, func, get_column_names)
 
     def datetime(
         self: Self,