Polars中字符串差异数计算的高效实现方案咨询

问题背景 #

我们有一个Polars DataFrame，其中每个单元格存储由单个数字组成的字符串序列，需要计算三类差异数：

1. 单个字符串内部所有两两元素的差异数（示例中的pi_1、pi_2）
2. 两个字符串之间所有跨序列元素的两两差异数（示例中的dxy）

原方法使用map_elements结合Python循环实现，但在大数据量下性能不足，需要改用Polars内置向量化操作优化。

核心思路：用统计计数替代暴力循环 #

暴力循环遍历所有两两组合效率极低，我们可以通过数学推导将问题转化为统计计数+数值计算，完全利用Polars的向量化能力：

1. 单个字符串内部差异数（pi） #

• 总两两组合数：$C(n,2) = \frac{n*(n-1)}{2}$，其中n是字符串长度
• 相同元素的两两组合数：对每个数字d，统计其出现次数c，计算$C(c,2) = \frac{c*(c-1)}{2}$，求和所有数字的该值
• 差异数 = 总组合数 - 相同元素组合数之和

2. 跨字符串差异数（dxy） #

• 总跨序列组合数：$len(str1)*len(str2)$
• 相同元素的跨序列组合数：对每个数字d，统计其在str1中的出现次数c1、str2中的出现次数c2，计算$c1*c2$，求和所有数字的该值
• 差异数 = 总组合数 - 相同元素跨组合数之和

向量化实现代码 #

初始化数据 #

import polars as pl

df = pl.DataFrame({"pop_1": ["100","0021"],"pop_2":["11002","0000",]})

计算pi列（单个字符串内部差异） #

def pi_expression(col_name: str) -> pl.Expr:
    # 计算字符串长度
    str_len = pl.col(col_name).str.len_bytes()
    # 总两两组合数
    total_pairs = str_len * (str_len - 1) // 2
    # 计算所有数字的相同元素组合数之和
    same_pairs = sum(
        pl.col(col_name).str.count_matches(str(d)) * (pl.col(col_name).str.count_matches(str(d)) - 1) // 2
        for d in range(10)
    )
    # 差异数 = 总组合数 - 相同组合数
    return (total_pairs - same_pairs).alias(f"pi_{col_name.split('_')[1]}")

# 生成pi_1和pi_2列
df = df.with_columns(pi_expression("pop_1"), pi_expression("pop_2"))

计算dxy列（跨字符串差异） #

# 计算总跨组合数
total_cross_pairs = pl.col("pop_1").str.len_bytes() * pl.col("pop_2").str.len_bytes()
# 计算相同元素的跨组合数之和
same_cross_pairs = sum(
    pl.col("pop_1").str.count_matches(str(d)) * pl.col("pop_2").str.count_matches(str(d))
    for d in range(10)
)
# 生成dxy列
df = df.with_columns((total_cross_pairs - same_cross_pairs).alias("dxy"))

最终结果 #

print(df)

输出：

shape: (2, 5)
┌───────┬───────┬──────┬──────┬─────┐
│ pop_1 ┆ pop_2 ┆ pi_1 ┆ pi_2 ┆ dxy │
│ ---   ┆ ---   ┆ ---  ┆ ---  ┆ --- │
│ str   ┆ str   ┆ i64  ┆ i64  ┆ i64 │
╞═══════╪═══════╪══════╪══════╪═════╡
│ 100   ┆ 11002 ┆ 2    ┆ 8    ┆ 9   │
│ 0021  ┆ 0000  ┆ 5    ┆ 0    ┆ 8   │
└───────┴───────┴──────┴──────┴─────┘

性能优势 #

• 完全基于Polars内置向量化函数，避免了Python循环的开销
• 所有计算在Polars的查询引擎中执行，支持并行化和内存优化
• 大数据量下，速度比原map_elements方法提升数个数量级

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