Repository: narwhals-dev/narwhals
Pull request #562 (2024-07-21 08:59:30)

feat: arrow selectors
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

The easiness of these implementations is really a joy 


MarcoGorelli: legend @FBruzzesi !

there's just a ci failure on old pyarrow versions
```
FAILED tests/test_selectors.py::test_categorical[pyarrow_table_constructor] - pyarrow.lib.ArrowNotImplementedError: Unsupported cast from string to dictionary using function cast_dictionary
```
maybe we can just skip that test for pyarrow < 12?
FBruzzesi: >maybe we can just skip that test for pyarrow < 12?

I am investigating this right now. It seems to fail until pyarrow 15.0 😭

Patch:
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index 40e1d1b7..bc045565 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -8,6 +8,7 @@ from typing import Iterable
 from narwhals import dtypes
 from narwhals._arrow.dataframe import ArrowDataFrame
 from narwhals._arrow.expr import ArrowExpr
+from narwhals._arrow.selectors import ArrowSelectorNamespace
 from narwhals._arrow.series import ArrowSeries
 from narwhals._arrow.utils import horizontal_concat
 from narwhals._arrow.utils import vertical_concat
@@ -198,3 +199,7 @@ class ArrowNamespace:
                 backend_version=self._backend_version,
             )
         raise NotImplementedError
+
+    @property
+    def selectors(self) -> ArrowSelectorNamespace:
+        return ArrowSelectorNamespace(backend_version=self._backend_version)
diff --git a/narwhals/_arrow/selectors.py b/narwhals/_arrow/selectors.py
new file mode 100644
index 00000000..b8ba5bf5
--- /dev/null
+++ b/narwhals/_arrow/selectors.py
@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import NoReturn
+
+from narwhals import dtypes
+from narwhals._arrow.expr import ArrowExpr
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from narwhals._arrow.dataframe import ArrowDataFrame
+    from narwhals._arrow.series import ArrowSeries
+    from narwhals.dtypes import DType
+
+
+class ArrowSelectorNamespace:
+    def __init__(self: Self, *, backend_version: tuple[int, ...]) -> None:
+        self._backend_version = backend_version
+
+    def by_dtype(self: Self, dtypes: list[DType | type[DType]]) -> ArrowSelector:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            return [df[col] for col in df.columns if df.schema[col] in dtypes]
+
+        return ArrowSelector(
+            func,
+            depth=0,
+            function_name="type_selector",
+            root_names=None,
+            output_names=None,
+            backend_version=self._backend_version,
+        )
+
+    def numeric(self: Self) -> ArrowSelector:
+        return self.by_dtype(
+            [
+                dtypes.Int64,
+                dtypes.Int32,
+                dtypes.Int16,
+                dtypes.Int8,
+                dtypes.UInt64,
+                dtypes.UInt32,
+                dtypes.UInt16,
+                dtypes.UInt8,
+                dtypes.Float64,
+                dtypes.Float32,
+            ],
+        )
+
+    def categorical(self: Self) -> ArrowSelector:
+        return self.by_dtype([dtypes.Categorical])
+
+    def string(self: Self) -> ArrowSelector:
+        return self.by_dtype([dtypes.String])
+
+    def boolean(self: Self) -> ArrowSelector:
+        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 ArrowSelector(
+            func,
+            depth=0,
+            function_name="type_selector",
+            root_names=None,
+            output_names=None,
+            backend_version=self._backend_version,
+        )
+
+
+class ArrowSelector(ArrowExpr):
+    def __repr__(self: Self) -> str:  # pragma: no cover
+        return (
+            f"ArrowSelector("
+            f"depth={self._depth}, "
+            f"function_name={self._function_name}, "
+            f"root_names={self._root_names}, "
+            f"output_names={self._output_names}"
+        )
+
+    def _to_expr(self: Self) -> ArrowExpr:
+        return ArrowExpr(
+            self._call,
+            depth=self._depth,
+            function_name=self._function_name,
+            root_names=self._root_names,
+            output_names=self._output_names,
+            backend_version=self._backend_version,
+        )
+
+    def __sub__(self: Self, other: Self | Any) -> ArrowSelector | Any:
+        if isinstance(other, ArrowSelector):
+
+            def call(df: ArrowDataFrame) -> list[ArrowSeries]:
+                lhs = self._call(df)
+                rhs = other._call(df)
+                return [x for x in lhs if x.name not in [x.name for x in rhs]]
+
+            return ArrowSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+            )
+        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 = self._call(df)
+                rhs = other._call(df)
+                return [x for x in lhs if x.name not in [x.name for x in rhs]] + rhs
+
+            return ArrowSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+            )
+        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 = self._call(df)
+                rhs = other._call(df)
+                return [x for x in lhs if x.name in [x.name for x in rhs]]
+
+            return ArrowSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+            )
+        else:
+            return self._to_expr() & other
+
+    def __invert__(self: Self) -> ArrowSelector:
+        return ArrowSelectorNamespace(backend_version=self._backend_version).all() - self
+
+    def __rsub__(self: Self, other: Any) -> NoReturn:
+        raise NotImplementedError
+
+    def __rand__(self: Self, other: Any) -> NoReturn:
+        raise NotImplementedError
+
+    def __ror__(self: Self, other: Any) -> NoReturn:
+        raise NotImplementedError


Test patch:
diff --git a/tests/test_selectors.py b/tests/test_selectors.py
index 3ab5c5e0..41c721cb 100644
--- a/tests/test_selectors.py
+++ b/tests/test_selectors.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import Any
 
 import pandas as pd
-import polars as pl
+import pyarrow as pa
 import pytest
 
 import narwhals.stable.v1 as nw
@@ -13,6 +13,7 @@ from narwhals.selectors import by_dtype
 from narwhals.selectors import categorical
 from narwhals.selectors import numeric
 from narwhals.selectors import string
+from narwhals.utils import parse_version
 from tests.utils import compare_dicts
 
 data = {
@@ -23,54 +24,43 @@ data = {
 }
 
 
-def test_selectors(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_selectors(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
-    result = nw.to_native(df.select(by_dtype([nw.Int64, nw.Float64]) + 1))
+    result = df.select(by_dtype([nw.Int64, nw.Float64]) + 1)
     expected = {"a": [2, 2, 3], "c": [5.1, 6.0, 7.0]}
     compare_dicts(result, expected)
 
 
-def test_numeric(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_numeric(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
-    result = nw.to_native(df.select(numeric() + 1))
+    result = df.select(numeric() + 1)
     expected = {"a": [2, 2, 3], "c": [5.1, 6.0, 7.0]}
     compare_dicts(result, expected)
 
 
-def test_boolean(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_boolean(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
-    result = nw.to_native(df.select(boolean()))
+    result = df.select(boolean())
     expected = {"d": [True, False, True]}
     compare_dicts(result, expected)
 
 
-def test_string(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_string(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
-    result = nw.to_native(df.select(string()))
+    result = df.select(string())
     expected = {"b": ["a", "b", "c"]}
     compare_dicts(result, expected)
 
 
-def test_categorical() -> None:
-    df = nw.from_native(pd.DataFrame(data).astype({"b": "category"}))
-    result = nw.to_native(df.select(categorical()))
-    expected = {"b": ["a", "b", "c"]}
-    compare_dicts(result, expected)
-    df = nw.from_native(pl.DataFrame(data, schema_overrides={"b": pl.Categorical}))
-    result = nw.to_native(df.select(categorical()))
+def test_categorical(request: Any, constructor: Any) -> None:
+    if "pyarrow_table_constructor" in str(constructor) and parse_version(
+        pa.__version__
+    ) <= (15,):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
     expected = {"b": ["a", "b", "c"]}
+
+    df = nw.from_native(constructor(data)).with_columns(nw.col("b").cast(nw.Categorical))
+    result = df.select(categorical())
     compare_dicts(result, expected)
 
 
@@ -89,18 +79,16 @@ def test_categorical() -> None:
     ],
 )
 def test_set_ops(
-    request: Any, constructor: Any, selector: nw.selectors.Selector, expected: list[str]
+    constructor: Any, selector: nw.selectors.Selector, expected: list[str]
 ) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
     df = nw.from_native(constructor(data))
     result = df.select(selector).columns
     assert sorted(result) == expected
 
 
-def test_set_ops_invalid() -> None:
-    df = nw.from_native(pd.DataFrame(data))
+@pytest.mark.parametrize("invalid_constructor", [pd.DataFrame, pa.table])
+def test_set_ops_invalid(invalid_constructor: Any) -> None:
+    df = nw.from_native(invalid_constructor(data))
     with pytest.raises(NotImplementedError):
         df.select(1 - numeric())
     with pytest.raises(NotImplementedError):