ENH: Allow more flexible return type for UDFs #2776

emilyreff7 · 2021-05-12T17:33:12Z

Overview

This PR adds support for more flexible return types of single- and multi-column UDFs.

Proposed Changes

Before this PR:

Supported return types for single column UDF: pd.Series
Supported return types for multi column UDF: Tuple[pd.Series]
Supported return types for multi column aggregation UDF: tuple of scalar

After this PR:

Supported return types for single column UDF: pd.Series, list, np.ndarray
Supported return types for multi column UDF: Tuple[pd.Series], List[pd.Series], Tuple[np.ndarray], List[np.ndarray]
Supported return types for multi column aggregation UDF: tuple of scalar, list of scalar, np.ndarray of scalar

The proposed return types make the UDF API a bit more flexible for the user, and are also non-ambiguous when the return type of the UDF is not an Array type.

Tests

Added tests in test_vectorized_udf.py to exercise the above return types for elementwise, analytic, and reduction UDFs.

icexelloss · 2021-05-12T21:59:52Z

ibis/backends/tests/test_vectorized_udf.py

+def demean_struct_tuple_np_array(v, w):
+    assert isinstance(v, pd.Series)
+    assert isinstance(w, pd.Series)
+    return np.array(v - v.mean()), np.array(w - w.mean())


Is it too hard to parameterized the return shape / type?

You mean like define a single UDF and then have the test cases pass in the return type / shape (e.g. 'tuple', 'np.ndarray') and have the UDF figure out how to structure the output?

More like generating different test fixture for each return shape

demean_struct_udfs = [ create_demean_struct_udf(output_shape=np.array), create_demean_struct_udf(output_shape=tuple), create_demean_struct_udf(output_shape=pd.Series), create_demean_struct_udf(output_shape=pd.DataFrame) ]

Or sth like

demean_struct_udfs = [ create_demean_struct_udf(output_shape=lambda v1, v2: tuple(v1.values, v2.values)) # tuple of numpy array, create_demean_struct_udf(output_shape=lambda v1, v2: pd.DataFrame({'v1': v1, 'v2': v1})) # dataframe with matching names ]

After trying this out I think there are pros and cons of this approach - the pros are that the list (e.g. demean_struct_udfs) makes it very clear what is the set of UDFs and their different return types that we're testing, and would be easy to add to in the future. The con is that in order for create_demean_struct_udf to return a function that is decorated with the UDF annotation and then wrapped in the lambda that formats the result, we change the signature of the original UDF to be a regular function, which then affects other tests that reuse these UDFs - e.g. they also have to call the create function to use them.

icexelloss · 2021-05-12T22:02:05Z

ibis/udf/vectorized.py

+            result = self.func(*args, **kwargs)
+
+            if isinstance(self.output_type, dt.Struct):
+                if isinstance(result, pd.DataFrame) or (


We should probably have function "coerce_struct_output" to capture this logic here

Actually, why don't we do

if isinstance(self.output_type, dt.Struct): result = coerce_to_dataframe(result, self.output_type.names) else: # coerce non struct output to standard form

Sounds good, I can encapsulate some of this logic to make it cleaner.

Before encapsulate, maybe explain first why checking array is needed here?

Because if the intended return type of the UDF is an array, then we would mis-shape the result data here by coercing to DataFrame or Series.

For example, if result is
[np.array(1, 2, 3), np.array(4, 5, 6)] with column names 'A' and 'B'
and the output type of my UDF is Struct(Array, Array),
then what I want to return in my result is something like:

index A B 0 [1, 2, 3] [4, 5, 6]

Without this check, we would instead return

index A B 0 1 4 1 2 5 2 3 6

The same is true for single-col UDFs.

Can we pass the output type into coerce_to_dataframe to do the right thing?

Unfortunately if we check the type directly inside this function, we get a circular dependency:

ibis/__init__.py:5: in <module> import ibis.expr.types as ir ibis/expr/types.py:11: in <module> import ibis.util as util ibis/util.py:27: in <module> import ibis.expr.datatypes as dt ibis/expr/datatypes.py:146: in <module> class Null(DataType): ibis/expr/datatypes.py:147: in Null scalar = ir.NullScalar E AttributeError: module 'ibis.expr.types' has no attribute 'NullScalar'

Avoided the circular dependency by moving this function into schema.py

ibis/udf/vectorized.py

emilyreff7 · 2021-05-13T13:48:19Z

ibis/backends/tests/test_vectorized_udf.py

@@ -290,11 +397,13 @@ def foo1(v, amount):
            return v + 1


-@pytest.mark.only_on_backends(['pandas', 'pyspark', 'dask'])
+@pytest.mark.only_on_backends(['pandas', 'pyspark'])
+@pytest.mark.xfail_backends(['dask'])


Do we already have an existing issue for the dask destructure refactoring work that would be required here? @gerrymanoim

icexelloss · 2021-05-13T14:44:19Z

ibis/backends/tests/test_vectorized_udf.py

+    )(_format_struct_udf_return_type(demean_struct, result_formatter))
+
+
+demean_struct_udfs = [


I think DataFrames with matching names are also easy to support?

Perhaps also add a test for mismatching names (error out)

Based on our implementation I believe we already support this: https://github.com/ibis-project/ibis/blob/master/ibis/util.py#L119

Actually if demean_struct returns a 2-column pd.DataFrame, the test will currently pass whether or not the returned df has matching names because it just sets the names at the end.

I see so if there is no matching names, we will just match by position?

Correct, based on the current implementation, coerce_to_dataframe just cares if the result DataFrame has the correct number of columns based on output_type.names.

icexelloss · 2021-05-13T14:46:15Z

ibis/backends/tests/test_vectorized_udf.py

+    )(_format_struct_udf_return_type(mean_struct, result_formatter))
+
+
+mean_struct_udfs = [


Series can also be supported here I think:

pd.Series([v1, v2])

Updated to include a test for this as well.

Hmm maybe this isn't directly supported in the dask backend, will have to take a look.

jreback

some comments, also pls rebase

ibis/backends/tests/test_vectorized_udf.py

ibis/expr/schema.py

jreback · 2021-05-17T18:12:07Z

ibis/expr/schema.py

+            # Promote scalar to Series
+            result = pd.concat([pd.Series([v]) for v in data], axis=1)
+    else:
+        raise ValueError(f"Cannot coerce to DataFrame: {data}")


test hits this?

is this just a scalar?

For example, if we don't add np.ndarray here (which I think we might as well, as you suggest below), a test that returns np.ndarray of np.ndarray would hit this ValueError. Other than that, I'm not sure realistically what other non-0-len type outside of DataFrame, Series, list, tuple, np.ndarray could make it through here given that this function is only called on multi-column UDFs. I don't believe there are any unit tests for this function but I could add one.

ibis/expr/schema.py

ibis/udf/vectorized.py

jreback · 2021-05-17T18:15:38Z

ibis/udf/vectorized.py

+
+            if isinstance(self.output_type, dt.Struct):
+                if isinstance(result, pd.DataFrame) or (
+                    isinstance(result[0], (pd.Series, list, np.ndarray))


does a 0-len return type hit here? (it would error), tests?

Functionally this should never happen because it's guarded by

if isinstance(self.output_type, dt.Struct)

which by definition means the output should not be 0-len.

ibis/util.py

icexelloss · 2021-05-17T20:42:29Z

ibis/backends/pandas/execution/util.py

    elif isinstance(expr, (ir.DestructScalar, ir.StructScalar)):
        # Here there are two cases, if this is groupby aggregate,
        # then the result e a Series of tuple/list, or
        # if this is non grouped aggregate, then the result
-        return ibis.util.coerce_to_dataframe(result, expr.type().names)
+        return coerce_to_dataframe(
+            result, expr.type().names, expr.type().types


just passing expr.type()?

Yep can do that, I saw that dask defines a similar function so wanted to be consistent with it, but probably doesn't matter too much: https://github.com/ibis-project/ibis/blob/master/ibis/backends/dask/execution/util.py#L153

icexelloss · 2021-05-17T20:43:56Z

ibis/backends/pyspark/compiler.py

    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        result = func(*args, **kwargs)
-        return coerce_to_dataframe(result, output_cols)
+        return coerce_to_dataframe(result, output_cols, output_types)


Do we need to coerce for struct scalar?

jreback

@emilyreff7 can you also add a release note (also let's create a followon issue for updating the docs proper for all the fun things udfs can do now)

jreback

lgtm

ibis/udf/vectorized.py

icexelloss

LGTM

icexelloss · 2021-05-18T15:54:50Z

ibis/backends/pandas/udf.py

@@ -269,7 +269,7 @@ def aggregator(first, *rest):
                # because this is the inner loop and we do not want
                # to wrap a scalar value with a series.
                if isinstance(op._output_type, dt.Struct):
-                    return coerce_to_dataframe(result, op._output_type.names)


Can we remove this since coercion happens in the udf wrapper?

Same for the pyspark wrapper

Removed both

icexelloss · 2021-05-18T16:04:04Z

ibis/udf/vectorized.py

+            return coerce_to_series
+        # Case 4: Array output type, reduction UDF -> coerce to np.ndarray
+        elif isinstance(self.output_type, dt.Array):
+            return coerce_to_np_array


Do we have test to cover this?

Yes this test hits it: https://github.com/ibis-project/ibis/blob/master/ibis/backends/dask/tests/test_udf.py#L487
I can add a more general one for all backends in test_vectorized_udf.py

jreback · 2021-05-18T19:02:45Z

thanks @emilyreff7 very nice!

emilyreff7 changed the title ~~Allow more flexible return type for UDFs~~ ENH: Allow more flexible return type for UDFs May 12, 2021

icexelloss reviewed May 12, 2021

View reviewed changes

ibis/udf/vectorized.py Outdated Show resolved Hide resolved

emilyreff7 commented May 13, 2021

View reviewed changes

icexelloss reviewed May 13, 2021

View reviewed changes

jreback added feature Features or general enhancements udf Issues related to user-defined functions labels May 17, 2021

jreback suggested changes May 17, 2021

View reviewed changes

emilyreff7 force-pushed the flexible-udf-return-type branch from ba2ce93 to 2e62598 Compare May 17, 2021 20:12

emilyreff7 commented May 17, 2021

View reviewed changes

ibis/util.py Show resolved Hide resolved

icexelloss reviewed May 17, 2021

View reviewed changes

jreback added this to the Next release milestone May 17, 2021

jreback suggested changes May 17, 2021

View reviewed changes

emilyreff added 11 commits May 18, 2021 10:32

Allow more flexible return type for UDFs

5992a46

Do not coerce for Array output types

28ab0dc

Parameterize tests

ee10e40

Parameterize test functions and clean up logic

3d524d9

Fix failing tests

03d9c03

Add DataFrame tests

24f3dfb

Move coerce_to_dataframe

0754222

Fix Spark output_type

cc4b7ee

Clean up tests and code

e4662c0

Move branching logic to UDF initialization

26ae35c

Small fixes

65a10e8

emilyreff7 force-pushed the flexible-udf-return-type branch from 2e62598 to 65a10e8 Compare May 18, 2021 14:35

Fix dask tests

58764b5

jreback approved these changes May 18, 2021

View reviewed changes

icexelloss reviewed May 18, 2021

View reviewed changes

ibis/udf/vectorized.py Show resolved Hide resolved

icexelloss approved these changes May 18, 2021

View reviewed changes

icexelloss reviewed May 18, 2021

View reviewed changes

emilyreff added 2 commits May 18, 2021 12:24

Remove unneeded coerce_to_dataframe calls

e973260

Add test for reduction UDF returning array

d3abb7a

jreback merged commit b46c60b into ibis-project:master May 18, 2021

This was referenced May 18, 2021

test: dask Backend specific xfailed tests #2553

Closed

Remove unneeded check about arrays in coerce_to_dataframe #2787

Merged

gerrymanoim mentioned this pull request May 27, 2021

ENH: Add dt.Struct support back to elementwise udfs #2797

Merged

emilyreff7 deleted the flexible-udf-return-type branch June 3, 2021 18:52

emilyreff7 mentioned this pull request Aug 6, 2021

[ENH] Improve efficiency of pandas multi-col UDF aggregation processing logic #2706

Closed

cpcloud removed this from the Next release milestone Jan 7, 2022

		)(_format_struct_udf_return_type(demean_struct, result_formatter))


		demean_struct_udfs = [

		)(_format_struct_udf_return_type(mean_struct, result_formatter))


		mean_struct_udfs = [

ENH: Allow more flexible return type for UDFs #2776

ENH: Allow more flexible return type for UDFs #2776

Conversation

emilyreff7 commented May 12, 2021 • edited Loading

Overview

Proposed Changes

Tests

Choose a reason for hiding this comment

Choose a reason for hiding this comment

icexelloss May 12, 2021 • edited Loading

Choose a reason for hiding this comment

icexelloss May 12, 2021 • 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

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

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

jreback left a 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

jreback commented May 18, 2021

emilyreff7 commented May 12, 2021 •

edited

Loading

icexelloss May 12, 2021 •

edited

Loading

icexelloss May 12, 2021 •

edited

Loading