Unnesting a 2D array of structs into a struct of 2D arrays - apache-spark

Question

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

Answer 1

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