PySpark DataFrame固定行数展开数组列并补全缺失值

核心思路 #

先给每行的数组列补全到固定长度（比如示例的5，实际场景的14），缺失的结构体用空值填充，再通过explode展开即可得到固定行数的结果。

步骤实现（以最大元素数5为例） #

假设原DataFrame的结构和示例数据如下：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, ArrayType

spark = SparkSession.builder.appName("FillArray").getOrCreate()

# 定义结构体Schema
attr_schema = StructType([
    StructField("version", StringType()),
    StructField("approval_status", StringType()),
    StructField("interest_rate", DoubleType())
])

# 示例数据
data = [
    (1, [{"version": "v1", "approval_status": "approved", "interest_rate": 3.5}, {"version": "v2", "approval_status": "pending", "interest_rate": 4.0}]),
    (2, [{"version": "v1", "approval_status": "rejected", "interest_rate": None}]),
    (3, [])
]

df = spark.createDataFrame(data, schema=["id", "attributes"])

1. 计算需要补充的空结构体数量 #

用size函数获取当前数组长度，计算与固定长度的差值，得到需要补充的空结构体数量：

from pyspark.sql.functions import size, lit, array, struct, when, expr

max_elements = 5  # 实际场景替换为14
df = df.withColumn("missing_count", when(size("attributes") < max_elements, max_elements - size("attributes")).otherwise(0))

2. 生成空结构体数组并拼接 #

先构造一个所有属性为null的空结构体，再用array_repeat生成对应数量的空结构体数组，最后和原数组拼接：

# 创建空结构体（匹配原结构体的所有属性）
empty_attr = struct(
    lit(None).alias("version"),
    lit(None).alias("approval_status"),
    lit(None).alias("interest_rate")
)

# 拼接原数组和空结构体数组
df = df.withColumn("filled_attributes", expr("concat(attributes, array_repeat({}, missing_count))".format(empty_attr)))

3. 展开数组并清理临时列 #

用explode展开填充后的数组，再拆分结构体属性、删除临时列：

df_final = df.select("id", explode("filled_attributes").alias("attributes")) \
             .select("id", "attributes.version", "attributes.approval_status", "attributes.interest_rate")

4. 验证结果 #

df_final中每个id都会对应5行数据，缺失位置的所有属性均为null。

适配14个元素的场景 #

只需将max_elements改为14，同时确保空结构体包含全部7个属性即可，示例代码如下：

# 包含7个属性的空结构体
empty_attr = struct(
    lit(None).alias("attr1"),
    lit(None).alias("attr2"),
    lit(None).alias("attr3"),
    lit(None).alias("attr4"),
    lit(None).alias("attr5"),
    lit(None).alias("attr6"),
    lit(None).alias("attr7")
)

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