How to transform a column with an array of json to separate columns based on the json keys in Spark?

user_18180624 0
Viewed 39

I have a column like this:

column
[{"key":1,"value":"aaaaa"},{"key":2,"value":"bbbbb"},{"key":3,"value":"ccccc"}]
[{"key":1,"value":"abcde"},{"key":2,"value":"bcdef"}]
[{"key":1,"value":"edcba"},{"key":3,"value":"zxcvb"},{"key":4,"value":"qwert"}]

I want to separate such column base on the keys, with each key having their column. I've tried something like this but it didn't work:

test_schema = ArrayType(StructType([StructField("key", IntegerType()), StructField("value", StringType())]))

teste = (hits_raw
  .withColumn("keys", get_json_object("hitsCustomDimensions", "$[*].key"))
  .withColumn("teste", explode(from_json("hitsCustomDimensions", test_schema)))      
  ).display()

The output I want is something like this:

column_1 column_2 column_3 column_4
aaaaa bbbbb ccccc null
abcde bcdef null null
edcba null zxcvb qwert
pyspark apache-spark-sql apache-spark dataframe
1 Answers

Parse into array of structs then pivot the key column. You'll need some ID column to group by for the pivot, here I used monotonically_increasing_id function to add an id before inlining the array of structs.

from pyspark.sql import functions as F

df = spark.createDataFrame([
    ('[{"key":1,"value":"aaaaa"},{"key":2,"value":"bbbbb"},{"key":3,"value":"ccccc"}]',),
    ('[{"key":1,"value":"abcde"},{"key":2,"value":"bcdef"}]',),
    ('[{"key":1,"value":"edcba"},{"key":3,"value":"zxcvb"},{"key":4,"value":"qwert"}]',)
], ["column"])

test = (df.withColumn("column", F.from_json("column",  test_schema))
        .withColumn("id", F.monotonically_increasing_id())
        .selectExpr("id", "inline(column)")
        .groupBy("id").pivot("key").agg(F.first("value"))
        .drop("id")
        )

test.show()
#+-----+-----+-----+-----+
#|    1|    2|    3|    4|
#+-----+-----+-----+-----+
#|aaaaa|bbbbb|ccccc| null|
#|abcde|bcdef| null| null|
#|edcba| null|zxcvb|qwert|
#+-----+-----+-----+-----+

You can then rename to columns to add prefix column_* if you want:


test = test.select(*[F.col(c).alias(f"column_{c}") for c in test.columns])
user_13865510
Related