Incorrect results for `with_row_index(order_by=...)`

[dangotbanned]

Noticed while looking at the pyarrow version for (#3224) that the index seems to be wrong when the column isn't already ordered.

Possibly a regression from #3239? (cc @MarcoGorelli)
Just guessing as that's the most recent change

Repro

pandas, pyarrow and polars all produce the same result, which I think is wrong?

import narwhals as nw

data = {"a": ["A", "B", "A"], "b": [1, 2, 3], "c": [9, 2, 4]}

df = nw.from_dict(data, backend="polars")
df.with_row_index(order_by="c").sort("index")

┌───────────────────────────┐
|    Narwhals DataFrame     |
|---------------------------|
|shape: (3, 4)              |
|┌───────┬─────┬─────┬─────┐|
|│ index ┆ a   ┆ b   ┆ c   │|
|│ ---   ┆ --- ┆ --- ┆ --- │|
|│ i64   ┆ str ┆ i64 ┆ i64 │|
|╞═══════╪═════╪═════╪═════╡|
|│ 0     ┆ A   ┆ 3   ┆ 4   │|
|│ 1     ┆ A   ┆ 1   ┆ 9   │|
|│ 2     ┆ B   ┆ 2   ┆ 2   │|
|└───────┴─────┴─────┴─────┘|
└───────────────────────────┘

Expected

duckdb, ibis, sqlframe, dask all produce this - which is an index ordered by "c"

df.lazy("duckdb").with_row_index(order_by="c").sort("index").collect("polars")

┌───────────────────────────┐
|    Narwhals DataFrame     |
|---------------------------|
|shape: (3, 4)              |
|┌───────┬─────┬─────┬─────┐|
|│ index ┆ a   ┆ b   ┆ c   │|
|│ ---   ┆ --- ┆ --- ┆ --- │|
|│ i64   ┆ str ┆ i64 ┆ i64 │|
|╞═══════╪═════╪═════╪═════╡|
|│ 0     ┆ B   ┆ 2   ┆ 2   │|
|│ 1     ┆ A   ┆ 3   ┆ 4   │|
|│ 2     ┆ A   ┆ 1   ┆ 9   │|
|└───────┴─────┴─────┴─────┘|
└───────────────────────────┘

[MarcoGorelli]

thanks, taking a look

[MarcoGorelli]

for polars we can do

-            end = pl.count() if self._backend_version < (0, 20, 5) else pl.len()
-            result = frame.select(
-                pl.int_range(start=0, end=end).sort_by(order_by).alias(name), pl.all()
-            )
+            result = frame.select(pl.int_range(pl.len()).over(order_by=order_by).alias(name), pl.all())

[MarcoGorelli]

I don't think the original logic was correct either:

data = {
    'i': [2, 1, 1, 2],
    'j': [5, 3, 4, 1]
}

df = nw.from_dict(data, backend="pandas")
print(df.with_row_index(order_by=['j', 'i']).sort("i"))
print(df.lazy("duckdb").with_row_index(order_by=['j', 'i']).sort("i").collect("pandas"))

┌──────────────────┐
|Narwhals DataFrame|
|------------------|
|     index  i  j  |
|  1    1.0  1  3  |
|  2    2.0  1  4  |
|  3    3.0  2  5  |
|  0    0.0  2  1  |
└──────────────────┘
┌──────────────────┐
|Narwhals DataFrame|
|------------------|
|     index  i  j  |
|  0      1  1  3  |
|  1      2  1  4  |
|  2      0  2  1  |
|  3      3  2  5  |
└──────────────────┘

[MarcoGorelli]

ah nevermind, 'j' doesn't match either

not sure, thinking about htis

[dangotbanned]

#3289 (comment)

There's some examples of int_range for pandas and pyarrow in (#2895) if that sparks any ideas 😄

[MarcoGorelli]

ok i can't break it 😄 I think the previous implementation was fine - sorry @FBruzzesi ! will revert + add some meaner tests

[FBruzzesi]

ok i can't break it 😄 I think the previous implementation was fine - sorry @FBruzzesi ! will revert + add some meaner tests

Nothing to be sorry about! Thanks both for spotting the issue, fixing it, and provide additional tests 🙏🏼