利用Polars库加速元组序列间的最小Levenshtein距离计算

我的pandas DataFrame的sequence列存储着元组形式的序列，示例数据如下：

id       sequence
33268    [(59, 2), (91, 2), (112, 2), (126, 2), (0, 3),...
44360    [(137, 0), (208, 0), (251, 0), (0, 1), (46, 1)...
51943    [(117, 0), (153, 1), (190, 1), (0, 7), (218, 7...
53888    [(0, 0), (18, 0), (58, 0), (69, 0), (86, 0), (...
38505    [(0, 1), (96, 1), (131, 1), (188, 1), (229, 1)...

我需要计算每个序列与列表l=[...]中同形式序列的最小Levenshtein距离。
已实现以下Levenshtein距离计算函数：

def levenshtein(s, t):
    if s == t: return 0
    elif len(s) == 0: return len(t)
    elif len(t) == 0: return len(s)
    v0 = [None] * (len(t) + 1)
    v1 = [None] * (len(t) + 1)
    for i in range(len(v0)):
        v0[i] = i
    for i in range(len(s)):
        v1[0] = i + 1
        for j in range(len(t)):
            cost = 0 if s[i] == t[j] else 1
            v1[j + 1] = min(v1[j] + 1, v0[j + 1] + 1, v0[j] + cost)
        for j in range(len(v0)):
            v0[j] = v1[j]

    return v1[len(t)]

目前用pandas.DataFrame.apply完成计算，但速度极慢，请问能否用Polars库提升计算速度？

当然可以用Polars提升计算速度，核心是利用Polars的向量化操作和并行处理能力，同时优化计算逻辑，具体步骤如下：

1. 转换数据到Polars DataFrame #

先将pandas DataFrame转为Polars格式，操作简单：

import polars as pl

pl_df = pl.from_pandas(pandas_df)

2. 结合map_elements优化计算 #

Polars的map_elements比pandas的apply效率更高，默认支持多线程并行。建议加上剪枝逻辑减少不必要的计算：

# 预处理目标列表，提前存储序列长度，避免重复计算
preprocessed_l = [(seq, len(seq)) for seq in l]

def min_levenshtein(row_seq):
    min_dist = float('inf')
    row_len = len(row_seq)
    for t, t_len in preprocessed_l:
        # 剪枝：长度差大于当前最小距离时直接跳过
        if abs(row_len - t_len) >= min_dist:
            continue
        dist = levenshtein(row_seq, t)
        if dist < min_dist:
            min_dist = dist
            # 找到0距离直接返回，无需继续计算
            if min_dist == 0:
                break
    return min_dist

# 应用到sequence列，生成最小距离列
pl_df = pl_df.with_columns(
    pl.col("sequence").map_elements(min_levenshtein, return_dtype=pl.Int32).alias("min_lev_dist")
)

3. 批量计算方案（适合小规模目标列表） #

如果列表l规模不大，可以通过笛卡尔积批量计算距离，再按id取最小值，更好利用Polars的向量化优势：

# 将目标列表转为Polars DataFrame
target_df = pl.DataFrame({"target_seq": l})

# 笛卡尔积关联原数据与目标序列
cross_df = pl_df.join(target_df, how="cross")

# 批量计算每对序列的Levenshtein距离
cross_df = cross_df.with_columns(
    pl.struct(["sequence", "target_seq"])
      .map_elements(lambda x: levenshtein(x["sequence"], x["target_seq"]), return_dtype=pl.Int32)
      .alias("lev_dist")
)

# 按id分组取最小距离
result_df = cross_df.group_by("id").agg(pl.col("lev_dist").min().alias("min_lev_dist"))

4. 极致提速：结合Numba编译Levenshtein函数 #

Python循环是性能瓶颈，用Numba编译Levenshtein函数可以大幅降低开销：

from numba import jit

@jit(nopython=True)
def numba_levenshtein(s, t):
    # Numba对元组支持更好，调用前需将列表转成tuple
    if s == t:
        return 0
    len_s = len(s)
    len_t = len(t)
    if len_s == 0:
        return len_t
    if len_t == 0:
        return len_s
    v0 = list(range(len_t + 1))
    v1 = [0]*(len_t + 1)
    for i in range(len_s):
        v1[0] = i + 1
        for j in range(len_t):
            cost = 0 if s[i] == t[j] else 1
            v1[j+1] = min(v1[j] + 1, v0[j+1] + 1, v0[j] + cost)
        v0, v1 = v1, v0
    return v0[len_t]

# 调整min_levenshtein函数使用编译后的版本
def min_levenshtein(row_seq):
    min_dist = float('inf')
    row_seq_tuple = tuple(row_seq)
    for t in l:
        t_tuple = tuple(t)
        dist = numba_levenshtein(row_seq_tuple, t_tuple)
        if dist < min_dist:
            min_dist = dist
            if min_dist == 0:
                break
    return min_dist

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