Skip to content

Commit a341154

Browse files
authored
fix: raise for rank followed by over with order_by for sql-like backends (#3178)
1 parent e2c4fd5 commit a341154

File tree

3 files changed

+65
-35
lines changed

3 files changed

+65
-35
lines changed

narwhals/_ibis/expr.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -336,7 +336,18 @@ def _rank(expr: ir.Column) -> ir.Value:
336336

337337
return ibis.cases((expr.notnull(), rank_))
338338

339-
return self._with_callable(_rank)
339+
def window_f(df: IbisLazyFrame, inputs: WindowInputs[ir.Value]) -> list[ir.Value]:
340+
if inputs.order_by:
341+
msg = "`rank` followed by `over` with `order_by` specified is not supported for Ibis backend."
342+
raise NotImplementedError(msg)
343+
return [
344+
_rank(cast("ir.Column", expr)).over(
345+
ibis.window(group_by=inputs.partition_by)
346+
)
347+
for expr in self(df)
348+
]
349+
350+
return self._with_callable(_rank, window_f)
340351

341352
@property
342353
def str(self) -> IbisExprStringNamespace:

narwhals/_sql/expr.py

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -755,17 +755,15 @@ def rank(self, method: RankMethod, *, descending: bool) -> Self:
755755
def _rank(
756756
expr: NativeExprT,
757757
partition_by: Sequence[str | NativeExprT] = (),
758-
order_by: Sequence[str | NativeExprT] = (),
759758
*,
760-
descending: Sequence[bool],
761-
nulls_last: Sequence[bool],
759+
descending: bool,
762760
) -> NativeExprT:
763761
count_expr = self._count_star()
764762
window_kwargs: dict[str, Any] = {
765763
"partition_by": partition_by,
766-
"order_by": (expr, *order_by),
767-
"descending": descending,
768-
"nulls_last": nulls_last,
764+
"order_by": (expr,),
765+
"descending": [descending],
766+
"nulls_last": [True],
769767
}
770768
count_window_kwargs: dict[str, Any] = {"partition_by": (*partition_by, expr)}
771769
window = self._window_expression
@@ -791,21 +789,16 @@ def _rank(
791789
return self._when(~F("isnull", expr), rank_expr) # type: ignore[operator]
792790

793791
def _unpartitioned_rank(expr: NativeExprT) -> NativeExprT:
794-
return _rank(expr, descending=[descending], nulls_last=[True])
792+
return _rank(expr, descending=descending)
795793

796794
def _partitioned_rank(
797795
df: SQLLazyFrameT, inputs: WindowInputs[NativeExprT]
798796
) -> Sequence[NativeExprT]:
799-
# node: when `descending` / `nulls_last` are supported in `.over`, they should be respected here
800-
# https://github.com/narwhals-dev/narwhals/issues/2790
797+
if inputs.order_by:
798+
msg = "`rank` followed by `over` with `order_by` specified is not supported for SQL-like backends."
799+
raise NotImplementedError(msg)
801800
return [
802-
_rank(
803-
expr,
804-
inputs.partition_by,
805-
inputs.order_by,
806-
descending=[descending] + [False] * len(inputs.order_by),
807-
nulls_last=[True] + [False] * len(inputs.order_by),
808-
)
801+
_rank(expr, inputs.partition_by, descending=descending)
809802
for expr in self(df)
810803
]
811804

tests/expr_and_series/rank_test.py

Lines changed: 44 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -340,19 +340,39 @@ def test_rank_with_order_by(
340340
if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
341341
pytest.skip(reason="too old version")
342342

343+
context = (
344+
pytest.raises(NotImplementedError)
345+
if any(x in str(constructor) for x in ("pyspark", "duckdb", "ibis"))
346+
else does_not_raise()
347+
)
348+
343349
df = nw.from_native(
344350
constructor(
345351
{"a": [1, 1, 2, 2, 3, 3], "b": [3, None, 4, 3, 5, 6], "i": list(range(6))}
346352
)
347353
)
348-
result = df.with_columns(c=nw.col("a").rank("ordinal").over(order_by="b")).sort("i")
349-
expected = {
350-
"a": [1, 1, 2, 2, 3, 3],
351-
"b": [3, None, 4, 3, 5, 6],
352-
"i": [0, 1, 2, 3, 4, 5],
353-
"c": [2, 1, 4, 3, 5, 6],
354-
}
355-
assert_equal_data(result, expected)
354+
with context:
355+
result = df.with_columns(c=nw.col("a").rank("ordinal").over(order_by="b")).sort(
356+
"i"
357+
)
358+
expected = {
359+
"a": [1, 1, 2, 2, 3, 3],
360+
"b": [3, None, 4, 3, 5, 6],
361+
"i": [0, 1, 2, 3, 4, 5],
362+
"c": [2, 1, 4, 3, 5, 6],
363+
}
364+
assert_equal_data(result, expected)
365+
366+
with context:
367+
# gh 3177
368+
df = nw.from_native(constructor({"i": [0, 1, 2], "j": [1, 2, 1]}))
369+
result = (
370+
df.with_columns(z=nw.col("j").rank("min").over(order_by="i"))
371+
.sort("i")
372+
.select("z")
373+
)
374+
expected = {"z": [1.0, 3.0, 1.0]}
375+
assert_equal_data(result, expected)
356376

357377

358378
def test_rank_with_order_by_and_partition_by(
@@ -386,14 +406,20 @@ def test_rank_with_order_by_and_partition_by(
386406
}
387407
)
388408
)
389-
result = df.with_columns(c=nw.col("a").rank("ordinal").over("g", order_by="b")).sort(
390-
"i"
409+
context = (
410+
pytest.raises(NotImplementedError)
411+
if any(x in str(constructor) for x in ("pyspark", "duckdb", "ibis"))
412+
else does_not_raise()
391413
)
392-
expected = {
393-
"a": [1, 1, 2, 2, 3, 3],
394-
"b": [3, None, 4, 3, 5, 6],
395-
"i": [0, 1, 2, 3, 4, 5],
396-
"g": ["x", "x", "x", "y", "y", "y"],
397-
"c": [2, 1, 3, 1, 2, 3],
398-
}
399-
assert_equal_data(result, expected)
414+
with context:
415+
result = df.with_columns(
416+
c=nw.col("a").rank("ordinal").over("g", order_by="b")
417+
).sort("i")
418+
expected = {
419+
"a": [1, 1, 2, 2, 3, 3],
420+
"b": [3, None, 4, 3, 5, 6],
421+
"i": [0, 1, 2, 3, 4, 5],
422+
"g": ["x", "x", "x", "y", "y", "y"],
423+
"c": [2, 1, 3, 1, 2, 3],
424+
}
425+
assert_equal_data(result, expected)

0 commit comments

Comments
 (0)