如何动态替换Spark DataFrame结构体字段中的指定值

问题概述 #

使用PySpark的DataFrame.replace方法时，仅能替换顶层字符串字段的指定值，无法处理结构体（Struct）内部的字符串字段。例如示例中my_struct.struct_string的"null"字符串无法被替换为None，且直接指定嵌套字段名会抛出不支持的异常。

问题复现 #

以下代码完整复现问题场景：

from awsglue.context import GlueContext
from pyspark.context import SparkContext
from pyspark.sql.functions import col
from pyspark.sql.types import StringType, StructType, StructField

glueContext = GlueContext(SparkContext.getOrCreate())

data = [
    ("null", {"struct_string": "null"}),
]

schema = StructType([
    StructField("a_string", StringType(), True),
    StructField(
        "my_struct",
        StructType([
           StructField("struct_string", StringType(), True),
        ]),
        True
    )
])

df = spark.createDataFrame(data, schema)

# 尝试替换所有"null"为None，但结构体内部字段未被处理
df = df.replace("null", None)

df_astring = df.filter(col("a_string").isNotNull())
df_struct_string = df.filter(col("my_struct.struct_string").isNotNull())

print("My df_astring")
df_astring.show()
print("My df_struct_string")
df_struct_string.show()

当前执行结果 #

My df_astring
+--------+---------+
|a_string|my_struct|
+--------+---------+
+--------+---------+

My df_struct_string
+--------+---------+
|a_string|my_struct|
+--------+---------+
|    null|   {null}|
+--------+---------+

可见顶层a_string字段的"null"被成功替换为None，但结构体内部的my_struct.struct_string仍保留"null"字符串。

无效尝试 #

当尝试直接指定嵌套字段名时：

df = df.replace("null", None, ["a_string", "my_struct.struct_string"])

会抛出异常：

java.lang.UnsupportedOperationException: Nested field my_struct.struct_string is not supported

动态解决方案 #

通过递归遍历DataFrame的所有字段，针对所有字符串类型字段（包括结构体内部的）统一替换"null"为None，无需手动指定所有字段名：

from pyspark.sql.functions import when, col, struct
from pyspark.sql.types import StructType, StringType

def replace_null_strings(col_name, col_type):
    """递归处理字段，替换字符串类型的"null"为None"""
    if isinstance(col_type, StructType):
        # 递归处理结构体内部字段，生成新的结构体
        struct_fields = [
            replace_null_strings(f.name, f.dataType).alias(f.name)
            for f in col_type.fields
        ]
        return struct(*struct_fields).alias(col_name)
    elif isinstance(col_type, StringType):
        # 对字符串字段替换"null"为None
        return when(col(col_name) == "null", None).otherwise(col(col_name)).alias(col_name)
    else:
        # 非字符串类型字段直接返回
        return col(col_name).alias(col_name)

# 生成所有字段的转换逻辑
transformed_cols = [
    replace_null_strings(field.name, field.dataType)
    for field in df.schema.fields
]

# 应用转换到DataFrame
df = df.select(*transformed_cols)

# 验证结果
df_astring = df.filter(col("a_string").isNotNull())
df_struct_string = df.filter(col("my_struct.struct_string").isNotNull())

print("My df_astring")
df_astring.show()
print("My df_struct_string")
df_struct_string.show()

期望执行结果 #

My df_astring
+--------+---------+
|a_string|my_struct|
+--------+---------+
+--------+---------+

My df_struct_string
+--------+---------+
|a_string|my_struct|
+--------+---------+
+--------+---------+

此时结构体内部的my_struct.struct_string的"null"也被成功替换为None，两个过滤结果均为空。

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