获取pyspark中数组列中所有True元素的索引

from pyspark.sql.functions import arrays_zip, expr

df = df.withColumn("indexed_array", arrays_zip(df.array_column, expr("sequence(size(array_column))")))

df = df.withColumn("filtered_array", expr("transform(indexed_array, x -> IF(x[0], x[1], NULL))"))

df = df.withColumn("indexes", expr("transform(filtered_array, x -> array_position(array_column, x))"))

from pyspark.sql import SparkSession
from pyspark.sql.functions import arrays_zip, expr

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例数据
data = [("A", [True, False, True]),
        ("B", [False, True, True]),
        ("C", [False, False, False])]

df = spark.createDataFrame(data, ["id", "array_column"])

# 获取所有True元素的索引
df = df.withColumn("indexed_array", arrays_zip(df.array_column, expr("sequence(size(array_column))")))
df = df.withColumn("filtered_array", expr("transform(indexed_array, x -> IF(x[0], x[1], NULL))"))
df = df.withColumn("indexes", expr("transform(filtered_array, x -> array_position(array_column, x))"))

df.show(truncate=False)

+---+--------------+-----------------+-----------+
|id |array_column  |indexed_array    |indexes    |
+---+--------------+-----------------+-----------+
|A  |[true, false, true]|[[true, 0], [false, 1], [true, 2]]|[0, 2]     |
|B  |[false, true, true]|[[false, 0], [true, 1], [true, 2]]|[1, 2]     |
|C  |[false, false, false]|[[false, 0], [false, 1], [false, 2]]|[]         |
+---+--------------+-----------------+-----------+