spark dataframe：如何爆炸IntegerType列

Question

val schema = StructType(Array(StructField("id", IntegerType, false),StructField("num", IntegerType, false)))

我希望每个id生成从0到num的连续数字。 我不知道怎么办.. 感谢

data and result here !!!

Answer 1

您可以使用UDF和explode功能：

import org.apache.spark.sql.functions.{udf, explode}

val range = udf((i: Int) => (0 to i).toArray)
df.withColumn("num", explode(range($"num")))

Answer 2

尝试DataFrame.explode：

df.explode(col("id"), col("num")) {case row: Row =>
    val id = row(0).asInstanceOf[Int]
    val num = row(1).asInstanceOf[Int]
    (0 to num).map((id, _))
}

或者在RDD领域，您可以使用flatmap：

df.rdd.flatMap(x => (0 to x._2).map((x._1, _)))