Polars中使用.when.then时新列类型的判定规则及类型保持方法

问题背景 #

使用Polars 0.18.2版本时，通过pl.when.then.otherwise为DataFrame添加或修改列时，对新列的类型继承规则存在困惑，以下是示例场景及现象：

首先构造示例DataFrame：

import polars as pl
df = pl.DataFrame({"a": ["x", "y"], "b": [1, 2]}, schema={"a": pl.String, "b": pl.UInt8})

输出：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 2   │
└─────┴─────┘

三种场景现象与原因解析 #

Polars对when.then.otherwise表达式的类型推导，核心规则是选择所有分支（所有then和otherwise）都能兼容的最小数据类型，同时字面量默认类型、数据范围兼容性会影响最终结果：

场景1：单when分支+字面量otherwise，类型保持UInt8

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .otherwise(5)
)

输出：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 5   │
└─────┴─────┘

原因：then分支是UInt8类型的列b，otherwise的字面量5在UInt8的取值范围（0-255）内，Polars会自动将字面量5转换为UInt8，所有分支类型统一，最终列类型保持UInt8。

场景2：多when分支+字面量otherwise，类型变为Int32

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(0)
    .otherwise(5)
)

输出：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ i32 │ # <-- Int32 instead of UInt8
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 0   │
└─────┴─────┘

原因：Polars中整数字面量（如0、5）的默认类型是Int32。此时三个分支的类型分别是：UInt8（b列）、Int32（0）、Int32（5）。由于UInt8可以安全向上转换为Int32，但Int32无法全部兼容UInt8的反向转换，所以Polars选择所有分支都能兼容的最小类型Int32。

场景3：多when分支+引用列的otherwise，类型回到UInt8

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(0)
    .otherwise(pl.col("b")) # <-- changed to pl.col('b') instead of a number
)

输出：

┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 0   │
└─────┴─────┘

原因：此时分支类型为UInt8（b列）、Int32（0）、UInt8（b列）。Polars检查到字面量0在UInt8的取值范围内，会自动将其转换为UInt8，所有分支类型统一为UInt8，最终列类型保持一致。

确保新列与原列类型一致的方法 #

不需要每次都显式转换，但可以通过以下方式避免类型意外变化：

• 显式转换字面量类型：将字面量用pl.lit()包裹并转换为目标类型，比如把0改为pl.lit(0).cast(pl.UInt8)，5改为pl.lit(5).cast(pl.UInt8)，确保所有分支类型与原列一致。
• 整体表达式转换：在when表达式末尾添加.cast(pl.UInt8)，强制将最终结果转换为目标类型，但要注意数据溢出风险（比如字面量超出UInt8范围会抛出错误）。
• 优先引用原列：在otherwise或then分支中尽量引用原列而非字面量，减少类型推导的不确定性。

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