A prototype of vectorized UDAF No. 1.

[ueshin]

What changes were proposed in this pull request?

This is a prototype of vectorized UDAF.

Proposed API

Introduce @pandas_udaf decorator (annotation) to define vectorized UDAFs which takes one or more pandas.Series and returns one or more scalar values.

We can define vectorized UDAFs if the function is algebraic as:

@pandas_udaf(LongType(), algebraic=True)
def p_sum(v):
    return v.sum()

or if the function is not algebraic with additional parameters partial and bufferType as:

@pandas_udaf(
    DoubleType(),
    algebraic=False,
    partial=lambda v: (v.sum(), len(v)),
    bufferType=StructType().add("sum", LongType()).add("count", LongType()))
def p_avg(sum, count):
    return (sum.sum() / count.sum())

We can use it similar to aggregate functions as:

df.groupBy(col('g')).agg(p_sum(col('n')), expr('count(n)'), p_avg(col('n')))