UDF using DataFrame.apply

[AbdealiLoKo]

These are the patterns that I have seen with the apply() function:

df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7], 'b': [7, 6, 5, 4, 3, 2, 1]})

# Option 1: Use a series and apply on it to run a UDF (Technically a Series.apply)
df['c'] = df['a'].apply(lambda x: x*2)

# Option 2: Use the entire row from a dataframe
df['d'] = df.apply(lambda x: x.a*x.b, axis=1)

# Option 3: Use some columns in a row
df['e'] = df[['a', 'b']].apply(lambda x: x.a*x.b, axis=1)

# Option 4: Apply aggregates on every column to create a new dataframe
new_df = df.apply(lambda x: x.sum())

# Option 5: Apply aggregates but return multiple values - here it adjusts the shorted columns with NaNs
new_df = df.apply(lambda x: x[:x[0]])

[HyukjinKwon]

There are some problems to use it directly.

1. Current Scalar UDF doesn't support aggregation across whole data. If we do the same thing, it will aggregate each batch unit that's internally used (arrow batch).

2. If we restrict it's 1 to 1 then we can consider transform instead. But we shouldn't forget we should check if the input & output length are same.

3. When axis is different, it will need a different execution style to row by row, which will probably be like d397c1a way.

[HyukjinKwon]

Actually, I think we could hack this via df.groupby(F.spark_partition_id()).apply(pands_udf). Let me take a look soon.

[chunyang]

Would groupby-apply introduce exchange/shuffle on every apply? If so, seems like that would be an expensive operation.

[HyukjinKwon]

Yeah, we should avoid. With Spark 3.0, we can avoid a shuffle with mapInPandas API.