PySpark: Add a column to DataFrame when column is a list

user_2352319 0
Viewed 16503

I have read similar questions but couldn't find a solution to my specific problem.

I have a list 

l = [1, 2, 3]

and a DataFrame

df = sc.parallelize([
    ['p1', 'a'],
    ['p2', 'b'],
    ['p3', 'c'],
]).toDF(('product', 'name'))

I would like to obtain a new DataFrame where the list l is added as a further column, namely

+-------+----+---------+
|product|name| new_col |
+-------+----+---------+
|     p1|   a|     1   |
|     p2|   b|     2   |
|     p3|   c|     3   |
+-------+----+---------+

Approaches with JOIN, where I was joining df with an 

 sc.parallelize([[1], [2], [3]])

have failed. Approaches using withColumn, as in

new_df = df.withColumn('new_col', l)

have failed because the list is not a Column object.

pyspark dataframe python
3 Answers

If the product column is unique then consider the following approach:

original dataframe:

df = spark.sparkContext.parallelize([
    ['p1', 'a'],
    ['p2', 'b'],
    ['p3', 'c'],
]).toDF(('product', 'name'))

df.show()

+-------+----+
|product|name|
+-------+----+
|     p1|   a|
|     p2|   b|
|     p3|   c|
+-------+----+

new column (and new index column):

lst = [1, 2, 3]
indx = ['p1','p2','p3']

create a new dataframe from the list above (with an index):

from pyspark.sql.types import *
myschema= StructType([ StructField("indx", StringType(), True),
                       StructField("newCol", IntegerType(), True)                       
                     ])
df1=spark.createDataFrame(zip(indx,lst),schema = myschema)
df1.show()
+----+------+
|indx|newCol|
+----+------+
|  p1|     1|
|  p2|     2|
|  p3|     3|
+----+------+

join this to the original dataframe, using the index created:

dfnew = df.join(df1, df.product == df1.indx,how='left')\
          .drop(df1.indx)\
          .sort("product")

to get:

dfnew.show()

+-------+----+------+
|product|name|newCol|
+-------+----+------+
|     p1|   a|     1|
|     p2|   b|     2|
|     p3|   c|     3|
+-------+----+------+
user_60443120

This is achievable via RDDs.

1 Convert dataframes to indexed rdds:

df_rdd = df.rdd.zipWithIndex().map(lambda row: (row[1], (row[0][0], row[0][1])))
l_rdd = sc.parallelize(l).zipWithIndex().map(lambda row: (row[1], row[0]))

2 Join two RDDs on index, drop index and rearrange elements:

res_rdd = df_rdd.join(l_rdd).map(lambda row: [row[1][0][0], row[1][0][1], row[1][1]])

3 Convert result to Dataframe:

res_df = res_rdd.toDF(['product', 'name', 'new_col'])
res_df.show()

+-------+----+-------+
|product|name|new_col|
+-------+----+-------+
|     p1|   a|      1|
|     p2|   b|      2|
|     p3|   c|      3|
+-------+----+-------+
user_103223200
Related