在Pyspark上模拟UDAF以进行封装

Question

我正在使用PySpark学习Spark，并在试图让事情变得更清洁时碰壁。

假设有一个看起来像这样的数据框。 （当然，有更多的列和行）

A | B |   C
--+---+------
a | 1 | 1.300
a | 2 | 2.500
a | 3 | 1.000
b | 1 | 120.0
b | 4 | 34.20
c | 2 | 3.442

我希望使用基本的groupby -> agg运行一堆pyspark.sql.functions，例如count()和mean()，如下所示：

df.groupby("A")\
    .agg(mean("B").alias("B_mean"),
         sum("C").alias("C_sum"),
         (countDistinct("B")/avg("C")).alias("New_metric"))

它工作正常，运行速度相对较快，并给我预期的效果。

但是，最终，需要稍微复杂的功能，而且，我们希望这些功能更容易测试。

如何封装这些功能？使用lambda？某种方式围绕UDF？

我知道UDAF，并且可以在SCALA中编写它们并将代码导入PySpark，但是，由于我们所有的代码库都已经在Python中，我想探索其他选项。

P.S。：我们正在运行Spark 1.6.0

Answer 1

可以将函数定义为pyspark.sql.functions：

的组合

• 是 - 走这条路。例如：

  def sum_of_squares(col):
      return sum(col * col)
  
  df.select(sum_of_squares(df["foo"]])
  
  df.groupBy("foo").agg(sum_of_squares(df["bar"]])
  

• 否 - 使用RDD。