Lower order Spark Dataframe Array concatenation of individual elements

Viewed 45

I am using Spark 3.x higher-order array functions like this:

%scala
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.sql.Column

val arrayStructureData = Seq(
Row(1,List(2,5,1,3),List(0.1, 0.5, 0.7, 0.8)),
Row(2,List(2,1),List(0.2, 0.3)),
Row(1,List(1,5),List(0.4, 0.3)),
Row(2,List(3,2),List(0.0, 0.1)) 
)
// Just a single StructType for the Row
val arrayStructureSchema = new StructType()
    .add("id",IntegerType)
    .add("prop1", ArrayType(IntegerType))
    .add("values", ArrayType(DoubleType))
val df = spark.createDataFrame(
    spark.sparkContext.parallelize(arrayStructureData),arrayStructureSchema)
df.printSchema()
df.show()

val resDF = df.withColumn(
  "jCols",
  zip_with(
    col("prop1"),
    col("values"),
      (left: Column, right: Column) => array(left, right)
  )
)

resDF.show(false)
resDF.printSchema()

so as to concat individual array elements positionally over 2 arrays to get a new sub-array. That works fine.

E.g.

[3, 2]|          [0.0, 0.1] 

returns:

[[3.0, 0.0], [2.0, 0.1]]

My question is as I cannot immediately see how I would do this without zip_with, how would we do that the easiest way? UDF?

Nice reference: https://mungingdata.com/spark-3/array-exists-forall-transform-aggregate-zip_with/ but interested in the harder bare-metal way.

0 Answers
Related