PySpark：查找数组列中首个正数的索引（非UDF方案）

针对你的需求，利用PySpark内置函数即可高效实现，避免UDF的性能瓶颈，以下是两种可行方案：

方案一：结合transform+array_position #

利用transform将数组转换为布尔标记数组（标记元素是否为正数），再通过array_position定位第一个正数的位置，最后转换为0-based索引：

from pyspark.sql import functions as F

# 假设你的DataFrame名为df
df = df.withColumn(
    "first_positive_element_index",
    F.when(
        # 判断数组中是否存在正数
        F.exists(F.col("arr"), lambda x: x != -1),
        # array_position返回1-based位置，减1转为0-based索引
        F.array_position(F.transform(F.col("arr"), lambda x: x != -1), True) - 1
    )
)

逻辑说明： #

1. transform(F.col("arr"), lambda x: x != -1)：将原数组转换为布尔数组，正数对应True，-1对应False
2. array_position(..., True)：找到布尔数组中第一个True的位置（从1开始计数）
3. exists(...)：判断数组是否存在正数，不存在则返回null
4. 最后减1得到符合要求的0-based索引

方案二：使用aggregate遍历数组 #

通过aggregate函数遍历数组，直接记录第一个正数的索引：

from pyspark.sql import functions as F

df = df.withColumn(
    "first_positive_element_index",
    F.aggregate(
        # 将数组元素与索引配对，生成(是否为正数, 索引)的结构体数组
        F.transform(F.col("arr"), lambda x, i: F.struct((x != -1).alias("is_pos"), i.alias("idx"))),
        # 初始值设为null，表示尚未找到正数
        F.lit(None).cast("integer"),
        # 遍历逻辑：如果还没找到正数且当前元素是正数，就更新索引，否则保留原有值
        lambda acc, elem: F.when(acc.isNull() & elem.is_pos, elem.idx).otherwise(acc)
    )
)

逻辑说明： #

1. transform(..., lambda x, i: struct(...))：给每个数组元素绑定对应的索引，并标记是否为正数
2. aggregate初始化值为null，遍历过程中仅在首次遇到正数时更新索引
3. 遍历结束后，acc即为第一个正数的索引，无正数则保持null

两种方案均基于PySpark内置的分布式优化函数，无需Python UDF的序列化开销，在大数据量场景下性能远优于UDF实现。

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