ENH: Support multi args window UDF for pandas backend #2035

icexelloss · 2019-11-19T15:33:28Z

This PR addresses issue #1998

Currently, when using UDAF with rolling window, pandas backend will throw an exception:

import ibis
import pandas as pd
import numpy as np

from ibis.pandas.udf import udf
import ibis.expr.datatypes as dt

client = ibis.pandas.connect({'table': pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'key': ['a', 'a', 'a']})})
t = client.table('table')
w = ibis.trailing_window(preceding=1, order_by='key', group_by='key')
#w = ibis.window(group_by='key')


@udf.reduction(input_type=[dt.double, dt.double], output_type=dt.double)
def my_average(v, w):
    return np.average(v, weights=w)

t = t.mutate(new_col=my_average(t.a, t.b).over(w))

t.execute()

With this PR, it will support this use case.

The main idea of this PR is that instead of using groupby().rolling().apply(func) to compute the result, we use groupby().rolling().apply(len, raw=True) to get the size of each window, and then manually apply func to each window in a Python for-loop. This way, we work around issue that groupby().rolling().apply(func) can only take function that apply on a single series.

Benchmarks

Before the change:

In [4]: %time pandas.time_low_card_window_analytics_udf()                                                                                                     
CPU times: user 898 ms, sys: 248 ms, total: 1.15 s
Wall time: 1.15 s

In [5]: %time pandas.time_high_card_window_analytics_udf()                                                                                                    
CPU times: user 12.4 s, sys: 324 ms, total: 12.7 s
Wall time: 12.8 s

In [6]: %time pandas.time_low_card_grouped_rolling()                                                                                                          
CPU times: user 7.99 s, sys: 1.95 s, total: 9.94 s
Wall time: 10.1 s

In [7]: %time pandas.time_high_card_grouped_rolling()                                                                                                         
CPU times: user 15.3 s, sys: 1.83 s, total: 17.1 s
Wall time: 17.3 s

In [8]: %time pandas.time_low_card_grouped_rolling_udf()                                                                                                      
CPU times: user 2min 1s, sys: 1.96 s, total: 2min 3s
Wall time: 2min 4s

In [9]: %time pandas.time_high_card_grouped_rolling_udf()                                                                                                     
CPU times: user 1min 41s, sys: 2.17 s, total: 1min 43s
Wall time: 1min 43s

After the change:

time pandas.time_low_card_window_analytics_udf()
CPU times: user 856 ms, sys: 235 ms, total: 1.09 s
Wall time: 1.08 s
time pandas.time_high_card_window_analytics_udf()
CPU times: user 12.7 s, sys: 316 ms, total: 13 s
Wall time: 13.1 s
time pandas.time_low_card_grouped_rolling()
CPU times: user 9.3 s, sys: 3.15 s, total: 12.4 s
Wall time: 13 s
time pandas.time_high_card_grouped_rolling()
CPU times: user 15.5 s, sys: 1.9 s, total: 17.4 s
Wall time: 17.5 s
time pandas.time_low_card_grouped_rolling_udf()
CPU times: user 1min 26s, sys: 1.74 s, total: 1min 27s
Wall time: 1min 28s
time pandas.time_high_card_grouped_rolling_udf()
CPU times: user 1min 4s, sys: 1.79 s, total: 1min 6s
Wall time: 1min 6s

icexelloss · 2019-11-19T15:35:55Z

cc @jreback @toryhaavik

icexelloss · 2019-11-19T15:36:35Z

ibis/pandas/aggcontext.py

-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore", message=".+raw=True.+", category=FutureWarning
+            # get the DataFrame from which the operand originated


These are the original logic for built-in aggregation functions.

icexelloss · 2019-11-19T15:36:53Z

ibis/pandas/tests/test_udf.py

+def test_udaf_window_multi_params():
+    @udf.reduction(['double', 'double'], 'double')
+    def my_wm(v, w):
+        print("v")


Remove these

icexelloss · 2019-11-19T15:37:04Z

ibis/pandas/udf.py

@@ -530,25 +531,33 @@ def execute_udaf_node_groupby(op, *args, **kwargs):
                #
                # If the argument is not a SeriesGroupBy then keep
                # repeating it until all groups are exhausted.
+


Remove white space

icexelloss · 2019-11-19T15:37:16Z

ibis/pandas/aggcontext.py

@@ -361,6 +364,7 @@ def agg(self, grouped_data, function, *args, **kwargs):
        else:
            # do mostly the same thing as if we did NOT have a grouping key,
            # but don't call the callable just yet. See below where we call it.
+


Remove white space

icexelloss · 2019-11-19T21:43:22Z

ibis/pandas/execution/tests/test_arrays.py

@@ -52,15 +52,8 @@ def test_array_collect(t, df):
    tm.assert_frame_equal(result, expected)


-@pytest.mark.xfail(


This is now supported as well

jreback

some comments

jreback · 2019-11-22T01:46:22Z

ibis/pandas/aggcontext.py

 import pandas as pd
+from pandas import Series


prob can just use of.Series

jreback · 2019-11-22T01:54:50Z

ibis/pandas/aggcontext.py

+                # create a generator for each input series
+                # the generator will yield a slice of the
+                # input series for each valid window
+                data = getattr(grouped_series, 'obj', grouped_series).values


you shouldn’t use .values here a that coerces to a ndarray

rather leave as a Series and use .iloc

I choose to do ndarray here because:

This preserve the same udf API as before.

Passing Series is about 15x slower than ndarray

I don't want to introduce API change and performance regression in this PR. I think we can have a separate chat whether window UDF should take ndarray or Series.

you shouldn’t show any perf issues must be something odd going on

does it currently take ndarray?

Yes it currently take ndarray (By using raw=True I think)

ok can you create an issue to fix this, meaning to use iloc. .values is not type preserving and generally a bad idea.

Created: #2045

jreback · 2019-11-22T01:55:30Z

ibis/pandas/aggcontext.py

+            inputs = args if len(args) > 0 else [grouped_data]
+
+            input_gens = list(
+                create_input_gen(arg, window_size)


what are you trying to do here?

Here I am creating generators for each inputs so later I can just call next(input). This hides the details of how next is implemented and unifies how we send data inputs and arg inputs to the user function.

jreback · 2019-11-22T01:56:12Z

ibis/pandas/aggcontext.py

+            result[mask] = valid_result
+            result.index = obj.index
+        else:
+            with warnings.catch_warnings():


i would move these out to 2 module level functions rather than nesting like this

Actually this is not needed anymore because we don't use raw=True - it was used for UDF and now UDF is handled separately.

ibis/pandas/execution/tests/test_arrays.py

jreback · 2019-11-22T01:57:17Z

ibis/pandas/tests/test_udf.py

@@ -271,6 +272,87 @@ def my_mean(series):
    tm.assert_frame_equal(result, expected)


+def test_udaf_window_interval():
+    @udf.reduction(['double'], 'double')


isn’t this defined above?

jreback · 2019-11-22T01:59:18Z

ibis/pandas/udf.py

+                    # the custom rolling logic.
+                    result = aggcontext.agg(args[0], func, *args, **kwargs)
+                else:
+                    iters = (


i would move this to a module level function

ibis/pandas/udf.py

jreback · 2019-11-27T16:34:17Z

ibis/pandas/aggcontext.py

+                # create a generator for each input series
+                # the generator will yield a slice of the
+                # input series for each valid window
+                data = getattr(grouped_series, 'obj', grouped_series).values


ok can you create an issue to fix this, meaning to use iloc. .values is not type preserving and generally a bad idea.

jreback · 2019-11-27T16:34:45Z

ibis/pandas/aggcontext.py

+            raw_window_size = windowed.apply(len, raw=True).reset_index(
+                drop=True
+            )
+            mask = ~(raw_window_size.isna())


pls make this a separate function, this is too hard to grok inline here.

jreback · 2019-11-27T16:35:37Z

ibis/pandas/tests/test_udf.py

@@ -271,6 +274,88 @@ def my_mean(series):
    tm.assert_frame_equal(result, expected)


+def test_udaf_window_interval():


are all of the window udf tests here or in test_window?

all winndow udf tests are here

jreback · 2019-11-27T16:35:54Z

ibis/pandas/udf.py

@@ -100,6 +101,29 @@ def arguments_from_signature(signature, *args, **kwargs):
    return args, new_kwargs


+def create_gens_from_args_groupby(args):


can you type args

jreback · 2019-11-27T16:36:22Z

ibis/pandas/udf.py

+    """ Create generators for each args for groupby udaf.
+
+    If the arg is SeriesGroupBy, return a generator that outputs each group.
+    If the arg is not SeriesGroupBy, return a generator that repeats the arg.


what else could this be? can you type it

Improved docstring

ibis/pandas/udf.py

jreback · 2019-11-27T16:38:23Z

ibis/pandas/aggcontext.py

            )
+
+            valid_result = pd.Series(valid_result)
+            valid_result.index = window_size.index


are you testing the output indexes are correct?

yes there are tests (test_udaf_window_interval) that cover out of order indices and make sure the output is correct

jreback · 2019-11-27T16:38:40Z

ibis/pandas/aggcontext.py

+
+            valid_result = pd.Series(valid_result)
+            valid_result.index = window_size.index
+            result = pd.Series(np.repeat(None, len(obj)))


this is really strange to do, what are you trying?

jreback · 2019-11-27T16:39:11Z

ibis/pandas/aggcontext.py

+            result = pd.Series(np.repeat(None, len(obj)))
+            result[mask] = valid_result
+            result.index = obj.index
+        else:
            result = method(windowed)


can you add some commments on what is this case (again would be much better to split these out to free functions)

I have separated this into separate functions

jreback · 2019-12-03T21:57:12Z

if you'd rebase can look again

icexelloss · 2019-12-04T15:19:48Z

@jreback Haven't addressed all comments. Will ping again when it's done.

jreback · 2019-12-04T20:17:49Z

ibis/pandas/aggcontext.py

@@ -326,6 +327,83 @@ def __init__(self, kind, *args, **kwargs):
        )
        self.construct_window = operator.methodcaller(kind, *args, **kwargs)

+    def _agg_built_in(self, frame, windowed, function, *args, **kwargs):


I would make both of these module level functions (don't pass self, I think you can just pass max_lookback as as kord arg).

I feel that these methods are very particular to Window aggregation context and therefore probably belongs in the class. I curious why module level functions are better here?

Discussed offline. Moved to module level.

jreback

looks good. can you add a release note (likely point to a new doc section) & a doc section (can be a followup PR).

jreback · 2019-12-06T20:17:33Z

ibis/pandas/aggcontext.py

@@ -313,6 +314,92 @@ def compute_window_spec_interval(_, expr):
    return pd.tseries.frequencies.to_offset(value)


+def _window_agg_built_in(
+    frame, windowed, function, max_lookback, *args, **kwargs


ideally if you can type these arguments

ibis/pandas/aggcontext.py

jreback · 2019-12-06T20:19:07Z

ibis/pandas/aggcontext.py

+
+
+def _window_agg_udf(
+    grouped_data, windowed, function, dtype, max_lookback, *args, **kwargs


can you type this

jreback · 2019-12-06T20:19:27Z

ibis/pandas/aggcontext.py

+):
+    """Apply window aggregation with UDFs.
+    """
+    # Use custom logic to computing rolling window UDF instead of


include this in the doc-string under Notes

jreback · 2019-12-06T20:31:14Z

ibis/pandas/aggcontext.py

+    # This is because pandas's rolling function doesn't support
+    # multi param UDFs.
+
+    def create_input_gen(grouped_series, window_size):


can you type

jreback · 2019-12-06T20:46:58Z

ibis/pandas/aggcontext.py

+            obj = getattr(grouped_data, 'obj', grouped_data)
+            name = obj.name
+            if frame[name] is not obj:
+                name = "{}_{}".format(name, ibis.util.guid())


can use an fstring here

jreback · 2019-12-06T20:47:16Z

ibis/pandas/aggcontext.py

+            # TODO: see if we can do this in the caller, when the context
+            # is constructed rather than pulling out the data
+            columns = group_by + order_by + [name]
+            indexed_by_ordering = frame.loc[:, columns].set_index(order_by)


use frame[columns].set_index(...)

jreback · 2019-12-06T20:48:04Z

ibis/pandas/tests/test_udf.py

+
+
+def test_udaf_window_multi_args():
+    @udf.reduction(['double', 'double'], 'double')


can you add the issue number as a comment (or this PR number)

jreback · 2019-12-09T23:03:49Z

ibis/pandas/aggcontext.py

+            name = obj.name
+            if frame[name] is not obj:
+                name = f"{name}_{ibis.util.guid()}"
+                frame[name] = obj


do this like

frame = frame.assign(name=obj)

to avoid mutating the input

I see you are just copying original code (but this should change anyhow)

ibis/pandas/aggcontext.py

jreback · 2019-12-09T23:09:10Z

ibis/pandas/tests/test_udf.py

        }
    )
    con = ibis.pandas.connect({'df': df})
    t = con.table('df')
    window = ibis.trailing_window(2, order_by='a', group_by='key')
-    expr = t.mutate(rolled=my_mean(t.b).over(window))
+    expr = t.mutate(
+        wm_b=my_wm(t.b, t.d).over(window), wm_c=my_wm(t.c, t.d).over(window)


I assume that this will work if we use different windows? can you add a test for that, or if it doesn't work can you test the exception.

jreback · 2019-12-09T23:10:24Z

ibis/pandas/tests/test_udf.py

@@ -225,49 +250,107 @@ def my_corr2(lhs, **kwargs):
            pass


-def test_compose_udfs():
+def test_compose_udfs(t2, df2):


what happens if the udaf raises an exception? are these caught in a reasonable way? just asking for a test (not even sure what this should do)

Added test_udf_error

jreback

lgtm. minor comment & can you add a note in release.rst? ping on green.

jreback · 2019-12-10T21:06:22Z

ibis/pandas/aggcontext.py

+                    **kwargs,
+                )
+            try:
+                return result.astype(self.dtype, copy=False)


may want to add a comment on when this can fail

Fix tests and address comments

jreback · 2019-12-10T22:59:26Z

thanks @icexelloss very nice patch!

icexelloss commented Nov 19, 2019

View reviewed changes

jreback suggested changes Nov 22, 2019

View reviewed changes

icexelloss force-pushed the pandas-backend-multi-args-udf branch 2 times, most recently from 700a22a to ee64a17 Compare November 25, 2019 19:13

jreback added the feature Features or general enhancements label Nov 27, 2019

jreback added this to the Next Feature Release milestone Nov 27, 2019

jreback suggested changes Nov 27, 2019

View reviewed changes

icexelloss force-pushed the pandas-backend-multi-args-udf branch from cccfa5c to de81e5a Compare December 3, 2019 22:16

icexelloss changed the title ~~ENH: Support multi param window UDF for pandas backend~~ ENH: Support multi arg window UDF for pandas backend Dec 4, 2019

icexelloss changed the title ~~ENH: Support multi arg window UDF for pandas backend~~ ENH: Support multi args window UDF for pandas backend Dec 4, 2019

jreback reviewed Dec 4, 2019

View reviewed changes

icexelloss added the pandas The pandas backend label Dec 4, 2019

jreback suggested changes Dec 6, 2019

View reviewed changes

jreback suggested changes Dec 9, 2019

View reviewed changes

jreback approved these changes Dec 10, 2019

View reviewed changes

icexelloss added 8 commits December 10, 2019 17:02

Support multi params window UDF for pandas backend

40efe93

Fix array tests

3e1d2dd

Add benchmarks

82e3756

Fix rolling with interval

640b7a3

Fix perf

fcffa45

Address comments and fix tests

a38f098

Fix tests and address comments

Minor change

098cb19

Address comments

22b6237

icexelloss added 8 commits December 10, 2019 17:02

Move _agg_udf and _agg_built_to to module level

0e81f85

Address comments

23454be

Add quick test method in Makefile

75c7ba6

Fix tests

6a25c3f

Fix tests

41ef8b7

Address comments

71a8278

Fix tests

5fef9a2

Add release notes

e652fb7

icexelloss force-pushed the pandas-backend-multi-args-udf branch from 7fb3065 to e652fb7 Compare December 10, 2019 22:03

jreback merged commit a95e32f into ibis-project:master Dec 10, 2019

		@@ -52,15 +52,8 @@ def test_array_collect(t, df):
		tm.assert_frame_equal(result, expected)


		@pytest.mark.xfail(

		@@ -271,6 +274,88 @@ def my_mean(series):
		tm.assert_frame_equal(result, expected)


		def test_udaf_window_interval():

		@@ -100,6 +101,29 @@ def arguments_from_signature(signature, args, *kwargs):
		return args, new_kwargs


		def create_gens_from_args_groupby(args):



		def _window_agg_udf(
		grouped_data, windowed, function, dtype, max_lookback, args, *kwargs



		def test_udaf_window_multi_args():
		@udf.reduction(['double', 'double'], 'double')

ENH: Support multi args window UDF for pandas backend #2035

ENH: Support multi args window UDF for pandas backend #2035

Conversation

icexelloss commented Nov 19, 2019 • edited Loading

Benchmarks

icexelloss commented Nov 19, 2019

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

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

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

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 commented Dec 3, 2019

icexelloss commented Dec 4, 2019

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

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

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

jreback commented Dec 10, 2019

icexelloss commented Nov 19, 2019 •

edited

Loading