Spark:如何使用嵌套数组转置和分解列

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

// 创建SparkSession
val spark = SparkSession.builder()
  .appName("Nested Array Transpose and Column Decomposition")
  .getOrCreate()

// 创建示例数据集
val data = Seq(
  (1, Array(Array(1, 2, 3), Array(4, 5, 6))),
  (2, Array(Array(7, 8, 9), Array(10, 11, 12)))
)

// 将数据集转换为DataFrame
val df = spark.createDataFrame(data).toDF("id", "nested_array")

// 定义嵌套数组转置和分解列的UDF函数
val transposeAndDecompose = udf((nestedArray: Seq[Seq[Int]]) => {
  val transposedArray = nestedArray.transpose
  val decomposedColumns = transposedArray.map(column => column.head)
  decomposedColumns
})

// 使用UDF函数进行嵌套数组转置和分解列
val result = df.withColumn("transposed_columns", transposeAndDecompose(col("nested_array")))

// 显示结果
result.show(false)