Question

我想根据给定列的值复制行。例如，我得到了这个DataFrame：

+-----+
|count|
+-----+
|    3|
|    1|
|    4|
+-----+

我想得到：

+-----+
|count|
+-----+
|    3|
|    3|
|    3|
|    1|
|    4|
|    4|
|    4|
|    4|
+-----+

根据this answer，我尝试使用withColum方法。

val replicateDf = originalDf
    .withColumn("replicating", explode(array((1 until $"count").map(lit): _*)))
    .select("count")

但是$"count"是ColumnName，不能用于在上述表达式中表示其值。

（我也尝试过explode(Array.fill($"count"){1})，但这里也有同样的问题。）

我需要更改什么？有没有更清洁的方法？

Answer 1

您可以使用array_repeat功能：

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

val df = Seq(1, 2, 3).toDF

df.select(explode(array_repeat($"value", $"value"))).show()

+---+
|col|
+---+
|  1|
|  2|
|  2|
|  3|
|  3|
|  3|
+---+

Answer 2

array_repeat从2.4开始可用。如果需要较低版本的解决方案，则可以使用udf（）或rdd。对于Rdd，请查看此

import scala.collection.mutable._

val df = Seq(3,1,4).toDF("count")
val rdd1 = df.rdd.flatMap( x=> { val y = x.getAs[Int]("count"); for ( p <- 0 until y ) yield Row(y) }  )
spark.createDataFrame(rdd1,df.schema).show(false)

结果：

+-----+
|count|
+-----+
|3    |
|3    |
|3    |
|1    |
|4    |
|4    |
|4    |
|4    |
+-----+

仅使用df（）

scala> df.flatMap( r=> { (0 until r.getInt(0)).map( i => r.getInt(0)) } ).show
+-----+
|value|
+-----+
|    3|
|    3|
|    3|
|    1|
|    4|
|    4|
|    4|
|    4|
+-----+

对于udf（），下面的方法会起作用

val df = Seq(3,1,4).toDF("count")
def array_repeat(x:Int):Array[Int]={
  val y = for ( p <- 0 until x )yield x
  y.toArray
}
val udf_array_repeat = udf (array_repeat(_:Int):Array[Int] )
df.withColumn("count2", explode(udf_array_repeat('count))).select("count2").show(false)

编辑：

请查看下面的@ user10465355答案，以获取有关array_repeat的更多信息。

根据列中的值复制Spark数据框中的行

2 个答案: