如何在Polars中高效实现NLP预处理的英文缩略词扩展

我在开发NLP模型，当前重点优化预处理流程。最初采用基于正则表达式的自定义函数，通过Polars的map_elements方法实现英文缩略词（例如I'm转换为I am）的展开，但这个自定义函数无法被Polars并行处理，在大规模数据集上性能受限。我尝试过Polars的replace_all方法以及.when.then.otherwise语法，一开始没找到可行的优化方案，希望借助Polars原生API发挥其性能优势。

原实现代码 #

import re
import polars as pl

# 仅展示部分示例缩略词映射
cList = {
    "i'm": "i am",
    "i've": "i have",
    "isn't": "is not"
}

c_re = re.compile("(%s)" % "|".join(cList.keys()))

def expandContractions(text, c_re=c_re):
    def replace(match):
        return cList[match.group(0)]

    return c_re.sub(replace, text)


df = pl.DataFrame({"Text": ["i'm i've, isn't"]})
df["Text"].map_elements(expandContractions)

原实现输出结果 #

shape: (1, 1)
┌─────────────────────┐
│ Text                │
│ ---                 │
│ str                 │
╞═════════════════════╡
│ i am i have, is not │
└─────────────────────┘

性能测试方案 #

为验证性能，我设计了三种测试方案（实际测试中cList包含100+键值对）：

import functools
import polars as pl

cList = {
    "i'm": "i am",
    "i've": "i have",
    "isn't": "is not"
}

def base_case(sr: pl.Series) -> pl.Series:
    c_re = re.compile("(%s)" % "|".join(cList.keys()))
    def expandContractions(text, c_re=c_re):
        def replace(match):
            return cList[match.group(0)]

        return c_re.sub(replace, text)

    sr = sr.map_elements(expandContractions)
    return sr


def loop_case(sr: pl.Series) -> pl.Series:
    for old, new in cList.items():
        sr = sr.str.replace_all(old, new, literal=True)
    return sr


def iter_case(sr: pl.Series) -> pl.Series:
    sr = functools.reduce(
        lambda res, x: res.str.replace_all(x[0], x[1], literal=True),
        cList.items(),
        sr,
    )
    return sr

测试结果 #

经过15轮测试（每轮处理约10000条、单条长度约500字符的样本），平均耗时如下：

Base case: 16.112362766265868
Loop case: 7.028670716285705
Iter case: 7.112465214729309

结论 #

测试结果显示，使用Polars原生replace_all API的两种方案性能是基准方案的两倍以上，最终我选择了循环方案以减少依赖模块。

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