apache-spark - 将二维结构数组取消嵌套为二维数组结构

Question

我有一个类型为 array<array<struct<a: String, b: Int>>> 的列.

我想要一个 struct<a: array<array<String>>, b: array<array<Int>> 类型的列.

理想情况下，此过程应自动解除所有结构字段的嵌套(即无需我手动指定字段“a”和“b”)，但任何可行的方法在这里都将非常有用。

我的示例代码(我正在尝试将 ds 变成 expected )。

case class Struct(foo: String, bar: Int)
case class Schema(structs: Vector[Vector[Struct]])

val ss = spark
import ss.implicits._

val ds = Seq(Schema(Vector(Vector(Struct("a", 1), Struct("b", 2)), Vector(Struct("c", 3))))).toDS

val expected = Seq(
    (Vector(Vector("a", "b"), Vector("c")), Vector(Vector(1, 2), Vector(3)))
).toDF("foo", "bar")

Answer 1

最短的解决方案是使用transform 高阶函数(在 Spark 2.4 中引入):

ds.selectExpr(
  "transform(structs, xs -> transform(xs, x -> x.foo)) as foo",
  "transform(structs, xs -> transform(xs, x -> x.bar)) as bar"
)

在旧版本中，您需要等效的 udf* 或使用类型化的 map:

ds.as[Schema]
  .map(x => (
    x.structs.map(_.map(_.foo)), 
    x.structs.map(_.map(_.bar))
)).toDF("foo", "bar")

前一个解决方案可以推广:

import org.apache.spark.sql.types._ 
import org.apache.spark.sql.DataFrame

def expand(ds: DataFrame, col: String) = {

  val fields = ds.schema(col).dataType match {
    case ArrayType(ArrayType(s: StructType, _), _) => s.fieldNames
  }
  val exprs = fields.map {
    field => expr(
      s"transform(`$col`, xs -> transform(xs, x -> x.`$field`)) as `$field`"
    )
  }
  ds.select(exprs: _*)
}

expand(ds.toDF, "structs")

后者可能不多，除非你想使用 Scala 反射(这是一个严重的矫枉过正)。

---

* 这些行周围的东西应该可以解决问题:

import scala.reflect.runtime.universe.TypeTag
import org.apache.spark.sql.functions.udf

def extract[T : TypeTag](field: String) = udf(
  (xs: Seq[Seq[Row]]) => xs.map(_.map(_.getAs[T](field)))
)

val extractString = extract[String] _
val extractInt = extract[Int] _

ds.select(
  extractString("foo")($"structs").as("foo"),
  extractInt("bar")($"structs").as("bar")
)