Repository: ibis-project/ibis

Issue:
BUG: Same window op with different window size on table lead to incorrect results for pyspark backend
Currently, this test in `window.py` will fail
```
def test_multiple_windows(client):
    table = client.table('time_indexed_table')
    window1 = ibis.trailing_window(
        preceding=ibis.interval(hours=1), order_by='time', group_by='key'
    )
    window2 = ibis.trailing_window(
        preceding=ibis.interval(hours=2), order_by='time', group_by='key'
    )
    result = table.mutate(
        mean_1h=table['value'].mean().over(window1),
        mean_2h=table['value'].mean().over(window2),
    ).compile()
    result_pd = result.toPandas()

    df = table.compile().toPandas()
    expected_win_1 = (
        df.set_index('time')
        .groupby('key')
        .value.rolling('1h', closed='both')
        .mean()
        .rename('mean_1h')
    ).reset_index(drop=True)
    expected_win_2 = (
        df.set_index('time')
        .groupby('key')
        .value.rolling('2h', closed='both')
        .mean()
        .rename('mean_2h')
    ).reset_index(drop=True)
    tm.assert_series_equal(result_pd['mean_1h'], expected_win_1)
    tm.assert_series_equal(result_pd['mean_2h'], expected_win_2)
```
The compiled result in `result_pd` is wrong and the two columns of different window sizes are the same. This is due to a bug in compiling window ops in pyspark backend.


Discussion:

Resolving patch:
diff --git a/ibis/pyspark/compiler.py b/ibis/pyspark/compiler.py
--- a/ibis/pyspark/compiler.py
+++ b/ibis/pyspark/compiler.py
@@ -374,11 +374,8 @@ def compile_aggregator(
         src_col = F.when(condition, src_col)
 
     col = fn(src_col)
-    if context in (AggregationContext.ENTIRE, AggregationContext.GROUP):
+    if context:
         return col
-    elif context == AggregationContext.WINDOW:
-        window = kwargs['window']
-        return col.over(window)
     else:
         # We are trying to compile a expr such as some_col.max()
         # to a Spark expression.
@@ -410,9 +407,7 @@ def compile_any(t, expr, scope, timecontext, context=None, **kwargs):
 
 
 @compiles(ops.NotAny)
-def compile_notany(
-    t, expr, scope, timecontext, *, context=None, window=None, **kwargs
-):
+def compile_notany(t, expr, scope, timecontext, *, context=None, **kwargs):
     # The code here is a little ugly because the translation are different
     # with different context.
     # When translating col.notany() (context is None), we returns the dataframe
@@ -430,13 +425,7 @@ def fn(col):
         )
     else:
         return ~compile_any(
-            t,
-            expr,
-            scope,
-            timecontext,
-            context=context,
-            window=window,
-            **kwargs,
+            t, expr, scope, timecontext, context=context, **kwargs,
         )
 
 
@@ -448,9 +437,7 @@ def compile_all(t, expr, scope, timecontext, context=None, **kwargs):
 
 
 @compiles(ops.NotAll)
-def compile_notall(
-    t, expr, scope, timecontext, *, context=None, window=None, **kwargs
-):
+def compile_notall(t, expr, scope, timecontext, *, context=None, **kwargs):
     # See comments for opts.NotAny for reasoning for the if/else
     if context is None:
 
@@ -462,13 +449,7 @@ def fn(col):
         )
     else:
         return ~compile_all(
-            t,
-            expr,
-            scope,
-            timecontext,
-            context=context,
-            window=window,
-            **kwargs,
+            t, expr, scope, timecontext, context=context, **kwargs,
         )
 
 
@@ -1111,15 +1092,26 @@ def compile_window_op(t, expr, scope, timecontext, **kwargs):
         else:
             pyspark_window = pyspark_window.rowsBetween(start, end)
 
-    result = t.translate(
-        operand, scope, timecontext, window=pyspark_window, context=context
+    res_op = operand.op()
+    if isinstance(res_op, (ops.NotAll, ops.NotAny)):
+        # For NotAll and NotAny, negation must be applied after .over(window)
+        # Here we rewrite node to be its negation, and negate it back after
+        # translation and window operation
+        operand = res_op.negate().to_expr()
+    result = t.translate(operand, scope, timecontext, context=context).over(
+        pyspark_window
     )
-    return result
 
+    if isinstance(res_op, (ops.NotAll, ops.NotAny)):
+        return ~result
+    elif isinstance(res_op, (ops.MinRank, ops.DenseRank, ops.RowNumber)):
+        # result must be cast to long type for Rank / RowNumber
+        return result.astype('long') - 1
+    else:
+        return result
 
-def _handle_shift_operation(
-    t, expr, scope, timecontext, *, fn, window, **kwargs
-):
+
+def _handle_shift_operation(t, expr, scope, timecontext, *, fn, **kwargs):
     op = expr.op()
 
     src_column = t.translate(op.arg, scope, timecontext)
@@ -1127,37 +1119,37 @@ def _handle_shift_operation(
     offset = op.offset.op().value if op.offset is not None else op.offset
 
     if offset:
-        return fn(src_column, count=offset, default=default).over(window)
+        return fn(src_column, count=offset, default=default)
     else:
-        return fn(src_column, default=default).over(window)
+        return fn(src_column, default=default)
 
 
 @compiles(ops.Lag)
-def compile_lag(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_lag(t, expr, scope, timecontext, **kwargs):
     return _handle_shift_operation(
-        t, expr, scope, timecontext, fn=F.lag, window=window, **kwargs
+        t, expr, scope, timecontext, fn=F.lag, **kwargs
     )
 
 
 @compiles(ops.Lead)
-def compile_lead(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_lead(t, expr, scope, timecontext, **kwargs):
     return _handle_shift_operation(
-        t, expr, scope, timecontext, fn=F.lead, window=window, **kwargs
+        t, expr, scope, timecontext, fn=F.lead, **kwargs
     )
 
 
 @compiles(ops.MinRank)
-def compile_rank(t, expr, scope, timecontext, *, window, **kwargs):
-    return F.rank().over(window).astype('long') - 1
+def compile_rank(t, expr, scope, timecontext, **kwargs):
+    return F.rank()
 
 
 @compiles(ops.DenseRank)
-def compile_dense_rank(t, expr, scope, timecontext, *, window, **kwargs):
-    return F.dense_rank().over(window).astype('long') - 1
+def compile_dense_rank(t, expr, scope, timecontext, **kwargs):
+    return F.dense_rank()
 
 
 @compiles(ops.PercentRank)
-def compile_percent_rank(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_percent_rank(t, expr, scope, timecontext, **kwargs):
     raise com.UnsupportedOperationError(
         'Pyspark percent_rank() function indexes from 0 '
         'instead of 1, and does not match expected '
@@ -1166,29 +1158,29 @@ def compile_percent_rank(t, expr, scope, timecontext, *, window, **kwargs):
 
 
 @compiles(ops.NTile)
-def compile_ntile(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_ntile(t, expr, scope, timecontext, **kwargs):
     op = expr.op()
     buckets = op.buckets.op().value
-    return F.ntile(buckets).over(window)
+    return F.ntile(buckets)
 
 
 @compiles(ops.FirstValue)
-def compile_first_value(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_first_value(t, expr, scope, timecontext, **kwargs):
     op = expr.op()
     src_column = t.translate(op.arg, scope, timecontext)
-    return F.first(src_column).over(window)
+    return F.first(src_column)
 
 
 @compiles(ops.LastValue)
-def compile_last_value(t, expr, scope, timecontext, *, window, **kwargs):
+def compile_last_value(t, expr, scope, timecontext, **kwargs):
     op = expr.op()
     src_column = t.translate(op.arg, scope, timecontext)
-    return F.last(src_column).over(window)
+    return F.last(src_column)
 
 
 @compiles(ops.RowNumber)
-def compile_row_number(t, expr, scope, timecontext, *, window, **kwargs):
-    return F.row_number().over(window).cast('long') - 1
+def compile_row_number(t, expr, scope, timecontext, **kwargs):
+    return F.row_number()
 
 
 # -------------------------- Temporal Operations ----------------------------
@@ -1658,11 +1650,8 @@ def compile_reduction_udf(t, expr, scope, timecontext, context=None, **kwargs):
     func_args = (t.translate(arg, scope, timecontext) for arg in op.func_args)
 
     col = spark_udf(*func_args)
-    if context in (AggregationContext.ENTIRE, AggregationContext.GROUP):
+    if context:
         return col
-    elif context == AggregationContext.WINDOW:
-        window = kwargs['window']
-        return col.over(window)
     else:
         src_table = t.translate(op.func_args[0].op().table, scope, timecontext)
         return src_table.agg(col)


Test patch:
diff --git a/ibis/pyspark/tests/test_window.py b/ibis/pyspark/tests/test_window.py
--- a/ibis/pyspark/tests/test_window.py
+++ b/ibis/pyspark/tests/test_window.py
@@ -58,12 +58,6 @@ def test_time_indexed_window(client, ibis_window, spark_range):
     tm.assert_frame_equal(result_pd, expected)
 
 
-# TODO: multi windows don't update scope correctly
-@pytest.mark.xfail(
-    reason='Issue #2412 Same window op with different window size on table '
-    'lead to incorrect results for pyspark backend',
-    strict=True,
-)
 def test_multiple_windows(client):
     table = client.table('time_indexed_table')
     window1 = ibis.trailing_window(