Repository: narwhals-dev/narwhals
Pull request #772 (2024-08-10 20:32:49)

feat: dask namespace `lit` method
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #637 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 7231b01d..676f6407 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -31,6 +31,7 @@
         - shape
         - sort
         - tail
+        - to_arrow
         - to_dict
         - to_numpy
         - to_pandas
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 2479196a..7b7f62b8 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -47,6 +47,7 @@
         - std
         - sum
         - tail
+        - to_arrow
         - to_dummies
         - to_frame
         - to_list
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 738732e4..ceadfdac 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -495,3 +495,6 @@ class ArrowDataFrame:
 
     def gather_every(self: Self, n: int, offset: int = 0) -> Self:
         return self._from_native_dataframe(self._native_dataframe[offset::n])
+
+    def to_arrow(self: Self) -> Any:
+        return self._native_dataframe
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index fadc72b8..bce0f471 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -615,6 +615,9 @@ class ArrowSeries:
 
         return self._from_native_series(arr)
 
+    def to_arrow(self: Self) -> Any:
+        return self._native_series.combine_chunks()
+
     @property
     def shape(self) -> tuple[int]:
         return (len(self._native_series),)
diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py
index f94e4779..82ade973 100644
--- a/narwhals/_dask/namespace.py
+++ b/narwhals/_dask/namespace.py
@@ -8,6 +8,7 @@ from typing import NoReturn
 
 from narwhals import dtypes
 from narwhals._dask.expr import DaskExpr
+from narwhals._dask.selectors import DaskSelectorNamespace
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals.dependencies import get_dask_dataframe
 from narwhals.dependencies import get_pandas
@@ -38,6 +39,10 @@ class DaskNamespace:
     Duration = dtypes.Duration
     Date = dtypes.Date
 
+    @property
+    def selectors(self) -> DaskSelectorNamespace:
+        return DaskSelectorNamespace(backend_version=self._backend_version)
+
     def __init__(self, *, backend_version: tuple[int, ...]) -> None:
         self._backend_version = backend_version
 
@@ -63,6 +68,19 @@ class DaskNamespace:
             backend_version=self._backend_version,
         )
 
+    def lit(self, value: Any, dtype: dtypes.DType | None) -> DaskExpr:
+        # TODO @FBruzzesi: cast to dtype once `reverse_translate_dtype` is implemented.
+        # It should be enough to add `.astype(reverse_translate_dtype(dtype))`
+        return DaskExpr(
+            lambda df: [df._native_dataframe.assign(lit=value).loc[:, "lit"]],
+            depth=0,
+            function_name="lit",
+            root_names=None,
+            output_names=["lit"],
+            returns_scalar=False,
+            backend_version=self._backend_version,
+        )
+
     def min(self, *column_names: str) -> DaskExpr:
         return DaskExpr.from_column_names(
             *column_names,
diff --git a/narwhals/_dask/selectors.py b/narwhals/_dask/selectors.py
new file mode 100644
index 00000000..11d0f0c7
--- /dev/null
+++ b/narwhals/_dask/selectors.py
@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import NoReturn
+
+from narwhals import dtypes
+from narwhals._dask.expr import DaskExpr
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from narwhals._dask.dataframe import DaskLazyFrame
+    from narwhals.dtypes import DType
+
+
+class DaskSelectorNamespace:
+    def __init__(self: Self, *, backend_version: tuple[int, ...]) -> None:
+        self._backend_version = backend_version
+
+    def by_dtype(self: Self, dtypes: list[DType | type[DType]]) -> DaskSelector:
+        def func(df: DaskLazyFrame) -> list[Any]:
+            return [
+                df._native_dataframe[col]
+                for col in df.columns
+                if df.schema[col] in dtypes
+            ]
+
+        return DaskSelector(
+            func,
+            depth=0,
+            function_name="type_selector",
+            root_names=None,
+            output_names=None,
+            backend_version=self._backend_version,
+            returns_scalar=False,
+        )
+
+    def numeric(self: Self) -> DaskSelector:
+        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) -> DaskSelector:
+        return self.by_dtype([dtypes.Categorical])
+
+    def string(self: Self) -> DaskSelector:
+        return self.by_dtype([dtypes.String])
+
+    def boolean(self: Self) -> DaskSelector:
+        return self.by_dtype([dtypes.Boolean])
+
+    def all(self: Self) -> DaskSelector:
+        def func(df: DaskLazyFrame) -> list[Any]:
+            return [df._native_dataframe[col] for col in df.columns]
+
+        return DaskSelector(
+            func,
+            depth=0,
+            function_name="type_selector",
+            root_names=None,
+            output_names=None,
+            backend_version=self._backend_version,
+            returns_scalar=False,
+        )
+
+
+class DaskSelector(DaskExpr):
+    def __repr__(self: Self) -> str:  # pragma: no cover
+        return (
+            f"DaskSelector("
+            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) -> DaskExpr:
+        return DaskExpr(
+            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,
+            returns_scalar=self._returns_scalar,
+        )
+
+    def __sub__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any:
+        if isinstance(other, DaskSelector):
+
+            def call(df: DaskLazyFrame) -> list[Any]:
+                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 DaskSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+                returns_scalar=self._returns_scalar,
+            )
+        else:
+            return self._to_expr() - other
+
+    def __or__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any:
+        if isinstance(other, DaskSelector):
+
+            def call(df: DaskLazyFrame) -> list[Any]:
+                lhs = self._call(df)
+                rhs = other._call(df)
+                return [  # type: ignore[no-any-return]
+                    x for x in lhs if x.name not in [x.name for x in rhs]
+                ] + rhs
+
+            return DaskSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+                returns_scalar=self._returns_scalar,
+            )
+        else:
+            return self._to_expr() | other
+
+    def __and__(self: Self, other: DaskSelector | Any) -> DaskSelector | Any:
+        if isinstance(other, DaskSelector):
+
+            def call(df: DaskLazyFrame) -> list[Any]:
+                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 DaskSelector(
+                call,
+                depth=0,
+                function_name="type_selector",
+                root_names=None,
+                output_names=None,
+                backend_version=self._backend_version,
+                returns_scalar=self._returns_scalar,
+            )
+        else:
+            return self._to_expr() & other
+
+    def __invert__(self: Self) -> DaskSelector:
+        return DaskSelectorNamespace(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
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 39af2821..d36a381b 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -20,6 +20,7 @@ from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
+from narwhals.dependencies import get_pyarrow
 from narwhals.utils import Implementation
 from narwhals.utils import flatten
 from narwhals.utils import generate_unique_token
@@ -592,3 +593,11 @@ class PandasLikeDataFrame:
 
     def gather_every(self: Self, n: int, offset: int = 0) -> Self:
         return self._from_native_dataframe(self._native_dataframe.iloc[offset::n])
+
+    def to_arrow(self: Self) -> Any:
+        if self._implementation is Implementation.CUDF:  # pragma: no cover
+            msg = "`to_arrow` is not implemented for CuDF backend."
+            raise NotImplementedError(msg)
+
+        pa = get_pyarrow()
+        return pa.Table.from_pandas(self._native_dataframe)
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index e52a034e..9a80e26f 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -17,6 +17,7 @@ from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
+from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
 from narwhals.utils import Implementation
 
@@ -638,6 +639,14 @@ class PandasLikeSeries:
             self._native_series.clip(lower_bound, upper_bound)
         )
 
+    def to_arrow(self: Self) -> Any:
+        if self._implementation is Implementation.CUDF:  # pragma: no cover
+            msg = "`to_arrow` is not implemented for CuDF backend."
+            raise NotImplementedError(msg)
+
+        pa = get_pyarrow()
+        return pa.Array.from_pandas(self._native_series)
+
     @property
     def str(self) -> PandasLikeSeriesStringNamespace:
         return PandasLikeSeriesStringNamespace(self)
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index f8a12a39..b40aa160 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -2020,6 +2020,42 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().gather_every(n=n, offset=offset)
 
+    def to_arrow(self: Self) -> Any:
+        r"""
+        Convert to arrow table.
+
+        Examples:
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data = {"foo": [1, 2, 3], "bar": ["a", "b", "c"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            Let's define a dataframe-agnostic function that converts to arrow table:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.to_arrow()
+
+            >>> func(df_pd)  # doctest:+SKIP
+            pyarrow.Table
+            foo: int64
+            bar: string
+            ----
+            foo: [[1,2,3]]
+            bar: [["a","b","c"]]
+
+            >>> func(df_pl)  # doctest:+NORMALIZE_WHITESPACE
+            pyarrow.Table
+            foo: int64
+            bar: large_string
+            ----
+            foo: [[1,2,3]]
+            bar: [["a","b","c"]]
+        """
+        return self._compliant_frame.to_arrow()
+
 
 class LazyFrame(BaseFrame[FrameT]):
     """
diff --git a/narwhals/series.py b/narwhals/series.py
index 76dd71e6..a6351f20 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -2227,6 +2227,45 @@ class Series:
             self._compliant_series.gather_every(n=n, offset=offset)
         )
 
+    def to_arrow(self: Self) -> Any:
+        r"""
+        Convert to arrow.
+
+        Examples:
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data = [1, 2, 3, 4]
+            >>> s_pd = pd.Series(name="a", data=data)
+            >>> s_pl = pl.Series(name="a", values=data)
+
+            Let's define a dataframe-agnostic function that converts to arrow:
+
+            >>> @nw.narwhalify
+            ... def func(s):
+            ...     return s.to_arrow()
+
+            >>> func(s_pd)  # doctest:+NORMALIZE_WHITESPACE
+            <pyarrow.lib.Int64Array object at ...>
+            [
+                1,
+                2,
+                3,
+                4
+            ]
+
+            >>> func(s_pl)  # doctest:+NORMALIZE_WHITESPACE
+            <pyarrow.lib.Int64Array object at ...>
+            [
+                1,
+                2,
+                3,
+                4
+            ]
+        """
+
+        return self._compliant_series.to_arrow()
+
     @property
     def str(self) -> SeriesStringNamespace:
         return SeriesStringNamespace(self)
diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py
index 80ee5d7a..68c98008 100644
--- a/utils/check_api_reference.py
+++ b/utils/check_api_reference.py
@@ -148,6 +148,7 @@ if (
     .difference(expr)
     .difference(
         {
+            "to_arrow",
             "to_dummies",
             "to_pandas",
             "to_list",


Test patch:
diff --git a/tests/frame/lit_test.py b/tests/frame/lit_test.py
index 212cffd6..328e4d8e 100644
--- a/tests/frame/lit_test.py
+++ b/tests/frame/lit_test.py
@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 def test_lit(
     constructor: Any, dtype: DType | None, expected_lit: list[Any], request: Any
 ) -> None:
-    if "dask" in str(constructor):
+    if "dask" in str(constructor) and dtype == nw.String:
         request.applymarker(pytest.mark.xfail)
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
     df_raw = constructor(data)
diff --git a/tests/frame/to_arrow_test.py b/tests/frame/to_arrow_test.py
new file mode 100644
index 00000000..c1f395e5
--- /dev/null
+++ b/tests/frame/to_arrow_test.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import narwhals.stable.v1 as nw
+from narwhals.utils import parse_version
+
+
+def test_to_arrow(request: Any, constructor_eager: Any) -> None:
+    if "pandas" in str(constructor_eager) and parse_version(pd.__version__) < (1, 0, 0):
+        # pyarrow requires pandas>=1.0.0
+        request.applymarker(pytest.mark.xfail)
+
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]}
+    df_raw = constructor_eager(data)
+    result = nw.from_native(df_raw, eager_only=True).to_arrow()
+
+    expected = pa.table(data)
+    assert result == expected
diff --git a/tests/series_only/to_arrow_test.py b/tests/series_only/to_arrow_test.py
new file mode 100644
index 00000000..ebd90b7c
--- /dev/null
+++ b/tests/series_only/to_arrow_test.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+import narwhals.stable.v1 as nw
+
+
+def test_to_arrow(constructor_eager: Any) -> None:
+    data = [1, 2, 3]
+    result = nw.from_native(constructor_eager({"a": data}), eager_only=True)[
+        "a"
+    ].to_arrow()
+
+    assert pa.types.is_int64(result.type)
+    assert pc.all(pc.equal(result, pa.array(data, type=pa.int64())))
+
+
+def test_to_arrow_with_nulls(constructor_eager: Any, request: Any) -> None:
+    if "pandas_constructor" in str(constructor_eager) or "modin_constructor" in str(
+        constructor_eager
+    ):
+        request.applymarker(pytest.mark.xfail)
+
+    data = [1, 2, None]
+    result = (
+        nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"]
+        .cast(nw.Int64)
+        .to_arrow()
+    )
+
+    assert pa.types.is_int64(result.type)
+    assert pc.all(pc.equal(result, pa.array(data, type=pa.int64())))
diff --git a/tests/test_selectors.py b/tests/test_selectors.py
index 5c57a967..dcd9e7ec 100644
--- a/tests/test_selectors.py
+++ b/tests/test_selectors.py
@@ -7,6 +7,7 @@ import pyarrow as pa
 import pytest
 
 import narwhals.stable.v1 as nw
+from narwhals.dependencies import get_dask
 from narwhals.selectors import all
 from narwhals.selectors import boolean
 from narwhals.selectors import by_dtype
@@ -24,36 +25,28 @@ data = {
 }
 
 
-def test_selectors(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_selectors(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     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(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_numeric(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.select(numeric() + 1)
     expected = {"a": [2, 2, 3], "c": [5.1, 6.0, 7.0]}
     compare_dicts(result, expected)
 
 
-def test_boolean(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_boolean(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.select(boolean())
     expected = {"d": [True, False, True]}
     compare_dicts(result, expected)
 
 
-def test_string(constructor: Any, request: Any) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_string(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.select(string())
     expected = {"b": ["a", "b", "c"]}
@@ -74,6 +67,17 @@ def test_categorical(request: Any, constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
+@pytest.mark.skipif((get_dask() is None), reason="too old for dask")
+def test_dask_categorical() -> None:
+    import dask.dataframe as dd
+
+    expected = {"b": ["a", "b", "c"]}
+    df_raw = dd.from_dict(expected, npartitions=1).astype({"b": "category"})
+    df = nw.from_native(df_raw)
+    result = df.select(categorical())
+    compare_dicts(result, expected)
+
+
 @pytest.mark.parametrize(
     ("selector", "expected"),
     [
@@ -89,10 +93,8 @@ def test_categorical(request: Any, constructor: Any) -> None:
     ],
 )
 def test_set_ops(
-    constructor: Any, selector: nw.selectors.Selector, expected: list[str], request: Any
+    constructor: Any, selector: nw.selectors.Selector, expected: list[str]
 ) -> None:
-    if "dask" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
     df = nw.from_native(constructor(data))
     result = df.select(selector).collect_schema().names()
     assert sorted(result) == expected