Pyspark Generate rows depending on column value

Question

Below is the data input,

|       start       |   format_date     |    diff|
+-------------------+-------------------+--------+
|2019-11-15 20:30:00|2019-11-15 18:30:00|     4  |

Expected output:

start                     format_date                      Diff                    seq
2019-11-15 20:30:00     2019-11-15 18:30:00                  4                       1
2019-11-15 20:30:00     2019-11-15 18:30:00                  4                       2
2019-11-15 20:30:00     2019-11-15 18:30:00                  4                       3
2019-11-15 20:30:00     2019-11-15 18:30:00                  4                       4

how do i generate rows depending up on the value (diff) of a column?

[How do I format my code blocks?](https://meta.stackexchange.com/questions/22186/how-do-i-format-my-code-blocks) — pault, Nov 15 '19 at 20:43
Possible duplicate of [Pyspark Replicate Row based on column value](https://stackoverflow.com/questions/51109018/pyspark-replicate-row-based-on-column-value) — pault, Nov 15 '19 at 20:43

score 3 · Answer 1 · answered Nov 15 '19 at 22:27

Spark 2.4 or greater solution

from pyspark.sql import functions as F

from pyspark.sql.types import *

df= spark.createDataFrame([["2019-11-15 20:30:00","2019-11-15 18:30:00" ,4]], ["start", "format_date", "diff"])


df.select("*", F.explode(F.sequence(F.lit(1), F.col("diff"))).alias("seq")).show


+-------------------+-------------------+----+---+
|              start|        format_date|diff|seq|
+-------------------+-------------------+----+---+
|2019-11-15 20:30:00|2019-11-15 18:30:00|   4|  1|
|2019-11-15 20:30:00|2019-11-15 18:30:00|   4|  2|
|2019-11-15 20:30:00|2019-11-15 18:30:00|   4|  3|
|2019-11-15 20:30:00|2019-11-15 18:30:00|   4|  4|

Sagar · Accepted Answer · 2019-11-15T22:34:04.620

0

Spark < 2.4

You can use explode function

import pyspark.sql.functions as F
import pyspark.sql.types as Types

def rangeArr(diff):
  return range(1,diff+1)
rangeUdf = F.udf(rangeArr, Types.ArrayType(Types.IntegerType()))

df = df.withColumn('seqArr', rangeUdf('diff'))

df = df.withColumn('seq', F.explode('seqArr'))

edited Nov 15 '19 at 22:34

answered Nov 15 '19 at 21:57

Sagar

373
1
6

Pyspark Generate rows depending on column value

2 Answers2