Spark 3.3用高阶函数提取数组列，生成库/模式/表数组列

有一个Spark DataFrame的ArrayType列，列中元素格式示例为['db1.schema1.table1','schema5.table3','table4']。目标是生成三个ArrayType列：

• dbs：提取每个元素的数据库名（示例值[db1]）
• schemas：提取每个元素的模式名（示例值[schema1,schema5]）
• tables：提取每个元素的表名（示例值['table1','table3','table4']）

此前用Python UDF实现，但运行慢、内存效率低。当前环境为Spark 3.3，曾尝试过explode+窗口函数的方案，但存在shuffle风险，询问是否可以用高阶函数解决该问题。

之前尝试的代码（存在shuffle风险）：

from pyspark.sql.types import *
from pyspark.sql.functions import *

cSchema = StructType([StructField("WordList", ArrayType(StringType()))])
test_list = [[['db1.s1.t1']], [['t3','d1.s1.t1','s2.t2']]]
df = spark.createDataFrame(test_list,schema=cSchema)
df = df.withColumn("random",expr("uuid()"))

df=df.select('*',explode("WordList").alias("x"))
df=df.withColumn('x_split',split(col('x'), "\\."))
df=df.withColumn("size", size(col('x_split')))

df = df.withColumn("table", element_at(col('x_split'),col('size')))
df = df.withColumn("database", when(col('size')==3, element_at(col('x_split'),1)).otherwise(lit('na')))
df = df.withColumn("schema", when(col('size')>1, element_at(col('x_split'),col('size')-1)).otherwise(lit('na')))

# 后续按uuid分组collect_set，存在shuffle

完全不需要explode或窗口函数，直接用Spark 3.3支持的数组高阶函数处理，全程无shuffle，性能更优：

from pyspark.sql import functions as F

# 测试数据
test_data = [
    (['db1.schema1.table1','schema5.table3','table4'],),
    (['t3','d1.s1.t1','s2.t2'],)
]
df = spark.createDataFrame(test_data, schema=["WordList"])

# 遍历数组每个元素，生成包含table、schema、db的结构体数组
df = df.withColumn(
    "processed",
    F.transform(
        F.col("WordList"),
        lambda x: F.struct(
            # 提取表名：分割后取最后一个元素
            F.element_at(F.split(x, "\\."), -1).alias("table"),
            # 提取schema：分割后长度>=2时取倒数第二个，否则为null
            F.when(F.size(F.split(x, "\\.")) >= 2, F.element_at(F.split(x, "\\."), -2)).alias("schema"),
            # 提取db：分割后长度==3时取第一个，否则为null
            F.when(F.size(F.split(x, "\\.")) == 3, F.element_at(F.split(x, "\\."), 1)).alias("db")
        )
    )
)

# 从结构体数组中提取目标列，过滤null值
df = df.withColumn("dbs", F.filter(F.transform(F.col("processed"), lambda x: x.db), lambda x: x.isNotNull()))
df = df.withColumn("schemas", F.filter(F.transform(F.col("processed"), lambda x: x.schema), lambda x: x.isNotNull()))
df = df.withColumn("tables", F.transform(F.col("processed"), lambda x: x.table))

# 清理中间列
df = df.drop("processed")

df.show(truncate=False)

代码说明 #

1. transform遍历数组：用transform对WordList中的每个元素逐个处理，生成包含table、schema、db的结构体数组processed，全程在单条记录内处理，无shuffle。
2. 提取目标数组：
   • tables：直接从结构体数组中提取所有table字段
   • dbs/schemas：提取对应字段后用filter去掉null值（对应原元素中没有db/schema的情况）
3. 性能优势：所有操作都是单分区内的行级操作，完全避免了explode后的shuffle和窗口函数的分组开销，比原方案和UDF效率高得多。

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