PySpark: How to explode list into multiple columns with sequential naming?

user_14908234 0
Viewed 112

I have the following PySpark DF:

+--------+--------------------+--------------------+
|      id|           resoFacts|             heating|
+--------+--------------------+--------------------+
|90179090|[, [No Handicap A...|[Central Heat, Fo...|
+--------+--------------------+--------------------+

created by the following:

(data_filt
     .where(col('id') == '90179090')
     .withColumn('heating', col("resoFacts").getField('heating')))

I want to create a DF that expands the list in heating into sequentially named columns, as so:

+--------------+------------+----------+----------+---------+
|  id          |heating_1   |heating_2 | heating_3|heating_4|
+--------------+------------+----------+----------+---------+
|  90179090    |Central Heat|Forced Air|      Gas |Heat Pump|
+--------------+------------+----------+----------+---------+

My furthest attempt has generated the following DF:

+---+------------+----------+----+---------+
|pos|Central Heat|Forced Air| Gas|Heat Pump|
+---+------------+----------+----+---------+
|  1|        null|Forced Air|null|     null|
|  3|        null|      null| Gas|     null|
|  2|        null|      null|null|Heat Pump|
|  0|Central Heat|      null|null|     null|
+---+------------+----------+----+---------+

with this code:

(data_filt
     .where(col('id') == '90179090')
     .withColumn('heating', col("resoFacts").getField('heating'))
     .select("heating", posexplode("heating"))
     .groupBy('pos').pivot('col').agg(first('col')))

I'm likely doing something wrong with the line beginning with groupBy. Does anyone have thoughts?

pyspark apache-spark-sql apache-spark python
1 Answers

If you only have 4 elements in the array, you can simply do that :

from pyspark.sql import functions as F

data_filt.select(
    "id",
    *(
        F.col("heating").getItem(i).alias(f"heating_{i+1}")
        for i in range(4)
    )
)

Increase the range if you have more elements.

user_50137520
Related