scala - scala 的 "collect"采用偏函数的 Spark 数据集等价物

Question

常规 scala 集合有一个漂亮的 collect 方法，它让我可以使用偏函数一次性完成 filter-map 操作。 spark Dataset s 上是否有等效的操作？

我喜欢它有两个原因:

语法简单

它将 filter-map 样式的操作减少到单次传递(尽管在 spark 我猜有一些优化可以为您发现这些东西)

这是一个例子来说明我的意思。假设我有一个选项序列，我想提取并加倍定义的整数(在 Some 中的那些):

val input = Seq(Some(3), None, Some(-1), None, Some(4), Some(5)) 

方法 1 - collect

input.collect {
  case Some(value) => value * 2
} 
// List(6, -2, 8, 10)

collect 使这在语法上非常简洁，并执行一次。

方法 2 - filter-map

input.filter(_.isDefined).map(_.get * 2)

我可以将这种模式用于 spark，因为数据集和数据框有类似的方法。

但我不太喜欢这个，因为 isDefined 和 get 对我来说似乎是代码的味道。有一个隐含的假设 map 只接收 Some s。编译器无法验证这一点。在更大的示例中，开发人员更难发现该假设，并且开发人员可能会交换过滤器并进行映射，例如不会出现语法错误。

方法 3 - fold* 操作

input.foldRight[List[Int]](Nil) {
  case (nextOpt, acc) => nextOpt match {
    case Some(next) => next*2 :: acc
    case None => acc
  }
}

我没有使用足够的 Spark 来知道 fold 是否有等价物，所以这可能有点切线。

无论如何，模式匹配、折叠样板和列表的重建都混杂在一起，很难阅读。

所以总的来说，我发现 collect 语法最好，我希望 spark 有这样的东西。

Answer 1

collect 和 RDD 上定义的 Dataset 方法用于在驱动程序中实现数据。

尽管没有类似于 Collections API collect 方法的东西，但您的直觉是正确的:因为这两个操作都是惰性求值的，所以引擎有机会优化这些操作并将它们链接起来，以便它们以最大的局部性执行。

对于您特别提到的用例，我建议您考虑 flatMap ，它适用于 RDD 和 Dataset s:

// Assumes the usual spark-shell environment
// sc: SparkContext, spark: SparkSession
val collection = Seq(Some(1), None, Some(2), None, Some(3))
val rdd = sc.parallelize(collection)
val dataset = spark.createDataset(rdd)

// Both operations will yield `Array(2, 4, 6)`
rdd.flatMap(_.map(_ * 2)).collect
dataset.flatMap(_.map(_ * 2)).collect

// You can also express the operation in terms of a for-comprehension
(for (option <- rdd; n <- option) yield n * 2).collect
(for (option <- dataset; n <- option) yield n * 2).collect

// The same approach is valid for traditional collections as well
collection.flatMap(_.map(_ * 2))
for (option <- collection; n <- option) yield n * 2

编辑

正如在另一个问题中正确指出的那样， RDD s 实际上有 collect 方法，它通过应用偏函数来转换 RDD ，就像它在普通集合中发生的那样。然而，正如 Spark documentation 指出的那样，“只有在预期结果数组很小的情况下才应该使用这种方法，因为所有数据都加载到驱动程序的内存中。”