[SPARK-52905][PYTHON] Arrow UDF for window

[zhengruifeng]

What changes were proposed in this pull request?

Arrow UDF for window

Why are the changes needed?

to make Arrow UDF support window operation

Does this PR introduce any user-facing change?

Not, yet. Will make Arrow UDF public soon

In [1]: from typing import Iterator, Tuple
   ...: import pyarrow as pa
   ...: from pyspark.sql import Window
   ...: from pyspark.sql import functions as sf
   ...: from pyspark.sql.pandas.functions import arrow_udf
   ...:
   ...: import pandas as pd
   ...: from pyspark.sql.functions import pandas_udf
   ...: from pyspark.sql import Window
   ...:
   ...: df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
   ...:
   ...: w = Window.partitionBy('id').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
   ...:
   ...:

In [2]: @arrow_udf("double")
   ...: def arrow_mean_udf(v: pa.Array) -> float:
   ...:    assert isinstance(v, pa.Array), str(type(v))
   ...:    return pa.compute.mean(v)
   ...:
   ...: # df.select(arrow_mean_udf(df['v'])).show()
   ...: # df.groupby("id").agg(arrow_mean_udf('v')).show()
   ...:
   ...: df.withColumn('mean_v', arrow_mean_udf(df['v']).over(w)).show()
   ...:
   ...:
+---+----+------+
| id|   v|mean_v|
+---+----+------+
|  1| 1.0|   1.5|
|  1| 2.0|   1.5|
|  2| 3.0|   6.0|
|  2| 5.0|   6.0|
|  2|10.0|   6.0|
+---+----+------+

How was this patch tested?

New tests

Was this patch authored or co-authored using generative AI tooling?

No

[xinrong-meng]

LGTM! Would you fix the lint (too long line)?

[zhengruifeng]

thanks, merged to master

[viirya]

WindowInPandasExec? I guess you meant ArrowWindowPythonExec?

[zhengruifeng]

yes, will fix it in #51648

[viirya]

pandas_window_bound_types? Maybe we should have arrow_window_bound_types?

[zhengruifeng]

Arrow UDF reuses the same https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala#L71

Maybe we can give it a better name