如何在PySpark DataFrame中获取首个非零元素对应的索引

假设你预先添加的索引列名为row_index，待检查的数值列名为target_col，核心实现逻辑为过滤非零行后按索引升序取首条记录即可，具体实现如下：

• 依赖导入

from pyspark.sql import functions as F

• 核心实现代码

# 1. 过滤出所有目标列非零的行
non_zero_df = df.filter(F.col("target_col") != 0)

# 2. 按索引升序排序后取第一条记录，加limit(1)优化大表查询性能
first_non_zero_row = non_zero_df.orderBy(F.col("row_index").asc()).limit(1).first()

# 3. 提取索引值
if first_non_zero_row:
    first_non_zero_index = first_non_zero_row["row_index"]
else:
    # 所有元素均为0的场景返回空值，可根据需求调整默认返回值
    first_non_zero_index = None

如果你需要判断整行非零而非单个列，修改filter条件即可，例如需要同时检查col1、col2两列的话，过滤条件改为filter((F.col("col1") != 0) | (F.col("col2") != 0))即可匹配任意一列非零的行。

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