Spark explode array of string to columns

Viewed 2234

I am using Spark with Java and I have a dataframe like this:

id  | array_column
-------------------
12  | [a:123, b:125, c:456]
13  | [a:443, b:225, c:126]

I want to explode array_column with the same id, however explode doesn't work, because I want dataframe to become:

id  | a  | b  | c
-------------------
12  |123 |125 | 456 
13  |443 |225 | 126
3 Answers

The following approach will work on variable length lists in array_column. The approach uses explode to expand the list of string elements in array_column before splitting each string element using : into two different columns col_name and col_val respectively. Finally a pivot is used with a group by to transpose the data into the desired format.

The following example uses the pyspark api but can easily be translated to the java/scala apis as they are similar. I assumed your dataset is in a dataframe named input_df

from pyspark.sql import functions as F

output_df = (
    input_df.select("id",F.explode("array_column").alias("acol"))
            .select(
                "id",
                F.split("acol",":")[0].alias("col_name"),
                F.split("acol",":")[1].cast("integer").alias("col_val")
            )
            .groupBy("id")
            .pivot("col_name")
            .max("col_val")
)

Let me know if this works for you.

A very similar approach like ggordon's answer in Java:

import static org.apache.spark.sql.functions.*;

Dataset<Row> df = ...

df.withColumn("array_column", explode(col("array_column")))
        .withColumn("array_column", split(col("array_column"), ":"))
        .withColumn("key", col("array_column").getItem(0))
        .withColumn("value", col("array_column").getItem(1))
        .groupBy(col("id"))
        .pivot(col("key"))
        .agg(first("value")) //1
        .show();

Output:

+---+---+---+---+
| id|  a|  b|  c|
+---+---+---+---+
| 12|456|225|126|
| 11|123|125|456|
+---+---+---+---+

I assume that the combination of id and and the key field in the array is unique. That's why the aggregation function used at //1 is first. If this combination is not unique, the aggregation function could be changed to collect_list in order to get an array of all matching values.

Extracting column names from strings inside columns:

  • create a proper JSON string (with quote symbols around json objects and values)
  • create schema using this column
  • create struct and explode it into columns

Input example:

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [(12, ['a:123', 'b:125', 'c:456']),
     (13, ['a:443', 'b:225', 'c:126'])],
    ['id', 'array_col'])

df.show(truncate=0)
# +---+---------------------+
# |id |array_col            |
# +---+---------------------+
# |12 |[a:123, b:125, c:456]|
# |13 |[a:443, b:225, c:126]|
# +---+---------------------+

Script:

df = df.withColumn("array_col", F.expr("to_json(str_to_map(array_join(array_col, ',')))"))
json_schema = spark.read.json(df.rdd.map(lambda row: row.array_col)).schema
df = df.withColumn("array_col", F.from_json("array_col", json_schema))
df = df.select("*", "array_col.*").drop("array_col")

df.printSchema()
df.show()
# +---+---+---+---+
# | id|  a|  b|  c|
# +---+---+---+---+
# | 12|123|125|456|
# | 13|443|225|126|
# +---+---+---+---+
Related