PySpark 中的状态聚合函数

Question

在 PySpark 中，我正在尝试定义一个自定义聚合器正在累积状态。在 Spark 2.3 中可以吗？

AFAIK，现在可以通过使用 PandasUDFType.GROUPED_AGG 关键字调用 pandas_udf，从 Spark 2.3 开始在 PySpark 中定义自定义 UDAF（参见 How to define and use a User-Defined Aggregate Function in Spark SQL?）。但是，鉴于它只是将函数作为参数，我认为在聚合期间不可能携带状态。

从 Scala 中，我看到可以通过扩展 UserDefinedAggregateFunction 或 org.apache.spark.sql.expressions.Aggregator 来进行有状态聚合，但是我只能在 python 端做类似的事情吗？

Answer 1

您可以使用accumulator。

您可以利用内置的 Spark Streaming state management。

用于 SQL 的简单累加器示例

from  pyspark.sql.types import IntegerType

# have some data
df = spark.range(10).toDF("num")

# have a table
df.createOrReplaceTempView("num_table")

# have an accumulator
accSum = sc.accumulator(0)

# have a function that accumulates
def add_acc(int_val):
  accSum.add(int_val)
  return int_val

# register function as udf
spark.udf.register("reg_addacc", add_acc, IntegerType())

# use in sql
spark.sql("SELECT sum(reg_addacc(num)) FROM num_table").show()

# get value from accumulator
print(accSum.value)

45