PySpark中高效提取不等长JSON数组（避免Lambda）的优化咨询

需求：在PySpark中高效提取JSON嵌套数组数据，要求避免使用Lambda，且处理col_1/col_2/col_3这类长度不一致的数组（无法用array_zip），现有方案在大数据量下性能不足，需更高效实现。

数据样例 #

data = [
    Row(
        id="1",
        json_data=(
            """{
            "col_1": [{"date": {"time": "2012-10-09T16:54:32+00:00"}, "value": -0.03},
                      {"date": {"time": "2012-10-09T16:54:33+00:00"}, "value": -1.03}],
            "col_2": [{"date": {"time": "2012-10-09T16:54:32+00:00"}, "value": 0.06},
               {"date": {"time": "2012-10-09T16:54:33+00:00"}, "value": 1.06},
               {"date": {"time": "2012-10-09T16:54:34+00:00"}, "value": -0.02}],
            "col_3": [{"date": {"time": "2012-10-09T16:54:31+00:00"}, "value": 10},
               {"date": {"time": "2012-10-09T16:54:32+00:00"}, "value": 7},
            {"date": {"time": "2012-10-09T16:54:33+00:00"}, "value": 8.43}],
            }"""
        )
    )
]

期望输出 #

+--+------------------------+------+------+-----+
|id|time                    |col_1 |col_2 |col_3|
+--+------------------------+------+------+-----+
|1|2012-10-09T16:54:31+00:00|null  |null  |10.0 |
|1|2012-10-09T16:54:32+00:00|-0.03 |0.06  |7.0  |
|1|2012-10-09T16:54:33+00:00|-1.03 |1.06  |8.43 |
|1|2012-10-09T16:54:34+00:00|null  |-0.02 |null |
+--+------------------------+------+------+-----+

当前代码（性能瓶颈点） #

import pyspark.sql.functions as F
from pyspark.sql.types import StructType, StructField, ArrayType, StringType, DoubleType

json_schema = StructType([
    StructField("col_1", ArrayType(StructType([
        StructField("date", StructType([
            StructField("time", StringType())
        ])),
        StructField("value", DoubleType())
    ]))),
    StructField("col_2", ArrayType(StructType([
        StructField("date", StructType([
            StructField("time", StringType())
        ])),
        StructField("value", DoubleType())
    ]))),
    StructField("col_3", ArrayType(StructType([
        StructField("date", StructType([
            StructField("time", StringType())
        ])),
        StructField("value", DoubleType())
    ])))
])

parsed_df = df.withColumn("parsed", F.from_json(df.json_data, json_schema))
    
timestamps_df = parsed_df.select(
    parsed_df.id,
    F.explode(
        F.array_distinct(
            F.concat(
                F.transform("parsed.col_1", lambda x: x.date.time),
                F.transform("parsed.col_2", lambda x: x.date.time),
                F.transform("parsed.col_3", lambda x: x.date.time)
            )
        )
    ).alias("time")
)

def find_value_for_time(array_col, time_col):
    return F.expr(f"""
        CASE 
            WHEN size(filter({array_col}, x -> x.date.time = {time_col})) > 0
            THEN (filter({array_col}, x -> x.date.time = {time_col}))[0].value
            ELSE NULL
        END
    """)

result = timestamps_df.join(parsed_df, "id").select(
    "id",
    "time",
    find_value_for_time("parsed.col_1", "time").alias("col_1"),
    find_value_for_time("parsed.col_2", "time").alias("col_2"),
    find_value_for_time("parsed.col_3", "time").alias("col_3")
)

result = result.orderBy("time")

优化后的高效实现（无Lambda、性能更优） #

核心思路： #

• 对每个列的数组直接展开，生成(id, time, col_name, value)的长表格式，避免全量笛卡尔积join
• 用pivot转置成宽表，自动补全缺失时间的NULL值
• 全程使用PySpark原生函数，避免Lambda和重复的filter计算

import pyspark.sql.functions as F
from pyspark.sql.types import StructType, StructField, ArrayType, StringType, DoubleType

# 1. 定义JSON Schema（和原代码一致）
json_schema = StructType([
    StructField("col_1", ArrayType(StructType([
        StructField("date", StructType([StructField("time", StringType())])),
        StructField("value", DoubleType())
    ]))),
    StructField("col_2", ArrayType(StructType([
        StructField("date", StructType([StructField("time", StringType())])),
        StructField("value", DoubleType())
    ]))),
    StructField("col_3", ArrayType(StructType([
        StructField("date", StructType([StructField("time", StringType())])),
        StructField("value", DoubleType())
    ])))
])

# 2. 解析JSON，拆分成单列展开的长表
parsed_df = df.withColumn("parsed", F.from_json(df.json_data, json_schema))

# 处理col_1
col1_df = parsed_df.select(
    "id",
    F.explode("parsed.col_1").alias("item")
).select(
    "id",
    F.col("item.date.time").alias("time"),
    F.lit("col_1").alias("col_name"),
    F.col("item.value").alias("value")
)

# 处理col_2
col2_df = parsed_df.select(
    "id",
    F.explode("parsed.col_2").alias("item")
).select(
    "id",
    F.col("item.date.time").alias("time"),
    F.lit("col_2").alias("col_name"),
    F.col("item.value").alias("value")
)

# 处理col_3
col3_df = parsed_df.select(
    "id",
    F.explode("parsed.col_3").alias("item")
).select(
    "id",
    F.col("item.date.time").alias("time"),
    F.lit("col_3").alias("col_name"),
    F.col("item.value").alias("value")
)

# 3. 合并所有列的长表，pivot转置成宽表
combined_df = col1_df.unionByName(col2_df).unionByName(col3_df)

result_df = combined_df.groupBy("id", "time").pivot("col_name").agg(F.first("value")).orderBy("time")

result_df.show()

性能优势： #

• 避免了原方案中transform+concat+array_distinct的全量时间收集，以及后续join后重复filter数组的开销
• 展开数组后直接聚合转置，计算逻辑更简洁，Spark的Catalyst优化器能更好地执行计划
• 全程无Lambda，完全使用原生函数，避免UDF或Lambda带来的序列化开销

内容的提问来源于stack exchange，提问作者Panda