Unnesting a 2D array of structs into a struct of 2D arrays

user_7039310 0
Viewed 127

I have a column of type array<array<struct<a: String, b: Int>>>.

I want a column of type struct<a: array<array<String>>, b: array<array<Int>>.

Ideally, this procedure should unnest all struct fields automatically (i.e. without me having to specify fields "a" and "b" manually), but anything that works would be extremely helpful here.

Example code that I have (I'm trying to turn ds into 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")
apache-spark
1 Answers

The shortest solution is to use transform higher order function (introduced in 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"
)

In older version you'll need either equivalent udf* or use typed map:

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

The former solution can be generalized:

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")

The latter one probably not so much, unless you want to use Scala reflection (and that's a serious overkill).

* Something around these lines should do the trick:

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")
)
user_109383620
Related