在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()和{}，如下所示：

^{pr2}$

它运行良好，运行速度相对较快，并给我期望的结果。在

但是，最终将需要稍微复杂一些的函数，而且，我们希望使这些函数更易于测试。在

如何封装这些函数？使用lambda？有什么办法绕过UDF？在

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

注：我们正在运行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。