如何修改PySpark代码使空loc列也返回指定数组值？

from pyspark.sql import functions as F

L = {'L1': ['us'] }
# df1 = df1.withColumnRenamed("name","OriginalCompanyName")
for key, vals in L.items():
    # regex pattern for extracting vals
    pat = r'\b(%s)\b' % '|'.join(vals)

    # 先判断loc是否为空数组，空数组直接返回["us"]
    # 非空数组则执行正则提取，提取结果为空则返回null，否则返回提取结果
    df1 = df1.withColumn(
        key,
        F.when(
            F.size("loc") == 0,
            F.array(F.lit(vals[0]))
        ).otherwise(
            F.when(
                F.size(F.expr(f"regexp_extract_all(array_join(loc, ' '), '{pat}')")) == 0,
                None
            ).otherwise(
                F.expr(f"regexp_extract_all(array_join(loc, ' '), '{pat}')")
            )
        )
    )

关键修改说明 #

• 新增loc列空数组判断逻辑：用F.size("loc") == 0识别空数组行，直接返回["us"]
• 将原正则提取逻辑嵌套在非空数组分支中，确保只有非空且未匹配到"us"的行才返回null
• 解决了原代码处理空数组时的异常问题，同时满足空数组返回指定值的需求

验证结果 #

针对示例数据，输出如下：

完全符合预期要求。

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