apache spark - Pyspark Replicate Row based on column value

Question

I would like to replicate all rows in my DataFrame based on the value of a given column on each row, and than index each new row. Suppose I have:

Column A Column B
T1       3
T2       2

I want the result to be:

Column A Column B Index
T1       3        1
T1       3        2
T1       3        3
T2       2        1
T2       2        2

I was able to to something similar with fixed values, but not by using the information found on the column. My current working code for fixed values is:

idx = [lit(i) for i in range(1, 10)]
df = df.withColumn('Index', explode(array( idx ) ))

I tried to change:

lit(i) for i in range(1, 10) 

to

lit(i) for i in range(1, df['Column B'])

and add it into my array() function:

df = df.withColumn('Index', explode(array( lit(i) for i in range(1, df['Column B']) ) ))

but it does not work (TypeError: 'Column' object cannot be interpreted as an integer).

How should I implement this?

Answer 1

Unfortunately you can't iterate over a Column like that. You can always use a udf, but I do have a non-udf hack solution that should work for you if you're using Spark version 2.1 or higher.

The trick is to take advantage of pyspark.sql.functions.posexplode() to get the index value. We do this by creating a string by repeating a comma Column B times. Then we split this string on the comma, and use posexplode to get the index.

df.createOrReplaceTempView("df")  # first register the DataFrame as a temp table

query = 'SELECT '
    '`Column A`,'
    '`Column B`,'
    'pos AS Index '
    'FROM ( '
        'SELECT DISTINCT '
        '`Column A`,'
        '`Column B`,'
        'posexplode(split(repeat(",", `Column B`), ",")) '
        'FROM df) AS a '
    'WHERE a.pos > 0'
newDF = sqlCtx.sql(query).sort("Column A", "Column B", "Index")
newDF.show()
#+--------+--------+-----+
#|Column A|Column B|Index|
#+--------+--------+-----+
#|      T1|       3|    1|
#|      T1|       3|    2|
#|      T1|       3|    3|
#|      T2|       2|    1|
#|      T2|       2|    2|
#+--------+--------+-----+

Note: You need to wrap the column names in backticks since they have spaces in them as explained in this post: How to express a column which name contains spaces in Spark SQL