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

Question

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

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

我想得到：

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

根据此答案 ，我尝试使用withColumn方法。

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

从2.4开始，可以使用array_repeat 。 如果需要较低版本的解决方案，则可以使用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更多信息。