用户定义的功能要应用于PySpark中的Window？

Question

我试图将用户定义的函数应用于PySpark中的Window。我已经读过UDAF也许是要走的路，但是我找不到任何具体的东西。

举个例子（摘自：Xinh的技术博客，并针对PySpark进行了修改）：

from pyspark import SparkConf
from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import avg

spark = SparkSession.builder.master("local").config(conf=SparkConf()).getOrCreate()

a = spark.createDataFrame([[1, "a"], [2, "b"], [3, "c"], [4, "d"], [5, "e"]], ['ind', "state"])

customers = spark.createDataFrame([["Alice", "2016-05-01", 50.00],
                                    ["Alice", "2016-05-03", 45.00],
                                    ["Alice", "2016-05-04", 55.00],
                                    ["Bob", "2016-05-01", 25.00],
                                    ["Bob", "2016-05-04", 29.00],
                                    ["Bob", "2016-05-06", 27.00]],
                               ["name", "date", "amountSpent"])

customers.show()

window_spec = Window.partitionBy("name").orderBy("date").rowsBetween(-1, 1)

result = customers.withColumn( "movingAvg", avg(customers["amountSpent"]).over(window_spec))

result.show()

我正在应用avg已经内置的函数，pyspark.sql.functions但是如果avg我不想使用更复杂的函数并编写自己的函数，该怎么办？

Answer 1

未来的工作：

火花24561 - 与熊猫UDF（界定窗口）用户定义的窗口函数正在进行AA工作。有关详细信息，请遵循相关的JIRA。

火花> = 2.4：

火花22239 - 与熊猫UDF（无界的窗口）用户定义的窗口函数用于与无界基于Windows的熊猫窗函数引入支持。总体结构是

return_type: DataType

@pandas_udf(return_type, PandasUDFType.GROUPED_AGG)
def f(v):
  return ... 

w = (Window
    .partitionBy(grouping_column)
    .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing))

df.withColumn('foo', f('bar').over(w))

请参阅doctests和单元测试以获取详细示例。

火花<2.4

你不能。窗口函数需要UserDefinedAggregateFunction或等效对象，而不是UserDefinedFunction，并且无法在PySpark中定义一个。

但是，在PySpark 2.3或更高版本中，您可以定义vectorized pandas_udf，可以将其应用于分组数据。您可以在PySpark中找到一个有效的示例，将UDF应用于GroupedData（带有有效的python示例）。尽管Pandas不提供与窗口函数直接等效的功能，但它们具有足够的表现力来实现任何类似于窗口的逻辑，尤其是使用pandas.DataFrame.rolling。此外，与一起使用的函数GroupedData.apply可以返回任意数量的行。

您还可以从PySpark Spark调用Scala UDAF ：如何使用Scala或Java用户定义函数映射Python？。

Answer 2

从 Spark 3.0.0 开始，UDF 现在可以应用于 Window。

https://spark.apache.org/docs/3.1.2/api/python/reference/api/pyspark.sql.functions.pandas_udf.html

文档摘录：

from pyspark.sql import Window

@pandas_udf("double")
def mean_udf(v: pd.Series) -> float:
    return v.mean()

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
w = Window.partitionBy('id').orderBy('v').rowsBetween(-1, 0)
df.withColumn('mean_v', mean_udf("v").over(w)).show()

+---+----+------+
| id|   v|mean_v|
+---+----+------+
|  1| 1.0|   1.0|
|  1| 2.0|   1.5|
|  2| 3.0|   3.0|
|  2| 5.0|   4.0|
|  2|10.0|   7.5|
+---+----+------+