如何加速Polars DataFrame重复过滤与条件列生成操作？

原代码的核心问题是循环执行20万次全表过滤，每次过滤都要扫描百万行数据，加上频繁创建小DataFrame再拼接，导致IO和内存开销极大，速度自然慢。下面是针对性的优化方案：

核心思路：向量化批量匹配 + 单次表扫描 #

避免循环遍历每个条件组合，改用向量化操作一次性判断所有条件组合与每行数据的匹配关系，只扫描原表一次，同时用广播计算替代循环，大幅降低计算开销。

步骤1：预处理条件组合 #

先把所有条件组合转换成Polars DataFrame，同时保留索引方便后续匹配：

import polars as pl
import numpy as np

# 原数据
df = pl.DataFrame({"A": [31,32,73,24,15,26,57,98,79,10], 
                   "B": [11,22,53,44,53,16,27,38,49,10], 
                   "C": [41,12,23,44,25,46,27,48,29,10], 
                   "D": [71,52,13,34,53,36,27,48,39,10], 
                   "E": [81,82,63,24,15,56,47,68,49,10]})

# 条件列表
a12_l = [[13,16], [12,72], [18,22]]
b12_l = [[11,13], [14,55], [22,55]]
c12_l = [[23,76], [13,65], [23,56]]
d12_l = [[21,42], [18,25], [25,35]]

# 转换为条件Series
a_conds = pl.DataFrame(a12_l, schema=["a1", "a2"])
b_conds = pl.DataFrame(b12_l, schema=["b1", "b2"])
c_conds = pl.DataFrame(c12_l, schema=["c1", "c2"])
d_conds = pl.DataFrame(d12_l, schema=["d1", "d2"])

# 生成所有条件组合（笛卡尔积），保留索引
cond_combinations = (
    a_conds.with_row_index("a_idx")
    .join(b_conds.with_row_index("b_idx"), how="cross")
    .join(c_conds.with_row_index("c_idx"), how="cross")
    .join(d_conds.with_row_index("d_idx"), how="cross")
)

# 提取条件数组用于向量化计算
a1_arr = a_conds["a1"].to_numpy()
a2_arr = a_conds["a2"].to_numpy()
b1_arr = b_conds["b1"].to_numpy()
b2_arr = b_conds["b2"].to_numpy()
c1_arr = c_conds["c1"].to_numpy()
c2_arr = c_conds["c2"].to_numpy()
d1_arr = d_conds["d1"].to_numpy()
d2_arr = d_conds["d2"].to_numpy()

步骤2：批量向量化匹配 #

用map_batches处理原DataFrame的每个批次，借助numpy广播一次性计算所有条件组合与行的匹配关系：

# 定义批量处理函数
def batch_match(batch: pl.DataFrame) -> pl.DataFrame:
    batch_np = batch.to_numpy()
    A = batch_np[:, 0]
    B = batch_np[:, 1]
    C = batch_np[:, 2]
    D = batch_np[:, 3]
    
    # 广播计算每个列的区间匹配（行 × 列条件数）
    a_match = (A[:, None] >= a1_arr) & (A[:, None] <= a2_arr)
    b_match = (B[:, None] >= b1_arr) & (B[:, None] <= b2_arr)
    c_match = (C[:, None] >= c1_arr) & (C[:, None] <= c2_arr)
    d_match = (D[:, None] >= d1_arr) & (D[:, None] <= d2_arr)
    
    # 计算四个维度的组合匹配（行 × 所有条件组合数）
    ab_match = a_match[:, :, None] & b_match[:, None, :]
    abc_match = ab_match[:, :, :, None] & c_match[:, None, None, :]
    abcd_match = abc_match[:, :, :, :, None] & d_match[:, None, None, None, :]
    abcd_flat = abcd_match.reshape(len(batch), -1)
    
    # 提取匹配成功的行和条件索引
    row_idx, cond_idx = np.where(abcd_flat)
    
    # 拼接行数据与对应条件
    matched_rows = batch_np[row_idx]
    matched_conds = cond_combinations.to_numpy()[cond_idx]
    combined = np.hstack([matched_rows, matched_conds])
    
    return pl.DataFrame(combined, schema=batch.columns + cond_combinations.columns)

# 执行批量匹配（用lazy模式进一步优化）
result = df.lazy().map_batches(batch_match).collect()

额外优化建议 #

• 去重重复条件：如果条件列表中有重复的区间组合，先对a12_l/b12_l等去重，减少条件组合总数。
• 调整批次大小：通过map_batches的batch_size参数设置合适的批次（比如10万行/批），避免内存溢出。
• 利用Polars并行：确保Polars启用了多线程（默认开启），充分利用CPU多核资源。

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