Repository: narwhals-dev/narwhals
Pull request #590 (2024-07-23 10:27:04)

Feat: `Series.shift` Pyarrow Backend Implementation 
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [x] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.



MarcoGorelli: thanks @mfonekpo  - could you put a descriptive title please?
mfonekpo: > thanks @mfonekpo  - could you put a descriptive title please?

Yes Marco. Done.
MarcoGorelli: thanks @mfonekpo - could you name the test so it ends with `_test.py` please? else it won't run
mfonekpo: > thanks @mfonekpo - could you name the test so it ends with `_test.py` please? else it won't run

Ohhhh....how did I miss that. 
mfonekpo: Alright then.

On Thu, 25 Jul 2024, 06:26 Francesco Bruzzesi, ***@***.***>
wrote:

> ***@***.**** commented on this pull request.
> ------------------------------
>
> In narwhals/_arrow/series.py
> <https://github.com/narwhals-dev/narwhals/pull/590#discussion_r1690834396>
> :
>
> > +    def shift(self, n: int) -> Self:
> +        pc = get_pyarrow_compute()
> +        return self._from_native_series(pc.shift(self._native_series, n))
>
> This might end up being not so straightforward after all. We should
> *manually* slice the original array and build the new one, this will also
> be dependent on the sign of n
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/narwhals-dev/narwhals/pull/590#pullrequestreview-2198339943>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AOA4GJDXMJF737IFUBBU7QTZOCEA7AVCNFSM6AAAAABLKD7LZSVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDCOJYGMZTSOJUGM>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

mfonekpo: There is room for growth, right? Thank you. Is it ok? because you said it
is out of interest. What do I do?

On Sun, Jul 28, 2024 at 2:09 PM Marco Edward Gorelli <
***@***.***> wrote:

> ***@***.**** commented on this pull request.
> ------------------------------
>
> In narwhals/_arrow/series.py
> <https://github.com/narwhals-dev/narwhals/pull/590#discussion_r1694244426>
> :
>
> > +        # This condition shifts the series by n positions.
> +        if n > 0:
> +            # If n is positive, prepend nulls and truncate the end
> +            shifted_chunks = [pa.array([None] * n, type=dtype)] + [
> +                chunk.slice(0, len(chunk) - n) for chunk in chunks
> +            ]
> +        elif n < 0:
> +            # If n is negative, append nulls and truncate the start
> +            shifted_chunks = [chunk.slice(-n, len(chunk) + n) for chunk in chunks] + [
> +                pa.array([None] * -n, type=dtype)
> +            ]
> +        else:
> +            # If n is 0, return the original series
> +            shifted_chunks = chunks
>
> cool, thanks - I'm just a bit surprised, as it looks a bit more advanced
> than other work you've submitted. if you figured this all out, then well
> done 🙌
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/narwhals-dev/narwhals/pull/590#discussion_r1694244426>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AOA4GJGVBBK57PHKQ75XC43ZOTUQXAVCNFSM6AAAAABLKD7LZSVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDEMBTGQ4DKOBYGE>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

MarcoGorelli: yup, definitely room for growth :) it's OK, if you add the test I suggested and it passes, then it should be fine
mfonekpo: Thanks, I will!

On Sun, Jul 28, 2024 at 2:56 PM Marco Edward Gorelli <
***@***.***> wrote:

> yup, definitely room for growth :) it's OK, if you add the test I
> suggested and it passes, then it should be fine
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/narwhals-dev/narwhals/pull/590#issuecomment-2254527963>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AOA4GJGEKAN6PIUI3FGRY3DZOT2AZAVCNFSM6AAAAABLKD7LZSVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENJUGUZDOOJWGM>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>


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]


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)