Not able to write loop expression in withcolumn in pyspark

Question

i have dataframe where DealKeys has data like as [{"Charge_Type": "DET", "Country": "VN", "Tariff_Loc": "VNSGN"}]

expected out put could be [{"keyname": "Charge_Type", "value": "DET", "description": "..."}, {"keyname": "Country", "value": "VN", "description": "..."}, {"keyname": "Tariff_Loc", "value": "VNSGN", "description": "..."}]

when i create dataframe got bellow error df = df2.withColumn('new_column',({'keyname' : i, 'value' : dictionary[i],'description' : "..."} for i in col("Dealkeys")))

Errro: Column is not iterable

DF2 schema:

root
 |-- Charge_No: string (nullable = true)
 |-- Status: string (nullable = true)
 |-- DealKeys: array (nullable = false)
 |    |-- element: struct (containsNull = false)
 |    |    |-- Charge_Type: string (nullable = true)
 |    |    |-- Country: string (nullable = true)
 |    |    |-- Tariff_Loc: string (nullable = true)

user_18844585 · Answer

We cannot iterate through a dataframe column in Pyspark, hence the error occured.

To get the expected output, you need to follow the approach specified below.

Create new_column column with value as an empty string beforehand so that we can update its value as we iterate through each row.
Since a column cannot be iterated, we can use collect() method to get DealKeys values so that we can insert them into the corresponding new_column value.
Using df.collect() returns a list of rows (rows can be iterated through). As per schema, DealKeys is also of type Row. Using dealkey_row (DealKeys column values as a row) with asDict(), perform list comprehension to create the list of dictionaries that will be inserted into corresponding Charge_No value.

# df is the initial dataframe
df = df.withColumn("new_column", lit(''))

rows = df.collect()
for row in rows:
    key = row[0]    #Charge_No column value (string type)
    dealkey_row = row[2][0]    #DealKeys column value (row type)
    
    lst = [{'keyname' : i, 'value' : dealkey_row[i],'description' : "..."} for i in dealkey_row.asDict()]    #row.asDict() to get dictionary
    
    df = df.withColumn('new_column', when(col('Charge_No') == key, str(lst)).otherwise(col('new_column')))
    
df.show(truncate=False)

Row.asDict() converts a row into a dictionary so that the list comprehension can be easily. Using withColumn() along with when(,) function in pyspark, insert the output of your list comprehension into the new_column column (‘otherwise’ helps to retain the previous value if Charge_No value doesn’t match).

The above code produced the following output when I used it.

Not able to write loop expression in withcolumn in pyspark

1 Answers