BUG: Fix same window op with different window size on table lead to incorrect results for pyspark backend

[LeeTZ]

Overview

This PR fixes #2412 where the same window op with different window size on tables will lead to incorrect results for pyspark backend.
A refactor of param window in pyspark translation is proposed.

Example

Let's say we create two windows with different sizes to calculate Mean() on the same column of a table

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)

Currently, this will fail. The compiled result in result_pd is wrong and the two columns of different window sizes are the same.

The Problem

The reason why these two columns are the same is: the value we stored in scope is wrong. The second window is not compiled.

Here we are calculating Mean() on 2 different windows, 1h and 2h. Note that in compile_window_op, we call

result = t.translate(operand, scope, timecontext, window=pyspark_window, context=context)
return result

passing window as a param and calculate window in the dispatched compile process. e.g. compile_aggregator

This is actually wrong. since we pass window as a param to compile_aggregator, then Mean() will be compiled to Col<... over 1h window> . When we translate the second window (2h). translate() will try to look up in scope first for Mean(). And Mean() is already compiled to Col<... over 1h window>. There is no window information in the key in scope for Mean(). So the 2h window will not be translated again and we get an incorrect result of 2 windows with the same size.

The Change

To make our scope cache correct and reliable, we should lift the logic to run .over(window) in compile_window_op.
This is a little complicated here:

There are operations like Lead, Lag, NotAny, NotAll, Rank, etc. These operations take window as a param, and they run .over(window) inside the dispatch function to compile these ops. If we move .over(window) outside of these ops, there will be several issues:

1. for compile_rank and compile_dense_rank, they do:

F.dense_rank().over(window).astype('long') - 1

So if we move .over(window) to compile_window_op, we will have to do some post processing to do typecast.

2. for compile_notany and compile_notall, they do:

if context is None:
        def fn(col):
            return ~(F.max(col))

        return compile_aggregator(
            t, expr, scope, timecontext, fn=fn, context=context, **kwargs
        )
    else:
        return ~compile_any(
            t,
            expr,
            scope,
            timecontext,
            context=context,
            window=window,
            **kwargs,
        )

The negation ~ here is very tricky to lift out here. We need to call .over(window) and then negate it. The solution we propose here is a rewrite in pre-translation to pass Any and All op to translation, and negate them back after window operation.

How is this tested

For window operations, cases are covered in ibis/tests/all/test_window.py. For multi-window, tests are added in ibis/pyspark/tests/test_window.py (Actually it's a re-enable of previously marked failure test)

[icexelloss]

Let's have cleaner branching like this:

if isinstance(op, (ops.NotAll, ops.NotAny)):
    return ...
elif isinstance(op, (ops.MinRank,ops.DenseRank, ops.RowNumber)):
    return ...
else:
    return ...

[LeeTZ]

That's indeed more readable. However, I argue that this may bring duplication in code:

result = t.translate(operand, scope, timecontext, context=context).over(
        pyspark_window
    )

will repeat 3 times.
The motivation for the current structure is pre_processing, translation, post_processing.

[LeeTZ]

I revised as you suggested for more readability

[icexelloss]

Looks good in general minor comments.

[jreback]

factor out this to a variable
t.translate(operand, scope, timecontext, context=context)

[jreback]

lgtm ping on green

[LeeTZ]

@jreback Green now

[icexelloss]

Reviewed another around

[icexelloss]

LGTM

[jreback]

thanks @LeeTZ