如何高效简洁地在Polars DataFrame中组合多字符串操作

问题背景 #

处理Polars DataFrame时，需要对engine列执行两步操作：先用str.replace()统一字符串格式，再提取多个字段生成新列。现有两种实现方式，希望找到更高效简洁、同时保持代码整洁的写法。

原始实现 #

import polars as pl

df = pl.DataFrame(
    {
        "engine": ["172.0HP 1.6L 4 Cylinder Engine Gasoline Fuel",
                    "3.0 Liter Twin Turbo",
                    "429.0HP 5.0L 8 Cylinder Engine Gasoline Fuel"],
    }
)

(
    df
    .with_columns(
        pl.col("engine").str.replace(r'\sLiter', "L")
    )
    .with_columns(
        pl.col("engine").str.extract(r'(\d+)\.\d+HP',1).alias("HP"),
        pl.col("engine").str.extract(r'(\d+\.\d+)L',1).alias("Displacement"),
        pl.col("engine").str.extract(r'(\d+)\sCylinder',1).alias("Cylinder"),
    )
)

函数式尝试 #

def get_engine() -> pl.Expr:
    return (
        pl.col("engine").str.extract(r'(\d+)\.\d+HP',1).alias("HP"),
        pl.col("engine").str.extract(r'(\d+\.\d+)L',1).alias("Displacement"),
        pl.col("engine").str.extract(r'(\d+)\sCylinder',1).alias("Cylinder"),
        pl.col("engine").str.contains("Electric").alias("Electric")
    )


(
    df
    .with_columns(
        pl.col("engine").str.replace(r'\sLiter', "L")
    )
    .with_columns(
        get_engine()
    )
)

优化方案 #

1. 合并with_columns调用，简化链式逻辑 #

把格式统一和字段提取放在同一个with_columns调用中，减少方法调用次数，同时优化正则兼容原始格式，避免额外列操作：

(
    df
    .with_columns(
        pl.col("engine").str.extract(r'(\d+)\.\d+HP', 1).alias("HP"),
        pl.col("engine").str.extract(r'(\d+\.\d+)(?:L| Liter)', 1).alias("Displacement"),
        pl.col("engine").str.extract(r'(\d+)\sCylinder', 1).alias("Cylinder"),
        pl.col("engine").str.contains("Electric").alias("Electric"),
        # 可选：生成统一格式的engine列
        pl.col("engine").str.replace(r'\sLiter', "L").alias("engine_clean")
    )
)

2. 封装完整处理逻辑，提升复用性 #

把格式清理和字段提取的所有逻辑封装到一个函数中，后续调用只需一行代码，维护修改更方便：

def process_engine() -> list[pl.Expr]:
    cleaned = pl.col("engine").str.replace(r'\sLiter', "L").alias("engine_clean")
    return [
        cleaned,
        pl.col("engine_clean").str.extract(r'(\d+)\.\d+HP', 1).alias("HP"),
        pl.col("engine_clean").str.extract(r'(\d+\.\d+)L', 1).alias("Displacement"),
        pl.col("engine_clean").str.extract(r'(\d+)\sCylinder', 1).alias("Cylinder"),
        pl.col("engine_clean").str.contains("Electric").alias("Electric")
    ]

# 使用示例
(
    df
    .with_columns(process_engine())
    # 可选：替换原始engine列
    # .rename({"engine_clean": "engine"})
)

3. 单正则分组一次性提取，提升性能 #

如果数据格式相对固定，用多分组正则一次性匹配所有字段，减少多次正则匹配的开销，适合大数据量场景：

(
    df
    .with_columns(
        pl.col("engine")
        .str.replace(r'\sLiter', "L")
        .str.extract(r'(?:(\d+)\.\d+HP)?.*?(\d+\.\d+)L.*?(?:(\d+)\sCylinder)?')
        .alias("engine_details")
    )
    .with_columns(
        pl.col("engine_details").list.get(0).alias("HP"),
        pl.col("engine_details").list.get(1).alias("Displacement"),
        pl.col("engine_details").list.get(2).alias("Cylinder"),
        pl.col("engine").str.contains("Electric").alias("Electric")
    )
    .drop("engine_details")
)

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