Question

我有一个Spark DataFrame，其中一个列是一个对象数组。我想做一个过滤该数组的操作。在我的下面的例子中，我有一个有孩子的父母，我想只让成年子女。

import spark.implicits._

case class Child(name: String, age: Int)
case class Parent(name: String, children: Array[Child])

val rawData = Seq(Parent("Mom", Array(Child("Jane", 9))), Parent("Dad", Array(Child("Hubert", 28), Child("David", 27), Child("Jim", 25))))
val data = spark.createDataFrame(rawData)

我能来的最近的是：

val adultChildren = udf((children: mutable.WrappedArray[Child]) => {
  val rowArray = children.asInstanceOf[mutable.WrappedArray[GenericRowWithSchema]]
  val ret = rowArray.filter(c => c.getAs[Int]("age") > 18)
  ret.asInstanceOf[mutable.WrappedArray[Child]]
})
data.select(adultChildren($"children")).show()

这有点烦人。我想，优点是Spark花费更少的时间（de）序列化对象，但它很冗长。

有更简洁的方法吗？

Answer 1

如果您可以使用数据集，那么它变得非常简单：

data.map(_.children.filter(_.age > 18).toList)

但是，如果你对DataFrames感激不尽：

data.select($"name", explode($"children").as("child"))
    .where($"child.age" > 18)
    .groupBy($"name").agg(collect_list($"child"))

Answer 2

一个改进是将样板封装在一个函数中：

import scala.reflect.runtime.universe._
def arrayFilterUDF[T: TypeTag](f: (GenericRowWithSchema) => Boolean) = udf((a: mutable.WrappedArray[T]) => {
    val rowArray = a.asInstanceOf[mutable.WrappedArray[GenericRowWithSchema]]
    rowArray.filter(f).asInstanceOf[mutable.WrappedArray[T]]
})

这允许你写：

val adultChildren = arrayFilterUDF[Child](_.getAs[Int]("age") > 18)
data.select(adultChildren($"children")).show()

什么是将过滤器应用于Spark中

2 个答案: