PySpark：移除collect_list生成的_source字段外层方括号方法

我需要把DataFrame以JSON格式导出到S3，但要求生成的_source字段仅为键值对结构，不要外层的方括号。

原代码 #

def export_data(example):
    example = example \
            .groupby(col("_id"), col("_index")) \
            .agg(F.collect_list(F.struct(
                col("colA1"),
                col("colA2"),
                col("colA3"),
                col("colA4"),
                col("colA5"),
                col("colA6"))).alias("_source"))
    return example.select(collect_list(F.struct(col("_id"), col("_source"))).alias("hits"))

def main():
    example_export = export_data(example)

    example_export.coalesce(1) \
        .write.format('json') \
        .mode("overwrite") \
        .save(
        "s3://aws-example/output/EXPORTER")

当前输出结果 #

{
    "hits": [
        {
            "_id": "22b9d653-c07a-46da-a929-42337f0f020d",
            "_source": [
                {
                    "colA1": "202211",
                    "colA2": "A",
                    "colA3": "020160",
                    "colA4": "MANAUS",
                    "colA5": "18",
                    "colA6": ""
                }
            ]
        },
        {
            "_id": "6fe3b950-3b09-4f11-b4d7-5e8c9aab8f3e",
            "_source": [
                {
                    "colA1": "202211",
                    "colA2": "A",
                    "colA3": "21345",
                    "colA4": "RORAIMA",
                    "colA5": "10",
                    "colA6": ""
                }
            ]
        }
    ]
}

期望输出结果 #

{
    "hits": [
        {
            "_id": "22b9d653-c07a-46da-a929-42337f0f020d",
            "_source": {
                "colA1": "202211",
                "colA2": "A",
                "colA3": "020160",
                "colA4": "MANAUS",
                "colA5": "18",
                "colA6": ""
            }
        },
        {
            "_id": "6fe3b950-3b09-4f11-b4d7-5e8c9aab8f3e",
            "_source": {
                "colA1": "202211",
                "colA2": "A",
                "colA3": "21345",
                "colA4": "RORAIMA",
                "colA5": "10",
                "colA6": ""
            }
        }
    ]
}

解决方案 #

问题根源是agg阶段使用了F.collect_list，它会把结构体包装成数组。根据你的分组逻辑选择对应方案：

• 如果每个_id+_index分组对应的colA1-colA6只有一条数据，直接用F.struct替代F.collect_list(F.struct(...))即可。
• 如果分组内有多条数据，但你只需要保留第一条，就用F.first(F.struct(...))来替代。

修改后的代码 #

from pyspark.sql import functions as F
from pyspark.sql.functions import col, collect_list

def export_data(example):
    # 方案1：每个分组只有一条数据时使用
    example = example \
            .groupby(col("_id"), col("_index")) \
            .agg(F.struct(
                col("colA1"),
                col("colA2"),
                col("colA3"),
                col("colA4"),
                col("colA5"),
                col("colA6")).alias("_source"))
    # 方案2：分组内有多条数据，取第一条时使用
    # example = example \
    #         .groupby(col("_id"), col("_index")) \
    #         .agg(F.first(F.struct(
    #             col("colA1"),
    #             col("colA2"),
    #             col("colA3"),
    #             col("colA4"),
    #             col("colA5"),
    #             col("colA6"))).alias("_source"))
    return example.select(collect_list(F.struct(col("_id"), col("_source"))).alias("hits"))

def main():
    example_export = export_data(example)

    example_export.coalesce(1) \
        .write.format('json') \
        .mode("overwrite") \
        .save("s3://aws-example/output/EXPORTER")

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