3

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?

pault
  • 41,343
  • 15
  • 107
  • 149
LuizF Gonçalves
  • 93
  • 2
  • 10

2 Answers2

3

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

pault
  • 41,343
  • 15
  • 107
  • 149
  • great hacking @pault . but I have a query that how `Column B` in repeat is treated as a primitive and not as column in sql expression. In api form it is treated as column, isn't it? – Ramesh Maharjan Jun 30 '18 at 03:11
  • @Ramesh i had trouble getting this to work using DataFrame functions. I don't know why it works in the query. – pault Jun 30 '18 at 14:02
  • @RameshMaharjan I posted a [question](https://stackoverflow.com/questions/51140470/using-a-column-value-as-a-parameter-to-a-dataframe-function) about this behavior – pault Jul 02 '18 at 16:42
0
You can try this:

    from pyspark.sql.window import Window
    from pyspark.sql.functions import *
    from pyspark.sql.types import ArrayType, IntegerType
    from pyspark.sql import functions as F
    df = spark.read.csv('/FileStore/tables/stack1.csv', header = 'True', inferSchema = 'True')

    w = Window.orderBy("Column A")
    df = df.select(row_number().over(w).alias("Index"), col("*"))

    n_to_array = udf(lambda n : [n] * n ,ArrayType(IntegerType()))
    df2 = df.withColumn('Column B', n_to_array('Column B'))
    df3= df2.withColumn('Column B', explode('Column B'))
    df3.show()
WarBoy
  • 146
  • 1
  • 11