PySpark如何不使用UDF仅用内置方法基于其他数组列修改数组列值

核心思路是利用PySpark内置的数组展开、条件判断、聚合函数完成逻辑，完全对齐你提供的UDF逻辑：

1. 给每一行生成唯一标识，避免数组展开后无法区分原始行
2. 将四个等长关联数组按位置展开，每行对应数组中单个位置的四个元素
3. 按你给出的判断逻辑对展开后的单条数据做条件修改
4. 按原始行分组、按位置排序后重新聚合为数组格式

完整实现代码 #

from pyspark.sql import functions as F

# 1. 生成唯一行ID，用于后续聚合区分不同原始行
df = df.withColumn("row_id", F.monotonically_increasing_id())

# 2. 按位置展开四个关联数组，保留原始posTags和tokens数组用于后续索引取值
df_exploded = df.select(
    "row_id",
    "tokens",
    "posTags",
    F.posexplode(F.arrays_zip("tokens", "posTags", "dependencies", "labelledDependencies")).alias("pos", "val")
).select(
    "row_id",
    "pos",
    F.col("val.tokens").alias("token"),
    F.col("val.posTags").alias("current_pos"),
    F.col("val.dependencies").alias("dependency"),
    F.col("val.labelledDependencies").alias("current_labelled_dep"),
    "posTags",
    "tokens"
)

# 3. 计算dependency对应token的pos tag，逻辑和Python list.index完全一致（取第一个匹配的位置）
df_exploded = df_exploded.withColumn(
    "dep_token_idx",
    F.array_position(F.col("tokens"), F.col("dependency")) - 1 # array_position返回1-based索引，转成数组用的0-based
).withColumn(
    "dep_pos",
    F.col("posTags")[F.col("dep_token_idx")]
)

# 4. 按规则修改labelled dependency值
df_exploded = df_exploded.withColumn(
    "new_labelled_dep",
    F.when(
        (F.col("current_labelled_dep") == "nsubj") &
        (F.col("current_pos") == "NN") &
        (F.col("dep_pos") == "VB"),
        "dobj"
    ).otherwise(F.col("current_labelled_dep"))
)

# 5. 按原始行分组，按位置排序后聚合回数组格式
result = df_exploded.orderBy("row_id", "pos").groupBy("row_id").agg(
    F.collect_list("token").alias("tokens"),
    F.collect_list("current_pos").alias("posTags"),
    F.collect_list("dependency").alias("dependencies"),
    F.collect_list("new_labelled_dep").alias("labelledDependencies")
).drop("row_id")

效果验证 #

针对你提供的样例数据，balance对应位置的三个判断条件全部满足：

• labelledDependencies值为nsubj
• 对应posTags值为NN
• dependency值为get，对应tokens数组中get的posTags值为VB
  修改后该位置的labelledDependencies会变为dobj，和你UDF的输出结果完全一致。

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