Repository: narwhals-dev/narwhals
Pull request #2269 (2025-03-22 22:40:38)

fix: safety assertion was being thrown for Expr.is_last_distinct with over and _order_by by pandas and pyarrow


closes https://github.com/narwhals-dev/narwhals/issues/2268

<!--
# 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
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #\<issue number\>

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] 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 f5397854..d914b644 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -159,7 +159,7 @@ class ArrowExpr(EagerExpr["ArrowDataFrame", ArrowSeries]):
                 df = df.with_row_index(token).sort(
                     *order_by, descending=False, nulls_last=False
                 )
-                result = self(df)
+                result = self(df.drop([token], strict=True))
                 # TODO(marco): is there a way to do this efficiently without
                 # doing 2 sorts? Here we're sorting the dataframe and then
                 # again calling `sort_indices`. `ArrowSeries.scatter` would also sort.
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index b8766fec..c7c7c57b 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -209,7 +209,7 @@ class PandasLikeExpr(EagerExpr["PandasLikeDataFrame", PandasLikeSeries]):
                 df = df.with_row_index(token).sort(
                     *order_by, descending=False, nulls_last=False
                 )
-                results = self(df)
+                results = self(df.drop([token], strict=True))
                 sorting_indices = df[token]
                 for s in results:
                     s._scatter_in_place(sorting_indices, s)
diff --git a/narwhals/_spark_like/expr.py b/narwhals/_spark_like/expr.py
index 77467bfb..1e87c61c 100644
--- a/narwhals/_spark_like/expr.py
+++ b/narwhals/_spark_like/expr.py
@@ -156,7 +156,7 @@ class SparkLikeExpr(LazyExpr["SparkLikeLazyFrame", "Column"]):
         return result
 
     def _cum_window_func(
-        self: Self, *, reverse: bool, func_name: Literal["sum", "max"]
+        self: Self, *, reverse: bool, func_name: Literal["sum", "max", "min"]
     ) -> WindowFunction:
         def func(
             _input: Column, partition_by: Sequence[str], order_by: Sequence[str]
@@ -584,6 +584,11 @@ class SparkLikeExpr(LazyExpr["SparkLikeLazyFrame", "Column"]):
             self._cum_window_func(reverse=reverse, func_name="max")
         )
 
+    def cum_min(self, *, reverse: bool) -> Self:
+        return self._with_window_function(
+            self._cum_window_func(reverse=reverse, func_name="min")
+        )
+
     def fill_null(
         self,
         value: Any | None,
@@ -652,6 +657,5 @@ class SparkLikeExpr(LazyExpr["SparkLikeLazyFrame", "Column"]):
     is_first_distinct = not_implemented()
     is_last_distinct = not_implemented()
     cum_count = not_implemented()
-    cum_min = not_implemented()
     cum_prod = not_implemented()
     quantile = not_implemented()


Test patch:
diff --git a/tests/expr_and_series/cum_min_test.py b/tests/expr_and_series/cum_min_test.py
index 13791480..da768d79 100644
--- a/tests/expr_and_series/cum_min_test.py
+++ b/tests/expr_and_series/cum_min_test.py
@@ -4,7 +4,9 @@ import pytest
 
 import narwhals.stable.v1 as nw
 from tests.utils import PANDAS_VERSION
+from tests.utils import POLARS_VERSION
 from tests.utils import PYARROW_VERSION
+from tests.utils import Constructor
 from tests.utils import ConstructorEager
 from tests.utils import assert_equal_data
 
@@ -37,6 +39,196 @@ def test_cum_min_expr(
     assert_equal_data(result, {name: expected[name]})
 
 
+@pytest.mark.parametrize(
+    ("reverse", "expected_a"),
+    [
+        (False, [1, 2, 1]),
+        (True, [1, 1, 3]),
+    ],
+)
+def test_lazy_cum_min_grouped(
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+    *,
+    reverse: bool,
+    expected_a: list[int],
+) -> None:
+    if "duckdb" in str(constructor):
+        # no window function support yet in duckdb
+        request.applymarker(pytest.mark.xfail)
+    if "pyarrow_table" in str(constructor):
+        # grouped window functions not yet supported
+        request.applymarker(pytest.mark.xfail)
+    if "modin" in str(constructor):
+        pytest.skip(reason="probably bugged")
+    if "dask" in str(constructor):
+        # https://github.com/dask/dask/issues/11806
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor) and POLARS_VERSION < (1, 9):
+        pytest.skip(reason="too old version")
+    if "cudf" in str(constructor):
+        # https://github.com/rapidsai/cudf/issues/18159
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(
+        constructor(
+            {
+                "a": [1, 2, 3],
+                "b": [1, 0, 2],
+                "i": [0, 1, 2],
+                "g": [1, 1, 1],
+            }
+        )
+    )
+    result = df.with_columns(
+        nw.col("a").cum_min(reverse=reverse).over("g", _order_by="b")
+    ).sort("i")
+    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2]}
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("reverse", "expected_a"),
+    [
+        (False, [1, 2, 1, 1, 1, 2, 4]),
+        (True, [1, 1, 2, 1, 2, 1, 1]),
+    ],
+)
+def test_lazy_cum_min_ordered_by_nulls(
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+    *,
+    reverse: bool,
+    expected_a: list[int],
+) -> None:
+    if "duckdb" in str(constructor):
+        # no window function support yet in duckdb
+        request.applymarker(pytest.mark.xfail)
+    if "pyarrow_table" in str(constructor):
+        # grouped window functions not yet supported
+        request.applymarker(pytest.mark.xfail)
+    if "modin" in str(constructor):
+        pytest.skip(reason="probably bugged")
+    if "dask" in str(constructor):
+        # https://github.com/dask/dask/issues/11806
+        request.applymarker(pytest.mark.xfail)
+    if "polars" in str(constructor) and POLARS_VERSION < (1, 9):
+        pytest.skip(reason="too old version")
+    if "cudf" in str(constructor):
+        # https://github.com/rapidsai/cudf/issues/18159
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(
+        constructor(
+            {
+                "a": [1, 2, 3, 1, 2, 3, 4],
+                "b": [1, -1, 3, 2, 5, 0, None],
+                "i": [0, 1, 2, 3, 4, 5, 6],
+                "g": [1, 1, 1, 1, 1, 1, 1],
+            }
+        )
+    )
+    result = df.with_columns(
+        nw.col("a").cum_min(reverse=reverse).over("g", _order_by="b")
+    ).sort("i")
+    expected = {
+        "a": expected_a,
+        "b": [1, -1, 3, 2, 5, 0, None],
+        "i": [0, 1, 2, 3, 4, 5, 6],
+    }
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("reverse", "expected_a"),
+    [
+        (False, [1, 2, 1]),
+        (True, [1, 1, 3]),
+    ],
+)
+def test_lazy_cum_min_ungrouped(
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+    *,
+    reverse: bool,
+    expected_a: list[int],
+) -> None:
+    if "duckdb" in str(constructor):
+        # no window function support yet in duckdb
+        request.applymarker(pytest.mark.xfail)
+    if "dask" in str(constructor) and reverse:
+        # https://github.com/dask/dask/issues/11802
+        request.applymarker(pytest.mark.xfail)
+    if "modin" in str(constructor):
+        pytest.skip(reason="probably bugged")
+    if "polars" in str(constructor) and POLARS_VERSION < (1, 9):
+        pytest.skip(reason="too old version")
+    if PYARROW_VERSION < (13, 0, 0) and "pyarrow_table" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(
+        constructor(
+            {
+                "a": [2, 3, 1],
+                "b": [0, 2, 1],
+                "i": [1, 2, 0],
+            }
+        )
+    ).sort("i")
+    result = df.with_columns(
+        nw.col("a").cum_min(reverse=reverse).over(_order_by="b")
+    ).sort("i")
+    expected = {"a": expected_a, "b": [1, 0, 2], "i": [0, 1, 2]}
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("reverse", "expected_a"),
+    [
+        (False, [1, 2, 1, 1, 1, 2, 4]),
+        (True, [1, 1, 2, 1, 2, 1, 1]),
+    ],
+)
+def test_lazy_cum_min_ungrouped_ordered_by_nulls(
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+    *,
+    reverse: bool,
+    expected_a: list[int],
+) -> None:
+    if "duckdb" in str(constructor):
+        # no window function support yet in duckdb
+        request.applymarker(pytest.mark.xfail)
+    if "dask" in str(constructor):
+        # https://github.com/dask/dask/issues/11806
+        request.applymarker(pytest.mark.xfail)
+    if "modin" in str(constructor):
+        pytest.skip(reason="probably bugged")
+    if "polars" in str(constructor) and POLARS_VERSION < (1, 9):
+        pytest.skip(reason="too old version")
+    if PYARROW_VERSION < (13, 0, 0) and "pyarrow_table" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(
+        constructor(
+            {
+                "a": [1, 2, 3, 1, 2, 3, 4],
+                "b": [1, -1, 3, 2, 5, 0, None],
+                "i": [0, 1, 2, 3, 4, 5, 6],
+            }
+        )
+    ).sort("i")
+    result = df.with_columns(
+        nw.col("a").cum_min(reverse=reverse).over(_order_by="b")
+    ).sort("i")
+    expected = {
+        "a": expected_a,
+        "b": [1, -1, 3, 2, 5, 0, None],
+        "i": [0, 1, 2, 3, 4, 5, 6],
+    }
+    assert_equal_data(result, expected)
+
+
 def test_cum_min_series(
     request: pytest.FixtureRequest, constructor_eager: ConstructorEager
 ) -> None:
diff --git a/tests/expr_and_series/is_last_distinct_test.py b/tests/expr_and_series/is_last_distinct_test.py
index 4bfa7e13..ee42528a 100644
--- a/tests/expr_and_series/is_last_distinct_test.py
+++ b/tests/expr_and_series/is_last_distinct_test.py
@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import pytest
+
 import narwhals.stable.v1 as nw
+from tests.utils import POLARS_VERSION
 from tests.utils import ConstructorEager
 from tests.utils import assert_equal_data
 
@@ -20,6 +23,21 @@ def test_is_last_distinct_expr(constructor_eager: ConstructorEager) -> None:
     assert_equal_data(result, expected)
 
 
+def test_is_last_distinct_expr_all(constructor_eager: ConstructorEager) -> None:
+    # https://github.com/narwhals-dev/narwhals/issues/2268
+    if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 9):
+        pytest.skip(reason="too old version")
+    data = {"a": [1, 1, 2, 3, 2], "b": [1, 2, 3, 2, 1], "i": [0, 1, 2, 3, 4]}
+    df = nw.from_native(constructor_eager(data))
+    result = df.select(nw.all().is_last_distinct().over(_order_by="i"))
+    expected = {
+        "a": [False, True, False, True, True],
+        "b": [False, False, True, True, True],
+        "i": [True, True, True, True, True],
+    }
+    assert_equal_data(result, expected)
+
+
 def test_is_last_distinct_series(constructor_eager: ConstructorEager) -> None:
     series = nw.from_native(constructor_eager(data), eager_only=True)["a"]
     result = series.is_last_distinct()