Репликация строк в кадре данных 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 доступно с 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.

Answer 2

Вы можете использовать функцию 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|
+---+