Scala / Spark：将方法应用于数据集的多个子集

Question

考虑具有以下结构的数据集。

case class Order(date: Date, customer: Int, product: Int, quantity: Int)
val data = ...
  .as[Order]

我有一种方法可以计算给定客户是否会在未来的某个时间段内为给定产品下订单。

//df only contains observations for one customer, and one product
def genPrediction(df: Dataset[Order]): Double = ...

data包含许多不同的customer / product组合。 我想做的是将data和customer分组。 product，并使用val groupedData = data.groupBy("customer", "product")对其进行汇总。 结果应该是具有以下结构的数据集：

genPrediction

是在写自己的case class Prediction(customer: Int, product: Int, probability: Double) 的方式，还是有一种更简单的方法可以利用我现有的方法？我正在处理大量数据，因此效率很重要。

编辑： 我找到了一个解决方案，很想获得有关这是否是一种好的方法的反馈。

Aggregator