如何在Polars中通过条件正则表达式为字符串列数据打标签

问题核心 #

多次链式调用str.replace会导致重复全列正则扫描，上百个正则对应上百次遍历，处理百万行数据时性能会急剧下降。

最优解决方案：向量化条件判断（when/then/otherwise） #

利用Polars的向量化条件表达式，在一次遍历中完成所有正则匹配判断，同时可灵活控制匹配优先级（先匹配的规则优先生效）。

步骤1：定义模式-标签映射（便于管理上百种规则）

按匹配优先级排序规则（比如先匹配邮箱，避免被数字类规则误判）：

pattern_label_pairs = [
    # 更严谨的邮箱正则（覆盖带特殊字符的邮箱）
    (r"^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}$", "EMAIL"),
    (r"^\d{3}-\d{3}-\d{4}$", "PHONE"),
    (r"^\d{5}$", "ZIP CODE"),
    (r"^[A-Z]{2}$", "STATE"),
    # 可继续添加更多模式...
]

步骤2：构建向量化标签表达式

通过链式when/then/otherwise实现一次遍历匹配：

import polars as pl

df = pl.DataFrame(
    {
        "data": [
            "123-123-1234",
            "someone@email.com",
            "345-345-3456",
            "456-456-4567",
            "12345",
            "charlie.brown@peanuts.org",
            "345-345-3456",
            "CA",
            "UT",
        ]
    }
)

# 初始化默认标签（未匹配任何规则时的标签）
label_expr = pl.lit("UNKNOWN")

# 按优先级依次添加匹配规则
for pattern, label in pattern_label_pairs:
    label_expr = pl.when(pl.col("data").str.contains(pattern)).then(label).otherwise(label_expr)

# 生成结果
result_df = df.with_columns(label=label_expr)
print(result_df)

输出结果

shape: (9, 2)
┌───────────────────────────┬──────────┐
│ data                      ┆ label    │
│ ---                       ┆ ---      │
│ str                       ┆ str      │
╞═══════════════════════════╪══════════╡
│ 123-123-1234              ┆ PHONE    │
│ someone@email.com         ┆ EMAIL    │
│ 345-345-3456              ┆ PHONE    │
│ 456-456-4567              ┆ PHONE    │
│ 12345                     ┆ ZIP CODE │
│ charlie.brown@peanuts.org ┆ EMAIL    │
│ 345-345-3456              ┆ PHONE    │
│ CA                        ┆ STATE    │
│ UT                        ┆ STATE    │
└───────────────────────────┴──────────┘

性能优势说明 #

• 时间复杂度优化：原方法为O(n*k)（n为行数，k为正则数量），新方法为O(n)（一次遍历完成所有匹配检查）。
• 向量化引擎加持：Polars的when/then/otherwise是底层优化的向量化操作，比Python层面循环或多次replace快数个数量级，适配百万级数据场景。
• 规则可维护性：模式-标签的列表形式便于添加、删除或调整规则优先级，管理上百种模式更清晰。

额外优化建议 #

1. 正则精准匹配：用^和$限定字符串首尾，避免部分匹配导致错误标签（比如防止包含5位数字的长字符串被误判为邮编）。
2. 极端复杂正则的适配：如果正则逻辑极复杂，可提前预编译正则，结合map_batches使用（仅在向量化操作无法满足时采用，因为会脱离Polars向量化引擎）：

   import re
   
   compiled_patterns = [(re.compile(p), l) for p, l in pattern_label_pairs]
   
   def assign_label(batch):
       labels = []
       for s in batch["data"]:
           for pat, label in compiled_patterns:
               if pat.match(s):
                   labels.append(label)
                   break
           else:
               labels.append("UNKNOWN")
       return pl.DataFrame({"label": labels})
   
   result_df = df.with_columns(pl.col("data").map_batches(assign_label))
   

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