Question

我正在尝试将用户定义的函数应用于PySpark中的Window。我读过UDAF可能是要走的路，但我找不到具体的东西。

举一个例子（取自这里：Xinh's Tech Blog并为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

未来工作：

SPARK-24561 - 用户定义的窗口函数与pandas udf（有界窗口）是一项正在进行的工作。请关注相关的JIRA。

Spark＆gt; = 2.4 ：

SPARK-22239 - 用户定义的窗口函数与pandas udf（无界窗口）引入了对具有无界窗口的基于Pandas的窗口函数的支持。一般结构是

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和unit tests。

Spark＆lt; 2.4

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

但是，在PySpark 2.3或更高版本中，您可以定义矢量化pandas_udf，它可以应用于分组数据。你可以找到一个有效的例子Applying UDFs on GroupedData in PySpark (with functioning python example)。虽然Pandas不提供直接等效的窗口函数，但有足够的表现力来实现任何类似窗口的逻辑，尤其是pandas.DataFrame.rolling。此外，与GroupedData.apply一起使用的函数可以返回任意数量的行。

您也可以从PySpark Spark: How to map Python with Scala or Java User Defined Functions?调用Scala UDAF。

Answer 2

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

https://spark.apache.org/docs/latest/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|
+---+----+------+

用户定义的函数将应用于PySpark中的Window？

2 个答案: