如何在Polars DataFrame中找到首个未按(a,b)排序元素的索引

问题分析 #

你尝试用结构体列实现字典序比较，但Polars不支持直接对结构体列使用</>等比较运算符——虽然结构体列支持排序（内部按字典序处理字段），但元素级比较需要手动拆解字段实现逻辑。以下是线性时间的解决方案，无需调用排序方法。

实现思路 #

直接针对a和b列实现字典序比较：

• 升序要求：当前行的a需≥前一行的a；若a相等，则当前行的b需≥前一行的b，不满足则为不符合排序的元素。
• 降序要求：当前行的a需≤前一行的a；若a相等，则当前行的b需≤前一行的b，不满足则为不符合排序的元素。

通过移位操作获取前一行数值，生成掩码标记不符合条件的行，再找到第一个符合条件的索引即可。

代码实现 #

情况1：期望升序排列

import polars as pl

df = pl.DataFrame({"a": [3,2,1], "b":[3,4,5]})

# 生成掩码：标记当前行是否小于前一行的(a,b)字典序（即不符合升序）
mask = (pl.col("a") < pl.col("a").shift()) | ((pl.col("a") == pl.col("a").shift()) & (pl.col("b") < pl.col("b").shift()))

# 获取首个不符合排序的元素索引
first_invalid_idx = df.with_columns(mask.alias("invalid")).select(pl.col("invalid").arg_true()).item()

print(first_invalid_idx)  # 输出：1

情况2：期望降序排列

# 生成掩码：标记当前行是否大于前一行的(a,b)字典序（即不符合降序）
mask = (pl.col("a") > pl.col("a").shift()) | ((pl.col("a") == pl.col("a").shift()) & (pl.col("b") > pl.col("b").shift()))

first_invalid_idx = df.with_columns(mask.alias("invalid")).select(pl.col("invalid").arg_true()).item()

# 处理所有行符合排序的情况，返回默认值-1
first_invalid_idx = first_invalid_idx if first_invalid_idx is not None else -1
print(first_invalid_idx)  # 输出：-1

关键说明 #

• 该方法时间复杂度为O(n)，仅需一次遍历即可完成判断，适合大数据量场景。
• Polars的arg_true()函数会返回第一个True值的索引，若所有值为False则返回None，可根据需求补充默认值处理。

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