如何在PySpark DataFrame中查找指定列字符串中所有X的出现位置

PySpark 实现代码 #

首先导入依赖，再执行对应处理逻辑即可：

from pyspark.sql import functions as F

# 初始化示例DataFrame
df = spark.createDataFrame([("A", "X-X-------------------------------X--X---XX-X--X-------")],["id", "value"])

# 处理逻辑：兼容所有PySpark版本的实现
result_df = df.select(
    "id",
    F.posexplode(F.split(F.col("value"), "")).alias("pos", "char")
).filter(
    F.col("char") == "X"
).groupBy("id").agg(
    F.collect_list(F.col("pos") + 1).alias("value")
)

# 查看结果
result_df.show(truncate=False)

输出结果

+---+--------------------------------+
|id |value                           |
+---+--------------------------------+
|A  |[1, 3, 35, 38, 42, 43, 45, 48]  |
+---+--------------------------------+

如果你使用的是PySpark 2.4及以上版本，可以用高阶函数实现，性能更优，无需shuffle：

result_df = df.withColumn(
    "char_arr", F.split(F.col("value"), "")
).withColumn(
    "value",
    F.filter(
        F.transform(
            F.sequence(F.lit(0), F.size(F.col("char_arr")) - 1),
            lambda i: F.when(F.col("char_arr")[i] == "X", i + 1)
        ),
        lambda x: x.isNotNull()
    )
).drop("char_arr")

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