Repository: narwhals-dev/narwhals
Pull request #733 (2024-08-07 13:40:18)

test: improve coverage for `from_native` (dask)
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [x] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## 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/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 844901da..a539040f 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -203,6 +203,9 @@ class ArrowExpr:
     def drop_nulls(self) -> Self:
         return reuse_series_implementation(self, "drop_nulls")
 
+    def shift(self, n: int) -> Self:
+        return reuse_series_implementation(self, "shift", n)
+
     def alias(self, name: str) -> Self:
         # Define this one manually, so that we can
         # override `output_names` and not increase depth
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index ff3ee94c..3c13be2d 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -221,6 +221,18 @@ class ArrowSeries:
         pc = get_pyarrow_compute()
         return self._from_native_series(pc.drop_null(self._native_series))
 
+    def shift(self, n: int) -> Self:
+        pa = get_pyarrow()
+        ca = self._native_series
+
+        if n > 0:
+            result = pa.concat_arrays([pa.nulls(n, ca.type), *ca[:-n].chunks])
+        elif n < 0:
+            result = pa.concat_arrays([*ca[-n:].chunks, pa.nulls(-n, ca.type)])
+        else:
+            result = ca
+        return self._from_native_series(result)
+
     def std(self, ddof: int = 1) -> int:
         pc = get_pyarrow_compute()
         return pc.stddev(self._native_series, ddof=ddof)  # type: ignore[no-any-return]
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 3d5cc281..00c2b05f 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -479,12 +479,10 @@ def from_native(  # noqa: PLR0915
     elif (dd := get_dask_dataframe()) is not None and isinstance(
         native_object, dd.DataFrame
     ):
-        if series_only:  # pragma: no cover
-            # TODO(unassigned): increase coverage
+        if series_only:
             msg = "Cannot only use `series_only` with dask DataFrame"
             raise TypeError(msg)
-        if eager_only or eager_or_interchange_only:  # pragma: no cover
-            # TODO(unassigned): increase coverage
+        if eager_only or eager_or_interchange_only:
             msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with dask DataFrame"
             raise TypeError(msg)
         if get_dask_expr() is None:  # pragma: no cover


Test patch:
diff --git a/tests/expr_and_series/shift_test.py b/tests/expr_and_series/shift_test.py
index 9460357d..02dbed6b 100644
--- a/tests/expr_and_series/shift_test.py
+++ b/tests/expr_and_series/shift_test.py
@@ -1,6 +1,6 @@
 from typing import Any
 
-import pytest
+import pyarrow as pa
 
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
@@ -13,10 +13,7 @@ data = {
 }
 
 
-def test_shift(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_shift(constructor: Any) -> None:
     df = nw.from_native(constructor(data))
     result = df.with_columns(nw.col("a", "b", "c").shift(2)).filter(nw.col("i") > 1)
     expected = {
@@ -28,21 +25,35 @@ def test_shift(request: Any, constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_shift_series(request: Any, constructor_eager: Any) -> None:
-    if "pyarrow_table" in str(constructor_eager):
-        request.applymarker(pytest.mark.xfail)
-
+def test_shift_series(constructor_eager: Any) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = df.with_columns(
+        df["a"].shift(2),
+        df["b"].shift(2),
+        df["c"].shift(2),
+    ).filter(nw.col("i") > 1)
     expected = {
         "i": [2, 3, 4],
         "a": [0, 1, 2],
         "b": [1, 2, 3],
         "c": [5, 4, 3],
     }
-    result = df.select(
-        df["i"],
-        df["a"].shift(2),
-        df["b"].shift(2),
-        df["c"].shift(2),
-    ).filter(nw.col("i") > 1)
+    compare_dicts(result, expected)
+
+
+def test_shift_multi_chunk_pyarrow() -> None:
+    tbl = pa.table({"a": [1, 2, 3]})
+    tbl = pa.concat_tables([tbl, tbl, tbl])
+    df = nw.from_native(tbl, eager_only=True)
+
+    result = df.select(nw.col("a").shift(1))
+    expected = {"a": [None, 1, 2, 3, 1, 2, 3, 1, 2]}
+    compare_dicts(result, expected)
+
+    result = df.select(nw.col("a").shift(-1))
+    expected = {"a": [2, 3, 1, 2, 3, 1, 2, 3, None]}
+    compare_dicts(result, expected)
+
+    result = df.select(nw.col("a").shift(0))
+    expected = {"a": [1, 2, 3, 1, 2, 3, 1, 2, 3]}
     compare_dicts(result, expected)
diff --git a/tests/translate/from_native_test.py b/tests/translate/from_native_test.py
index 191ad49b..3629b93f 100644
--- a/tests/translate/from_native_test.py
+++ b/tests/translate/from_native_test.py
@@ -169,3 +169,33 @@ def test_init_already_narwhals_unstable() -> None:
     s = df["a"]
     result_s = unstable_nw.from_native(s, allow_series=True)
     assert result_s is s
+
+
+def test_series_only_dask() -> None:
+    pytest.importorskip("dask")
+    pytest.importorskip("dask_expr")
+    import dask.dataframe as dd
+
+    dframe = dd.from_pandas(df_pd)
+
+    with pytest.raises(TypeError, match="Cannot only use `series_only`"):
+        nw.from_native(dframe, series_only=True)
+
+
+@pytest.mark.parametrize(
+    ("eager_only", "context"),
+    [
+        (False, does_not_raise()),
+        (True, pytest.raises(TypeError, match="Cannot only use `eager_only`")),
+    ],
+)
+def test_eager_only_lazy_dask(eager_only: Any, context: Any) -> None:
+    pytest.importorskip("dask")
+    pytest.importorskip("dask_expr")
+    import dask.dataframe as dd
+
+    dframe = dd.from_pandas(df_pd)
+
+    with context:
+        res = nw.from_native(dframe, eager_only=eager_only)
+        assert isinstance(res, nw.LazyFrame)