如何在Polars中替换值时避免数值类型被转为字符串？

Polars 默认会自动进行隐式类型转换来兼容替换值，但要实现类型不匹配时提前报错、仅允许安全向上转型的需求，可以通过类型预校验结合严格模式操作来实现，以下是两种基于 Expr 级的解决方案：

方法1：封装自定义严格替换函数 #

通过自定义 Expr 级函数，提前校验替换值与原列的类型兼容性，仅允许同类型或安全向上转型（如 Int32 → Int64、Float32 → Float64），不兼容则直接抛出错误。

import polars as pl
from polars.datatypes import (
    Int32, Int64, Float32, Float64, String, Enum, DataType
)

def strict_replace(expr: pl.Expr, mapping: dict) -> pl.Expr:
    orig_dtype = expr.dtype
    
    # 定义允许的兼容类型组：仅组内类型可互相替换/向上转型
    compatible_groups = {
        Int32: {Int32, Int64},
        Int64: {Int64},
        Float32: {Float32, Float64},
        Float64: {Float64},
        String: {String},
        Enum: {Enum}
    }
    
    # 遍历替换值，逐一校验类型兼容性
    for val in mapping.values():
        val_dtype = pl.lit(val).dtype
        orig_group = next((g for g in compatible_groups if isinstance(orig_dtype, g)), None)
        
        if not orig_group:
            raise TypeError(f"不支持的原列类型: {orig_dtype}")
        if not isinstance(val_dtype, tuple(compatible_groups[orig_group])):
            raise TypeError(
                f"替换值类型 {val_dtype} 与原列类型 {orig_dtype} 不兼容，禁止跨类型转换"
            )
    
    # 执行严格替换，自动处理合法的向上转型
    return expr.replace_strict(mapping)

测试示例 #

1. 字符串列用数值替换：触发报错

dat = pl.DataFrame({"x": ["a", "b"]})
# 执行后抛出 TypeError: 替换值类型 Int64 与原列类型 String 不兼容，禁止跨类型转换
dat.with_columns(y=strict_replace(pl.col("x"), {"a": 1}))

2. Int32列用Int64值替换：自动向上转型

dat_int = pl.DataFrame({"x": [1, 2]}, dtype=pl.Int32)
result = dat_int.with_columns(y=strict_replace(pl.col("x"), {1: 10000000000}))
print(result.schema)
# 输出：{'x': Int32, 'y': Int64}

3. 同类型替换：正常执行

result = dat.with_columns(y=strict_replace(pl.col("x"), {"a": "c"}))
print(result)
# shape: (2, 2)
# ┌─────┬─────┐
# │ x   ┆ y   │
# │ --- ┆ --- │
# │ str ┆ str │
# ╞═════╪═════╡
# │ a   ┆ c   │
# │ b   ┆ b   │
# └─────┴─────┘

方法2：原生Expr结合严格类型转换 #

如果不想封装函数，可直接在 when/then/otherwise 逻辑中加入严格模式的类型转换，强制校验替换值与原列的类型兼容性，不兼容则抛出计算错误。

dat = pl.DataFrame({"x": ["a", "b"]})
# 尝试用数值替换字符串，执行时抛出 ComputeError: strict casting from Int64 to String failed
dat.with_columns(
    y=pl.when(pl.col("x") == "a")
    .then(pl.lit(1).cast(pl.col("x").dtype, strict=True))
    .otherwise(pl.col("x"))
)

关键说明 #

• 两种方案均为 Expr 级操作，可无缝集成到 Polars 链式调用中
• 兼容类型组可根据业务需求扩展（如允许 UInt 转 Int），但需注意数值溢出风险
• 严格模式的核心是避免跨大类类型转换（如数值→字符串、字符串→枚举），仅保留合法的向上转型逻辑

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