FEAT: Support Ibis interval for window in pyspark backend #2409

LeeTZ · 2020-09-26T23:27:51Z

Overview

Currently, window size defined by ibis.interval is broken for pyspark backend. This PR aims to fix this issue.

The Problem

In pyspark backend, for a table generated by the following df:

df_time_indexed = client._session.createDataFrame(
        [
            [datetime(2017, 1, 2, 5, tzinfo=timezone.utc), 1, 1.0],
            [datetime(2017, 1, 2, 6, tzinfo=timezone.utc), 1, 3.0],
        ],
        ['time', 'key', 'value'],
    )

And calculate mean with a trailing_window, ordered by time column on this table:

window = ibis.trailing_window(
    preceding=ibis.interval(hours=2), order_by='time', group_by='key'
)
result = table.mutate(
    mean_2h=table['value'].mean().over(window)
).compile()

will result in error:

-> pyspark_window = pyspark_window.rowsBetween(start, end)  (compile_window_op inibis/pyspark/compiler.py)
TypeError: Arguments with datatype interval<int8>(unit='h') and int64 are not comparable

The Fix

In compiling window op for pyspark backend, we cast TimeStampColumn to long type (seconds since epoch), this is required for spark to filtering and comparing time.
Add support to pass ibis interval in preceding and following, now None, IntervalScalar and int are supported for preceding and following in pyspark window.

Tests

Tests are added in ibis/pyspark/tests/test_window.py.

Follow-ups

Once this is fixed, the next step is to add logic for context adjustment in pyspark backend. Again, to make things easier to review, will do context adjustment in follow-up PRs.

icexelloss · 2020-09-28T14:48:42Z

ibis/pyspark/compiler.py

    ordering_keys = [
-        key.to_expr().get_name()
+        F.col(key.to_expr().get_name()).cast('long')


This looks really weird. Why are we getting op from expr and then use to_expr to get expr from op again?

ordering_keys = [ F.col(expr.get_name()).cast('long') if isinstance(expr, types.TimestampColumn) else expr.get_name() for expr in order_by ]

Does this work?

This doesn't work but I get your point that we don't need to call op and to_expr at the same time. I refactored a little.

ibis/pyspark/compiler.py

icexelloss · 2020-09-28T14:52:09Z

ibis/pyspark/compiler.py

+        if window.following is None:
+            end = Window.unboundedFollowing
+        elif isinstance(window.following, ir.IntervalScalar):
+            end = int(execute(window.following).value / 1e9)


These are duplicate logic as window.preceding. Can you refactor?

Sure, I added a function to canonicalize these intervals.

icexelloss · 2020-09-28T14:52:54Z

ibis/pyspark/tests/test_timecontext.py

+
+def test_time_indexed_window(client):
+    table = client.table('time_indexed_table')
+    window1 = ibis.trailing_window(


Can you add tests for future window as well

Sure, I add one to test following and compare the result with df generated by pyspark window.

icexelloss

High level looks good. Left some comments.

LeeTZ · 2020-09-28T18:35:09Z

ibis/pyspark/compiler.py

    )
+
+    # These ops need type cast after running over window


As we agreed over.(window) should be lifted to this function instead of passing via kwargs.
This for sure works for compile_aggregator, but there are some other ops that is dispatched in this line. For example, lead, lag, row_numbers, first, last, etc. We used to pass window as a param to them.
I removed window param in these functions, but there are some ops like the three ones I listed here, they do type cast after running over window. Let me know if this change makes sense to you, or maybe there are better ways to do so? @icexelloss

I see. This is probably why over(window) was not in compile_window_op originally.

I don't think what you have here is much better than original implementation actually and I think we should roll back to the original implementation.

I agree branching too much isn't a good idea, I reverted and commented below for thoughts

LeeTZ · 2020-09-28T18:56:56Z

ibis/pyspark/compiler.py

-def compile_notany(
-    t, expr, scope, timecontext, *, context=None, window=None, **kwargs
-):
+def compile_notany(t, expr, scope, timecontext, *, context=None, **kwargs):


https://github.com/ibis-project/ibis/pull/2409/checks?check_run_id=1178201082#step:4:1284
We removed window here and seems tests are failing for this, and compile_notall, given the comments here, could you give some suggestions on how we should refactor? @icexelloss

This function looks fine and I don't think we need to change this. What exactly failed and how?

The stacktrace is:

pyspark.sql.utils.AnalysisException: "grouping expressions sequence is empty, and 'functional_alltypes.`index`' is not an aggregate function. Wrap '(min((functional_alltypes.`double_col` = CAST(0 AS DOUBLE))) AS `_w0`)' in windowing function(s) or wrap 'functional_alltypes.`index`'

I debugged into this issue
For col_in_selection_order in compile_selection_op`, The correct value should be:

['index', 'Unnamed: 0', 'id', 'bool_col', 'tinyint_col', 'smallint_col', 'int_col', 'bigint_col', 'float_col', 'double_col', 'date_string_col', 'string_col', 'timestamp_col', 'year', 'month', Column<b'(NOT min((double_col = 0)) OVER (PARTITION BY string_col ORDER BY id ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)) AS `val`'>]

while the wrong value is

['index', 'Unnamed: 0', 'id', 'bool_col', 'tinyint_col', 'smallint_col', 'int_col', 'bigint_col', 'float_col', 'double_col', 'date_string_col', 'string_col', 'timestamp_col', 'year', 'month', Column<b'(NOT min((double_col = 0))) OVER (PARTITION BY string_col ORDER BY id ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS `val`'>]

Basically we need to call .over(window) inside F.col calls, and note that there are negate operation that we should also take care of. The fix we purpose to lift .over(window) to compile_window_op breaks this behavior.

As per discussion with @icexelloss , NotAny must be compiled to ~compile(Any) otherwse table aggregation breaks. Also we have to keep NotAny to call compile_aggregator because it could be used with table aggregation. e.g.

table.groupby('key').aggregate(new_col=table['bool'].not_any())

So we cannot easily change the rule for compiling NotAny to replement the order to insert over.window() inside.
Let me just revert this change to remove window and address this in followups

ibis/pyspark/compiler.py

LeeTZ · 2020-09-29T19:32:43Z

Overview

I revert refactoring to remove window from param, and add a flag in translate to force translation without using cached result in scope.

The Problem

The problem of window passed in param is: same op with different window size will not be compiled twice.
For example, the test case in test_timecontext.py, we are calculating Mean() on 2 different windows, 1h and 2h.
Since we pass window as a param to compile_aggregator, then Mean() will be compiled, and the result will be 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>. So this 2h window will not be translated again and we get an incorrect result of 2 window with the same range.

A perfect fix will be removing window from param and lift the logic to run .over(window) in compile_window_op. But 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:

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 branch on type of op and do such typecast. Too much branching and multiple calls of instance is not a good implementation.
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. Also according to the stacktrace of failing tests in a previous commit in this PR,
Column<b'(NOT min((double_col = 0)) OVER (PARTITION BY string_col ORDER BY id ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)) AS val'>] is what we should get.
But if we lift .over(window) outside of compile_notany, column will be compiled to Column<b'(NOT min((double_col = 0))) OVER (PARTITION BY string_col ORDER BY id ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS val'>]
Note the differences in the location of (). We need to call .over(window) inside F.col calls, and note that there is a negate operation that we should also take care of. The fix we purpose to lift .over(window) to compile_window_op breaks this behavior.
3. A solution might be: don't call negate on compile_any for compile_notany, and don't call compile_aggregator, implement the full logic to handle compile with window in compile_notany, and other functions that are causing issues. This is not a good solution, not only it brings more duplication in code, but also breaks the behavior. we have to keep NotAny to call compile_aggregator because it could be used with table aggregation. e.g.

table.groupby('key').aggregate(new_col=table['bool'].not_any())

So we cannot easily change the rule for compiling NotAny to reimplement the order to insert over.window() inside.

The Change

Given the complexity of removing window in param, I proposed to add a flag in translate to force translation to not trust result cached in scope, so that the same op with different windows will be translated twice with different window ranges, and we could get the correct result. This sacrifices performance a little bit but guarantee the correctness. This flag only applies to translate in compile_window_op now.

Follow-ups

We should definitely think about how we could refactor all window related op compilation. Will create an issue to track this.

…:LeeTZ/ibis into tiezheng/fix_pyspark_time_indexed_window

LeeTZ · 2020-09-29T21:26:14Z

As per offline discussion with Li, I am leaving this PR simple as to support Ibis interval for pyspark window, let me create another PR to do the refactor.

icexelloss · 2020-09-29T21:26:46Z

Given the complexity of removing window in param, I proposed to add a flag in translate to force translation to not trust result cached in scope

Sorry, I am -1 to this idea. I think if we cannot trust the cache then the fundamental invariant of the compile algorithm breaks. (an op can uniquely define the compilation result)

Per discussion, we agree the best solution forward is to change the lift the pyspark .over(window) statement from the compilation result for aggregators (e.g., Min, Max, NotAny, etc) to the compilation of window op itself.

This increases the complexity in the "compile window op" but at least the complexity is local to that function. "local complexity" is better than "global complexity" which the flag solution will introduce.

icexelloss · 2020-09-29T21:33:11Z

ibis/pyspark/compiler.py

+    """
+    if isinstance(interval, ir.IntervalScalar):
+        # value is in nanosecond
+        return int(execute(interval).value / 1e9)


Do we have a translation rule for IntervalScalar? Perhaps we can add a rule for IntervalScalar and do

@compiles(IntervalScalar) def compile_interval_scalar(interval): return execute(interval) timedelta = t.translate(interval, ...) # This is a pd.Timedelta second_since_epoch = int(timedelta.value / 1e9)

IntervalScalar is not of type op (operations), while all other @ compiles are compiling ops. I found it a little strange to do so.

op is Literal

hmm i agree with @icexelloss here shouldn't we just be able to translate?

or should this be a common method with other backends?

I add this to compile(Literal)

ibis/pyspark/compiler.py

icexelloss · 2020-09-29T21:38:17Z

ibis/pyspark/tests/test_window.py

+pytestmark = pytest.mark.pyspark
+
+
+def test_time_indexed_window(client):


Can you parameterize this test by the window bounds?

(trailing_window(ibis.interval(hours=1), ...), Window.partitionBy(...).rangeBetween(...))

and

(ibis.range_window(..., ...), Window.partitionBy(...).rangeBetween(...))

We cannot simply put Window in param because that might trigger error of 'NoneType' object has no attribute '_jvm' where spark jvm is not initialized properly. Instead, I put range as param.

ibis/pyspark/compiler.py

ibis/pyspark/tests/test_window.py

icexelloss

Reviewed again.

jreback · 2020-09-30T02:31:27Z

ibis/pyspark/compiler.py

+
+    if isinstance(interval, ir.IntervalScalar):
+
+        return int(execute(interval).value / 1e9)


is there a concern if this is too big here? e.g. i think resolution on spark is only ms?

Here in window bounds spark uses seconds since epoch, execute gives us a Timedelta and value is in nanosecond

jreback · 2020-09-30T02:31:39Z

ibis/pyspark/compiler.py

+        return int(execute(interval).value / 1e9)
+    elif isinstance(interval, int):
+        return interval
+    else:


Sure I removed the else here

jreback · 2020-09-30T02:32:25Z

ibis/pyspark/compiler.py

+    """
+    if isinstance(interval, ir.IntervalScalar):
+        # value is in nanosecond
+        return int(execute(interval).value / 1e9)


hmm i agree with @icexelloss here shouldn't we just be able to translate?

jreback · 2020-09-30T02:32:32Z

ibis/pyspark/compiler.py

+        return interval
+    else:
+        raise com.UnsupportedOperationError(
+            'type {} is not supported in preceding/following in window'.format(


use f-strings

jreback · 2020-09-30T02:32:54Z

ibis/pyspark/compiler.py

+    """
+    if isinstance(interval, ir.IntervalScalar):
+        # value is in nanosecond
+        return int(execute(interval).value / 1e9)


or should this be a common method with other backends?

jreback · 2020-09-30T02:34:44Z

ibis/pyspark/tests/test_window.py

+
+
+# TODO: multi windows don't update scope correctly
+@pytest.mark.xfail


can you add this int he reason for the xfail and create an issue for this (and list it here)

np, this is just the next issue I am going to fix, #2412

jreback

lgtm minor comment in tests & can you add a release note with this PR number. ping on green.

jreback · 2020-09-30T15:57:59Z

ibis/pyspark/tests/test_window.py

+
+
+# TODO: multi windows don't update scope correctly
+@pytest.mark.xfail(reason='Issue #2412', strict=True)


can you add a description as well here

LeeTZ · 2020-09-30T17:08:54Z

@icexelloss @jreback CI is green for tests now but failed for impala clickhouse / docbuild. Could we ignore and proceed or should I run again?

jreback · 2020-09-30T17:17:38Z

yep this is fine

icexelloss

LGTM

Support time indexed windowing in pyspark backend

7a73fa1

LeeTZ force-pushed the tiezheng/fix_pyspark_time_indexed_window branch from 322f13d to 7a73fa1 Compare September 28, 2020 14:13

put pyspark_window back in param

95d30e9

icexelloss reviewed Sep 28, 2020

View reviewed changes

ibis/pyspark/compiler.py Outdated Show resolved Hide resolved

icexelloss reviewed Sep 28, 2020

View reviewed changes

icexelloss suggested changes Sep 28, 2020

View reviewed changes

refactor compile_window, add tests for forward looking window

171637e

LeeTZ commented Sep 28, 2020

View reviewed changes

icexelloss reviewed Sep 28, 2020

View reviewed changes

ibis/pyspark/compiler.py Show resolved Hide resolved

add a flag in translate to use_scope

41e2a39

LeeTZ force-pushed the tiezheng/fix_pyspark_time_indexed_window branch from 171637e to 41e2a39 Compare September 29, 2020 19:03

LeeTZ closed this Sep 29, 2020

LeeTZ reopened this Sep 29, 2020

icexelloss changed the title ~~Support time indexed windowing in pyspark backend~~ FEAT: Support time based windowing in pyspark backend Sep 29, 2020

LeeTZ changed the title ~~FEAT: Support time based windowing in pyspark backend~~ FEAT: Support Ibis interval for window in pyspark backend Sep 29, 2020

LeeTZ added 3 commits September 29, 2020 16:09

Merge branch 'tiezheng/fix_pyspark_time_indexed_window' of github.com…

6f3c3d8

…:LeeTZ/ibis into tiezheng/fix_pyspark_time_indexed_window

merge change to revert window in param

2620cbf

remove tests for multi windows

4647c42

icexelloss reviewed Sep 29, 2020

View reviewed changes

ibis/pyspark/compiler.py Outdated Show resolved Hide resolved

icexelloss reviewed Sep 29, 2020

View reviewed changes

ibis/pyspark/compiler.py Outdated Show resolved Hide resolved

icexelloss reviewed Sep 29, 2020

View reviewed changes

ibis/pyspark/compiler.py Outdated Show resolved Hide resolved

icexelloss reviewed Sep 29, 2020

View reviewed changes

ibis/pyspark/compiler.py Outdated Show resolved Hide resolved

icexelloss reviewed Sep 29, 2020

View reviewed changes

ibis/pyspark/tests/test_window.py Outdated Show resolved Hide resolved

icexelloss suggested changes Sep 29, 2020

View reviewed changes

parametize tests

46f400a

jreback added pyspark The Apache PySpark backend backends - spark labels Sep 30, 2020

jreback suggested changes Sep 30, 2020

View reviewed changes

add translation for ibis interval

7cee491

jreback suggested changes Sep 30, 2020

View reviewed changes

add release note

3820da9

jreback added this to the Next Feature Release milestone Sep 30, 2020

jreback approved these changes Sep 30, 2020

View reviewed changes

jreback merged commit 67bca4e into ibis-project:master Sep 30, 2020

icexelloss reviewed Sep 30, 2020

View reviewed changes

		pytestmark = pytest.mark.pyspark


		def test_time_indexed_window(client):


		if isinstance(interval, ir.IntervalScalar):

		return int(execute(interval).value / 1e9)



		# TODO: multi windows don't update scope correctly
		@pytest.mark.xfail



		# TODO: multi windows don't update scope correctly
		@pytest.mark.xfail(reason='Issue #2412', strict=True)

FEAT: Support Ibis interval for window in pyspark backend #2409

FEAT: Support Ibis interval for window in pyspark backend #2409

Conversation

LeeTZ commented Sep 26, 2020 • edited Loading

Overview

The Problem

The Fix

Tests

Follow-ups

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

icexelloss left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LeeTZ commented Sep 29, 2020

Overview

The Problem

The Change

Follow-ups

LeeTZ commented Sep 29, 2020

icexelloss commented Sep 29, 2020 • edited Loading

icexelloss Sep 29, 2020 • edited Loading

Choose a reason for hiding this comment

LeeTZ Sep 29, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

icexelloss Sep 29, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

icexelloss left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jreback left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LeeTZ commented Sep 30, 2020

jreback commented Sep 30, 2020

icexelloss left a comment

Choose a reason for hiding this comment

LeeTZ commented Sep 26, 2020 •

edited

Loading

icexelloss commented Sep 29, 2020 •

edited

Loading

icexelloss Sep 29, 2020 •

edited

Loading

LeeTZ Sep 29, 2020 •

edited

Loading

icexelloss Sep 29, 2020 •

edited

Loading