如何在PySpark DataFrame列中定位所有指定字符的出现位置？

问题背景 #

原始PySpark DataFrame：

需要新增一列，记录列A中所有'1'的索引位置（从0开始计数），期望结果如下：

尝试使用F.locate()函数仅能获取第一个'1'的位置，无法满足需求。

解决方案 #

方法1：使用PySpark内置函数（Spark 3.1+推荐） #

利用Spark内置的数组操作函数实现，无需自定义UDF，性能更优：

from pyspark.sql import functions as F

# 生成字符与对应索引的结构体数组
df_with_index = dfa.withColumn("char_index_pair",
    F.zip_with(
        F.split(F.col("A"), ""),  # 将字符串拆分为单个字符的数组
        F.sequence(F.lit(0), F.length(F.col("A")) - 1),  # 生成从0开始的索引数组
        lambda char, idx: F.struct(char.alias("char"), idx.alias("index"))
    )
)

# 过滤出字符为'1'的元素，提取索引并拼接为字符串
df_result = df_with_index.withColumn("出现位置",
    F.concat_ws(",",
        F.transform(
            F.filter(F.col("char_index_pair"), lambda x: x.char == "1"),
            lambda x: x.index
        )
    )
).drop("char_index_pair")

# 查看结果
df_result.show()

方法2：自定义UDF（兼容低版本Spark） #

如果你的Spark版本低于3.1，不支持zip_with等函数，可以用UDF实现：

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

def collect_1_positions(s):
    # 遍历字符串，收集所有'1'的索引并转为逗号分隔的字符串
    pos_list = [str(i) for i, c in enumerate(s) if c == '1']
    return ",".join(pos_list)

# 注册UDF
pos_udf = F.udf(collect_1_positions, StringType())

# 应用UDF生成目标列
df_result = dfa.withColumn("出现位置", pos_udf(F.col("A")))

# 查看结果
df_result.show()

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