scala - Scala 中的 Spark 分组 map UDF

Question

我正在尝试编写一些代码，以允许我对数据帧的一组行计算某些操作。在 PySpark 中，这可以通过定义 GROUPED_MAP 类型的 Pandas UDF 来实现。但是，在 Scala 中，我只找到了一种创建自定义聚合器 (UDAF) 或经典 UDF 的方法。

我的临时解决方案是生成一个键列表，这些键将对我的组进行编码，这将允许我过滤数据帧并对数据帧的每个子集执行我的操作。然而，这种方法不是最优的并且非常慢。所执行的 Action 是按顺序进行的，因此会花费很多时间。我可以将循环并行化，但我确信这会显示任何改进，因为 Spark 已经分发。

有没有更好的方法来做我想做的事？

编辑:尝试使用 Futures 进行并行化，但没有像预期的那样提高速度

Answer 1

据我所知，这在 Scala 中是不可能的。根据您的需要，我认为可能还有其他方法可以将转换应用于 Spark/Scala 中的一组行:

1. 做一个groupBy(...).agg(collect_list(<column_names>)) ，并使用对值数组进行操作的 UDF。如果需要，您可以将 select 语句与 explode(<array_column>) 结合使用恢复到原来的格式
2. 尝试使用窗口函数重写您想要实现的目标。您可以添加一个带有聚合表达式的新列，如下所示:

import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy('group)

val result = spark.range(100)
    .withColumn("group", pmod('id, lit(3)))
    .withColumn("group_sum", sum('id).over(w))