如何使用Polars基于另一DataFrame的行过滤目标DataFrame？

问题描述 #

现有两个Polars DataFrame：

• graph：存储基准区间，列名为START_RANGE和END_RANGE
• search：存储待匹配的目标区间，列名为START和END

需要筛选出graph中所有能包含至少一个search区间的行，即满足：存在某一行search，使得search.START ≥ graph.START_RANGE 且 search.END ≤ graph.END_RANGE。

示例数据 #

import polars as pl

graph = pl.DataFrame(
    {
        "START_RANGE": [1, 10, 20, 30],
        "END_RANGE": [5, 15, 25, 35],
    },
)

search = pl.DataFrame(
    {
        "START": [2, 11, 7],
        "END": [5, 14, 12], 
    },
)

预期输出：

shape: (2, 2)
┌─────────────┬───────────┐
│ START_RANGE ┆ END_RANGE │
│ ---         ┆ ---       │
│ i64         ┆ i64       │
╞═════════════╪═══════════╡
│ 1           ┆ 5         │
│ 10          ┆ 15        │
└─────────────┴───────────┘

解决方案 #

方法1：交叉连接+过滤+去重（直观易懂）

适合数据量较小的场景，逻辑清晰：

result = (
    graph
    # 让每个graph行与每个search行配对
    .join(search, how="cross")
    # 筛选满足包含条件的配对
    .filter(
        pl.col("START") >= pl.col("START_RANGE"),
        pl.col("END") <= pl.col("END_RANGE")
    )
    # 保留graph的列并去重（避免同一graph行匹配多个search行导致重复）
    .select("START_RANGE", "END_RANGE")
    .unique()
)

方法2：向量化判断（高性能）

适合数据量大的场景，避免生成笛卡尔积，减少内存开销：

# 先做初步过滤，排除完全不可能匹配的行
candidate_graph = graph.filter(
    pl.col("START_RANGE") <= search["START"].max(),
    pl.col("END_RANGE") >= search["END"].min()
)

# 对每个候选行，检查是否存在任意search区间被其包含
result = candidate_graph.filter(
    pl.any(
        (pl.col("START_RANGE") <= search["START"]) & (pl.col("END_RANGE") >= search["END"]),
        axis=1
    )
)

说明 #

• 方法1的优势是逻辑简单，容易理解，适合小规模数据；
• 方法2通过提前过滤候选集+向量化判断，大幅减少计算量，在数据量大时性能更优。

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