如何高效过滤Polars DataFrame中含指定前缀列表的行？

问题场景 #

我有一个包含数百万行的Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    "col1": ["12345", "12", "54467899", "5433523353","0024355"]
})

同时有一个包含数百个前缀的列表：

prefixes = ["123", "544", "55443345"]

需要过滤DataFrame，仅保留col1列值以prefixes中任意前缀开头的行，预期结果如下：

┌────────────┐
│ col1       │
│ ---        │
│ str        │
╞════════════╡
│ 12345      │
│ 54467899   │
└────────────┘

我曾尝试用正则表达式实现，但性能极差：

prefix_regex = '|'.join(f"^{prefix}" for prefix in prefixes)
ldf = pl.scan_parquet('./data.parquet')
ldf = ldf.filter(pl.col("col1").str.contains(prefix_regex)).collect()

使用600个前缀的真实列表时，该操作运行5分钟仍未完成，而不带过滤的collect()仅需几秒。

高效解决方案 #

方法1：str.starts_with + any_horizontal（推荐） #

Polars的str.starts_with支持直接传入单个前缀，结合any_horizontal可以批量判断是否匹配任意前缀。这个操作是向量化的Rust底层实现，性能远优于正则表达式：

import polars as pl

prefixes = ["123", "544", "55443345"]
ldf = pl.scan_parquet('./data.parquet')

# 生成每个前缀的匹配表达式，取任意一个为真
filter_expr = pl.any_horizontal(pl.col("col1").str.starts_with(p) for p in prefixes)
result = ldf.filter(filter_expr).collect()

原理：str.starts_with是专门针对前缀匹配的优化操作，跳过了正则引擎的复杂解析过程，Polars可以对整个列做批量处理。

方法2：字典序区间过滤（性能最优，需前缀满足条件） #

如果前缀可以按字典序排序，我们可以利用字符串的字典序特性，将前缀转换为区间范围，用is_between做过滤。这个操作的性能是所有方法中最高的：

import polars as pl

prefixes = ["123", "544", "55443345"]
# 生成每个前缀对应的区间：所有以p开头的字符串都在[p, p_next)范围内
prefix_ranges = []
for p in prefixes:
    if len(p) == 0:
        continue
    last_char = p[-1]
    next_char = chr(ord(last_char) + 1)
    upper_bound = p[:-1] + next_char
    prefix_ranges.append((p, upper_bound))

# 构建多区间匹配的过滤表达式
filter_expr = pl.any_horizontal(
    pl.col("col1").is_between(start, end, closed='left') for start, end in prefix_ranges
)

ldf = pl.scan_parquet('./data.parquet')
result = ldf.filter(filter_expr).collect()

注意：如果前缀存在包含关系（比如"12"和"123"），需要先去重保留最长前缀，避免区间重叠导致重复匹配。

方法3：前缀树（适合极多前缀场景） #

如果前缀数量超过上千个且存在大量重叠，可以用前缀树（Trie）减少匹配次数。需要借助第三方库pygtrie，但注意这是Python层面的逐行操作，性能略逊于前两种方法：

import polars as pl
import pygtrie

prefixes = ["123", "544", "55443345"]
# 构建前缀树
trie = pygtrie.StringTrie()
for p in prefixes:
    trie[p] = True

def matches_prefix(s: str) -> bool:
    # 检查字符串是否包含前缀树中的任意前缀
    return trie.has_node(s) or any(s.startswith(p) for p in trie.prefixes(s))

ldf = pl.scan_parquet('./data.parquet')
result = ldf.filter(pl.col("col1").map_elements(matches_prefix, return_dtype=pl.Boolean)).collect()

性能对比 #

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