分区

时间:2017-06-24 09:08:20

标签: scala apache-spark spark-dataframe udf

我有200 Mil的行,1K组看起来像这样

Group     X             Y             Z          Q           W
group1  0.054464866 0.002248819 0.299069804 0.763352879 0.395905106
group2  0.9986218   0.023649037 0.50762069  0.212225807 0.619571705
group1  0.839928517 0.290339179 0.050407454 0.75837838  0.495466007
group1  0.021003132 0.663366686 0.687928832 0.239132224 0.020848608
group1  0.393843426 0.006299292 0.141103438 0.858481036 0.715860852
group2  0.045960198 0.014858905 0.672267793 0.59750871  0.893646818

我想为每个群组运行相同的功能(对linear regression X[X, Z, Q, W])。我本可以做Window.partition等,但我有自己的功能。目前,我做了以下事情:

df.select("Group").distinct.collect.toList.foreach{group => 
val dfGroup = df.filter(col("Group")===group
dfGroup.withColumn("res", myUdf(col("X"), col("Y"), col("Z"), col("Q"), col("W"))}

想知道是否有更好的方法吗?

1 个答案:

答案 0 :(得分:0)

根据您的喜好,您至少有两个选项:DataFrame或Dataset。

使用UDAF的DataFrame

df
  .groupBy("group")
  .agg(myUdaf(col("col1"), col("col2")))

其中myUdaf是UDAF

您可以在此找到如何实施UDAF的示例:https://docs.databricks.com/spark/latest/spark-sql/udaf-scala.html

数据集

您可以使用数据集API中的groupByKeymapGroups转换:

ds
  .groupByKey(_.group)
  .mapGroups{case (group, values) =>
    (group, aggregator(values))
  }

其中aggregator是负责聚合对象集合的Scala函数。

如果您不需要汇总,则可以使用values转换来映射map,例如:

values.map(v => fun(...))