spark-dev mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From Li Jin <>
Subject Re: [DISCUSS] PySpark Window UDF
Date Wed, 05 Sep 2018 21:51:43 GMT
Hello again!

I recently implemented a proof-of-concept implementation of proposal above.
I think the results are pretty exciting so I want to share my findings with
the community. I have implemented two variants of the pandas window UDF -
one that takes pandas.Series as input and one that takes numpy array as
input. I benchmarked with rolling mean on 1M doubles and here are some

Spark SQL window function: 20s
Pandas variant: ~60s
Numpy variant: 10s
Numpy variant with numba: 4s

You can see the benchmark code here:

I think the results are quite exciting because:
(1) numpy variant even outperforms the Spark SQL window function
(2) numpy variant with numba has the best performance as well as the
flexibility to allow users to write window functions in pure python

The Pandas variant is not bad either (1.5x faster than existing UDF with
collect_list) but the numpy variant definitely has much better performance.

So far all Pandas UDFs interacts with Pandas data structure rather than
numpy data structure, but the window UDF result might be a good reason to
open up numpy variants of Pandas UDFs. What do people think? I'd love to
hear community's feedbacks.

You can reproduce benchmark with numpy variant by using the branch:

PR link:

On Wed, May 16, 2018 at 3:34 PM Li Jin <> wrote:

> Hi All,
> I have been looking into leverage the Arrow and Pandas UDF work we have
> done so far for Window UDF in PySpark. I have done some investigation and
> believe there is a way to do PySpark window UDF efficiently.
> The basic idea is instead of passing each window to Python separately, we
> can pass a "batch of windows" as an Arrow Batch of rows + begin/end indices
> for each window (indices are computed on the Java side), and then rolling
> over the begin/end indices in Python and applies the UDF.
> I have written my investigation in more details here:
> I think this is a pretty promising and hope to get some feedback from the
> community about this approach. Let's discuss! :)
> Li

View raw message