Question

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

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

理想情况下，此过程应该自动取消嵌套所有struct字段（即，无需我手动指定字段“ 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")
)

将结构的2D数组取消嵌套为2D数组的结构

1 个答案: