Polars高效统计大量子字符串出现次数的最优方案

我有一个包含百万行单字符串列的Polars DataFrame（命名为corpus），同时还有一个包含10000个子字符串的列表substrings。单个子字符串的出现行数可以用以下代码统计：

corpus.select(pl.col('contents').str.contains(substrings[0]).sum()).item()

但如果循环调用这个方法处理10000个子字符串，效率极低。我试过把substrings转成Polars DataFrame后通过内连接分组统计，但这个方案内存开销太大，受限于当前内存资源。请问在Polars中有没有更高效、简洁的实现方式？

当前低效实现代码：

import polars as pl

substrings = pl.DataFrame({'substring': ['a', 'b', 'c']})
corpus = pl.DataFrame({'contents': ['aBMMmcICmY', 'ORqkIJCwjV', 'JTQHufYApo', 'SNoqiJxpMY', 'SYbEsasrzt', 'XLinDPSRld', 'iInkOGqBDU', 'vBtykwGOqN', 'ZIpOdkkXBd', 'iUokuiefBS']})

def count_occurrences(substring):
    return corpus.select(pl.col('contents').str.contains(substring).sum()).item()

substrings = substrings.with_columns(pl.col('substring').map_elements(count_occurrences).alias('frequency'))

输出结果：

shape: (3, 2)  
┌───────────┬───────────┐  
│ substring ┆ frequency │  
│ ---       ┆ ---       │  
│ str       ┆ i64       │  
╞═══════════╪═══════════╡  
│ a         ┆ 2         │  
│ b         ┆ 1         │  
│ c         ┆ 1         │  
└───────────┴───────────┘  

可以利用Polars的向量化字符串操作，一次性处理所有子字符串，避免循环开销同时控制内存占用，以下两种方案可供选择：

方法1：批量模式匹配+转置统计 #

将所有子字符串作为批量模式传入str.contains，生成对应每个子串的布尔列，再对这些列求和，最后和原substrings数据关联：

import polars as pl

substrings = pl.DataFrame({'substring': ['a', 'b', 'c']})
corpus = pl.DataFrame({'contents': ['aBMMmcICmY', 'ORqkIJCwjV', 'JTQHufYApo', 'SNoqiJxpMY', 'SYbEsasrzt', 'XLinDPSRld', 'iInkOGqBDU', 'vBtykwGOqN', 'ZIpOdkkXBd', 'iUokuiefBS']})

# 获取子字符串列表
substr_list = substrings['substring'].to_list()

# 批量匹配所有子串并统计出现次数
counts = (
    corpus
    .select(pl.col('contents').str.contains(substr_list, literal=True).alias(substr_list))
    .sum()
    .transpose(include_header=True, column_names=['substring', 'frequency'])
)

# 关联原substrings数据（保留原顺序与结构）
result = substrings.join(counts, on='substring', how='left')
print(result)

方法2：正则分支拼接+提取统计（内存紧张场景适用） #

如果子字符串数量极大，一次性生成大量布尔列内存压力大，可以将子字符串拼接成正则分支模式，结合str.extract_all和explode完成统计：

import polars as pl
import re

substrings = pl.DataFrame({'substring': ['a', 'b', 'c']})
corpus = pl.DataFrame({'contents': ['aBMMmcICmY', 'ORqkIJCwjV', 'JTQHufYApo', 'SNoqiJxpMY', 'SYbEsasrzt', 'XLinDPSRld', 'iInkOGqBDU', 'vBtykwGOqN', 'ZIpOdkkXBd', 'iUokuiefBS']})

# 转义子串中的正则特殊字符，拼接成分支模式
escaped_substrs = [re.escape(s) for s in substrings['substring'].to_list()]
pattern = '|'.join(escaped_substrs)

# 提取匹配子串、去重后统计行数
counts = (
    corpus
    .with_columns(pl.col('contents').str.extract_all(pattern).alias('matches'))
    .filter(pl.col('matches').list.len() > 0)
    .with_columns(pl.col('matches').list.unique())
    .explode('matches')
    .group_by('matches')
    .agg(frequency=pl.count('contents'))
    .rename({'matches': 'substring'})
)

# 关联原数据并填充未匹配的子串频率为0
result = substrings.join(counts, on='substring', how='left').fill_null(0)
print(result)

方案优势 #

• 方法1完全利用Polars向量化操作，彻底避免Python循环，速度远快于逐个子串处理；
• 方法2通过正则分支减少内存占用，适合子串数量极大、内存资源有限的场景，同时保证统计准确性；
• 两种方法均无需内存密集型的内连接操作，内存开销可控。

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