I am struggling to transform my pyspark dataframe which looks like this:
df = spark.createDataFrame([('0018aad4',[300, 450], ['{"v1": "blue"}', '{"v2": "red"}']), ('0018aad5',[300], ['{"v1": "blue"}'])],[ "id","Tlist", 'Tstring']) df.show(2, False) +--------+----------+-------------------------------+ |id |Tlist |Tstring | +--------+----------+-------------------------------+ |0018aad4|[300, 450]|[{"v1": "blue"}, {"v2": "red"}]| |0018aad5|[300] |[{"v1": "blue"}] | +--------+----------+-------------------------------+
to this:
df_result = spark.createDataFrame([('0018aad4',[300, 450], 'blue', 'red'), ('0018aad5',[300], 'blue', None)],[ "id","Tlist", 'v1', 'v2']) df_result.show(2, False) +--------+----------+----+----+ |id |Tlist |v1 |v2 | +--------+----------+----+----+ |0018aad4|[300, 450]|blue|red | |0018aad5|[300] |blue|null| +--------+----------+----+----+
I tried to pivot and a bunch of others things but don’t get the result above.
Note that I don’t have the exact number of dict in the column Tstring
Do you know how I can do this?
Advertisement
Answer
Using transform
function you can convert each element of the array into a map type. After that, you can use aggregate
function to get one map, explode it then pivot the keys to get the desired output:
from pyspark.sql import functions as F df1 = df.withColumn( "Tstring", F.transform("Tstring", lambda x: F.from_json(x, "map<string,string>")) ).withColumn( "Tstring", F.aggregate( F.expr("slice(Tstring, 2, size(Tstring))"), F.col("Tstring")[0], lambda acc, x: F.map_concat(acc, x) ) ).select( "id", "Tlist", F.explode("Tstring") ).groupby( "id", "Tlist" ).pivot("key").agg(F.first("value")) df1.show() #+--------+----------+----+----+ #|id |Tlist |v1 |v2 | #+--------+----------+----+----+ #|0018aad4|[300, 450]|blue|red | #|0018aad5|[300] |blue|null| #+--------+----------+----+----+
I’m using Spark 3.1+, so the higher-order functions such as transform
are available in dataframe API but you can do the same using expr
for spark <3.1:
df1 = (df.withColumn("Tstring", F.expr("transform(Tstring, x-> from_json(x, 'map<string,string>'))")) .withColumn("Tstring", F.expr("aggregate(slice(Tstring, 2, size(Tstring)), Tstring[0], (acc, x) -> map_concat(acc, x))")) .select("id", "Tlist", F.explode("Tstring")) .groupby("id", "Tlist") .pivot("key") .agg(F.first("value")) )