本教程详细阐述了如何在polars中实现复杂的条件排序，特别适用于根据模型预测结果和置信度对数据进行优先级排序的场景。通过巧妙结合布尔表达式、算术运算和多键排序，polars能够高效地将数据按“高置信度错误预测”、“低置信度错误预测”以及“从低到高置信度的正确预测”进行分组和排序，避免了传统的分拆合并操作，极大提升了数据处理的优雅性和效率。
在数据分析和机器学习模型评估中，我们经常需要对数据进行复杂的条件排序，以便优先关注某些特定模式的样本。例如，在二分类任务中，可能需要优先审查那些模型预测错误但置信度高的样本，以及预测正确但置信度低的样本，以发现模型潜在的改进点。传统的做法可能涉及将数据集拆分成多个子集，分别排序后再合并，但这往往会导致代码冗长且效率低下。Polars作为一个高性能的DataFrame库，提供了强大且富有表达力的API，可以优雅地解决这类复杂排序问题。
Polars多键排序的核心原理
Polars的DataFrame.sort()方法支持传入一个表达式列表作为排序键。这意味着我们可以定义多个排序规则，Polars会按照这些表达式在列表中的顺序依次进行排序。一个关键的特性是，Polars中的布尔值（True和False）在算术运算中会被隐式转换为整数（1和0）。这一特性为实现条件排序提供了极大的灵活性。
构建条件排序表达式
为了实现“高置信度错误预测优先，其次是低置信度正确预测”的排序逻辑，我们将构建三个关键的排序表达式。
1. 数据准备
首先，我们创建一个包含模型预测结果的Polars DataFrame作为示例数据。
import polars as pl
df = pl.DataFrame({
"name": ["Alice", "Bob", "Caroline", "Dutch", "Emily", "Frank", "Gerald", "Henry", "Isabelle", "Jack"],
"truth": [1, 0, 1, 0, 1, 0, 0, 1, 1, 0],
"prediction": [1, 1, 1, 0, 0, 1, 0, 1, 1, 0],
"confidence": [0.343474, 0.298461, 0.420634, 0.125515, 0.772971, 0.646964, 0.833705, 0.837181, 0.790773, 0.144983]
}).with_columns(
(pl.col("truth") == pl.col("prediction")).alias("correct_prediction")
)
print("原始数据：")
print(df)原始数据：
shape: (10, 5) ┌──────────┬───────┬────────────┬────────────┬────────────────────┐ │ name ┆ truth ┆ prediction ┆ confidence ┆ correct_prediction │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ i64 ┆ f64 ┆ bool │ ╞══════════╪═══════╪════════════╪════════════╪════════════════════╡ │ Alice ┆ 1 ┆ 1 ┆ 0.343474 ┆ true │ │ Bob ┆ 0 ┆ 1 ┆ 0.298461 ┆ false │ │ Caroline ┆ 1 ┆ 1 ┆ 0.420634 ┆ true │ │ Dutch ┆ 0 ┆ 0 ┆ 0.125515 ┆ true │ │ Emily ┆ 1 ┆ 0 ┆ 0.772971 ┆ false │ │ Frank ┆ 0 ┆ 1 ┆ 0.646964 ┆ false │ │ Gerald ┆ 0 ┆ 0 ┆ 0.833705 ┆ true │ │ Henry ┆ 1 ┆ 1 ┆ 0.837181 ┆ true │ │ Isabelle ┆ 1 ┆ 1 ┆ 0.790773 ┆ true │ │ Jack ┆ 0 ┆ 0 ┆ 0.144983 ┆ true │ └──────────┴───────┴────────────┴────────────┴────────────────────┘
2. 第一排序键：区分错误与正确预测
我们的首要目标是将错误预测排在正确预测之前。我们可以利用correct_prediction列（布尔类型）作为第一个排序键。在Polars中，False（对应整数0）会排在True（对应整数1）之前。
# 定义第一个排序键：判断是否是正确预测
# False (0) 会排在 True (1) 之前，因此错误预测会优先出现
first_key = pl.col('truth').eq(pl.col('prediction'))
# 为了方便后续使用，我们可以将其赋值给一个变量
# good_pred = pl.col('truth').eq(pl.col('prediction'))这里我们直接在sort方法中使用表达式，并利用Python的赋值表达式:=将结果赋值给good_pred，以便在后续表达式中复用。
3. 第二排序键：处理错误预测的置信度（降序）
在错误预测的组内，我们希望置信度最高的错误预测排在最前面。这意味着我们需要对置信度进行降序排序。对于正确预测的样本，这个排序键不应该影响它们的顺序，而是让下一个排序键来决定。
我们可以通过一个巧妙的数学表达式来实现这一点： (good_pred - 1) * pl.col('confidence')
- 当good_pred为False（即0，表示错误预测）时： 表达式变为 (0 - 1) * confidence = -confidence。 对-confidence进行升序排序，实际上就是对confidence进行降序排序。这正是我们对错误预测所期望的。
- 当good_pred为True（即1，表示正确预测）时： 表达式变为 (1 - 1) * confidence = 0 * confidence = 0。 这意味着所有正确预测的样本在这个排序键上的值都为0，它们之间保持相对顺序，等待下一个排序键来决定。
4. 第三排序键：处理正确预测的置信度（升序）
对于正确预测的组内，我们希望置信度最低的预测排在最前面，即对置信度进行升序排序。由于第二个排序键对正确预测的样本都返回了0，因此它们将由第三个排序键来决定顺序。
# 定义第三个排序键：对置信度进行升序排序
# 这将应用于所有样本，但只有在前面键值相同的情况下才会生效
third_key = pl.col('confidence')完整示例代码
将以上三个排序键组合到DataFrame.sort()方法中：
sorted_df = df.sort([
(good_pred:=pl.col('truth').eq(pl.col('prediction'))), # 1. 错误预测 (False/0) 优先于正确预测 (True/1)
(good_pred-1)*pl.col('confidence'), # 2. 错误预测内部，按置信度降序
pl.col('confidence') # 3. 正确预测内部，按置信度升序
])
print("\n条件排序后的数据：")
print(sorted_df)条件排序后的数据：
shape: (10, 5) ┌──────────┬───────┬────────────┬────────────┬────────────────────┐ │ name ┆ truth ┆ prediction ┆ confidence ┆ correct_prediction │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ i64 ┆ f64 ┆ bool │ ╞══════════╪═══════╪════════════╪════════════╪════════════════════╡ │ Emily ┆ 1 ┆ 0 ┆ 0.772971 ┆ false │ │ Frank ┆ 0 ┆ 1 ┆ 0.646964 ┆ false │ │ Bob ┆ 0 ┆ 1 ┆ 0.298461 ┆ false │ │ Dutch ┆ 0 ┆ 0 ┆ 0.125515 ┆ true │ │ Jack ┆ 0 ┆ 0 ┆ 0.144983 ┆ true │ │ Alice ┆ 1 ┆ 1 ┆ 0.343474 ┆ true │ │ Caroline ┆ 1 ┆ 1 ┆ 0.420634 ┆ true │ │ Isabelle ┆ 1 ┆ 1 ┆ 0.790773 ┆ true │ │ Gerald ┆ 0 ┆ 0 ┆ 0.833705 ┆ true │ │ Henry ┆ 1 ┆ 1 ┆ 0.837181 ┆ true │ └──────────┴───────┴────────────┴────────────┴────────────────────┘
结果分析与验证
从输出结果可以看出，排序完全符合我们的预期：
- 错误预测优先： Emily、Frank、Bob 都属于 correct_prediction = False，它们排在所有正确预测之前。
