Question

我有一个如下的PySpark数据框：

+---+----+----+
| id|day1|day2|
+---+----+----+
|  1|   2|   4|
|  2|   1|   2|
|  3|   3|   3|
+---+----+----+

我想在n = day2 - day1处每行重复n次。生成的数据框如下所示：

+---+----+----+
| id|day1|day2|
+---+----+----+
|  1|   2|   4|
|  1|   2|   4|
|  1|   2|   4|
|  2|   1|   2|
|  2|   1|   2|
|  3|   3|   3|
+---+----+----+

我该怎么做？

Answer 1

这是一种方法。

from pyspark.sql import functions as F
from pyspark.sql.types import *

@F.udf(ArrayType(StringType()))
def gen_array(day1, day2):
  return ['' for i in range(day2-day1+1)]

df.withColumn(
  "dup", 
  F.explode(
    gen_array(F.col("day1"), F.col("day2"))
  )
).drop("dup").show()

+---+----+----+
| id|day1|day2|
+---+----+----+
|  1|   2|   4|
|  1|   2|   4|
|  1|   2|   4|
|  2|   1|   2|
|  2|   1|   2|
|  3|   3|   3|
+---+----+----+

Answer 2

另一个使用rdd.flatMap的选项：

df.rdd.flatMap(lambda r: [r] * (r.day2 - r.day1 + 1)).toDF().show()
+---+----+----+
| id|day1|day2|
+---+----+----+
|  1|   2|   4|
|  1|   2|   4|
|  1|   2|   4|
|  2|   1|   2|
|  2|   1|   2|
|  3|   3|   3|
+---+----+----+

基于两列之间的差异的Pyspark DataFrame动态次数中的重复行

2 个答案: