Pyspark group by collect list, to_json and pivot

There's definitely a better way to do this but I continued your to json experiment.

• Using UDFs:

After you get something like [{Alpha: {A1:1, A2:2}}, {AlphaNew: {A1:6, A2:8}}] you could create a UDF to flatten the dict. But since it's a JSON string you'll have to parse it to dict and then back again to JSON.

After that you would like to explode and pivot the table but that's not possible with JSON strings, so you have to use F.from_json with defined schema. That will give you MapType which you can explode and pivot.

Here's an example:

from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from collections import MutableMapping
import json
from pyspark.sql.types import (
    ArrayType,
    IntegerType,
    MapType,
    StringType,
)


def flatten_dict(d, parent_key="", sep="_"):
    items = []
    for k, v in d.items():
        new_key = parent_key + sep + k if parent_key else k
        if isinstance(v, MutableMapping):
            items.extend(flatten_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)


def flatten_groups(data):
    result = []
    for item in json.loads(data):
        result.append(flatten_dict(item))
    return json.dumps(result)


if __name__ == "__main__":
    spark = SparkSession.builder.master("local").appName("Test").getOrCreate()
    data = [
        (1, "Alpha", 1, 2, None, None),
        (1, "AlphaNew", 6, 8, None, None),
        (2, "Alpha", 7, 4, None, None),
        (2, "Beta", None, None, 3, 9),
    ]
    columns = ["Id", "group", "A1", "A2", "B1", "B2"]
    df = spark.createDataFrame(data, columns)
    vcols = df.columns[2:]
    df = (
        df.withColumn("json", F.struct(*vcols))
        .groupby("id")
        .agg(F.to_json(F.collect_list(F.create_map("group", "json"))).alias("json"))
    )

    # Flatten groups
    flatten_groups_udf = F.udf(lambda x: flatten_groups(x))
    schema = ArrayType(MapType(StringType(), IntegerType()))
    df = df.withColumn("json", F.from_json(flatten_groups_udf(F.col("json")), schema))

    # Explode and pivot
    df = df.select(F.col("id"), F.explode(F.col("json")).alias("json"))
    df = (
        df.select("id", F.explode("json"))
        .groupby("id")
        .pivot("key")
        .agg(F.first("value"))
    )

At the end dataframe looks like:

+---+-----------+-----------+--------+--------+-------+-------+                 
|id |AlphaNew_A1|AlphaNew_A2|Alpha_A1|Alpha_A2|Beta_B1|Beta_B2|
+---+-----------+-----------+--------+--------+-------+-------+
|1  |6          |8          |1       |2       |null   |null   |
|2  |null       |null       |7       |4       |3      |9      |
+---+-----------+-----------+--------+--------+-------+-------+

• Without UDFs:

vcols = df.columns[2:]
df = (
    df.withColumn("json", F.to_json(F.struct(*vcols)))
    .groupby("id")
    .agg(
        F.collect_list(
            F.create_map(
                "group", F.from_json("json", MapType(StringType(), IntegerType()))
            )
        ).alias("json")
    )
)
df = df.withColumn("json", F.explode(F.col("json")).alias("json"))
df = df.select("id", F.explode(F.col("json")).alias("root", "value"))
df = df.select("id", "root", F.explode(F.col("value")).alias("sub", "value"))
df = df.select(
    "id", F.concat(F.col("root"), F.lit("_"), F.col("sub")).alias("name"), "value"
)
df = df.groupBy(F.col("id")).pivot("name").agg(F.first("value"))

Result:

+---+-----------+-----------+--------+--------+-------+-------+                 
|id |AlphaNew_A1|AlphaNew_A2|Alpha_A1|Alpha_A2|Beta_B1|Beta_B2|
+---+-----------+-----------+--------+--------+-------+-------+
|1  |6          |8          |1       |2       |null   |null   |
|2  |null       |null       |7       |4       |3      |9      |
+---+-----------+-----------+--------+--------+-------+-------+

I found a slightly better way than the json approach:

1. Stack the input dataframe value columns A1, A2,B1, B2,.. as rows So the structure would look like id, group, sub, value where sub has the column name like A1, A2, B1, B2 and the value column has the value associated

2. Filter out the rows that have value as null

3. And, now we are able to pivot by the group. Since the null value rows are removed, we wont have the initial issue of the pivot making extra columns

import pyspark.sql.functions as F
data = [
        (1, "Alpha", 1, 2, None, None),
        (1, "AlphaNew", 6, 8, None, None),
        (2, "Alpha", 7, 4, None, None),
        (2, "Beta", None, None, 3, 9),
    ]
columns = ["id", "group", "A1", "A2", "B1", "B2"]
df = spark.createDataFrame(data, columns)

# Value columns that need to be stacked
vcols = df.columns[2:]
expr_str = ', '.join([f"'{i}', {i}" for i in vcols])
expr_str = f"stack({len(vcols)}, {expr_str}) as (sub, value)"

df = df\
.selectExpr("id", "group", expr_str)\
.filter(F.col("value").isNotNull())\
.select("id", F.concat("group", F.lit("_"), "sub").alias("group"), "value")\
.groupBy("id")\
.pivot("group")\
.agg(F.first("value"))

df.show()

Result:
+---+-----------+-----------+--------+--------+-------+-------+
| id|AlphaNew_A1|AlphaNew_A2|Alpha_A1|Alpha_A2|Beta_B1|Beta_B2|
+---+-----------+-----------+--------+--------+-------+-------+
|  1|          6|          8|       1|       2|   null|   null|
|  2|       null|       null|       7|       4|      3|      9|
+---+-----------+-----------+--------+--------+-------+-------+