Polars中when().then(map_elements)为何全量调用函数？如何优化

你遇到的情况不是操作错误，而是Polars的map_elements执行逻辑导致的：map_elements会对整个列的所有元素执行自定义函数，生成完整的结果列后，再由when/otherwise根据条件选择保留对应位置的值或设为None。所以哪怕你把map_elements放在then分支里，它依然会先处理全列，这才导致所有条目都触发了get_year的调用。

推荐优先用Polars原生表达式替代自定义函数（性能远高于map_elements），如果必须保留自定义函数，也可以通过调整逻辑实现按需调用。

方案1：用Polars原生表达式实现（推荐） #

提取年份的逻辑完全可以用Polars内置的字符串和列表操作实现，不需要自定义函数，既高效又不会有多余调用：

import polars as pl

df = pl.DataFrame({
    'Type': ['CAT_YY', 'CAT5', 'CAT_YY', 'CATX'],
    'ID':   ['DEBT.24', 'CASH.EU', 'DEBT.26', 'CASH.US'],
}).with_columns(
    year=pl.when(pl.col('Type') == 'CAT_YY')
           # 分割ID字符串，取第二部分并转为整数
           .then(pl.col('ID').str.split('.').list.get(1).cast(pl.Int64))
           .otherwise(None)
)

print(df)

执行后只会对Type为CAT_YY的行处理，没有多余的函数调用，输出结果和原代码一致。

方案2：保留自定义函数，实现按需调用 #

如果你的实际逻辑无法用原生表达式替代，可以通过以下两种方式优化：

方式A：在自定义函数内部添加条件判断 #

修改get_year，让它先检查Type是否满足要求（需要把Type和ID打包成结构体传入）：

import polars as pl

def get_year(data: dict) -> int | None:
    entry_type = data['Type']
    entry_id = data['ID']
    if entry_type != 'CAT_YY':
        return None
    arr = entry_id.split('.')
    print(f'> Function is called for {entry_id} <')
    if arr[0] != 'DEBT':
        return None
    else:
        return int(arr[1])

df = pl.DataFrame({
    'Type': ['CAT_YY', 'CAT5', 'CAT_YY', 'CATX'],
    'ID':   ['DEBT.24', 'CASH.EU', 'DEBT.26', 'CASH.US'],
}).with_columns(
    year=pl.struct(['Type', 'ID']).map_elements(get_year, return_dtype=pl.Int64)
)

print(df)

此时函数依然会遍历所有行，但会跳过不满足Type == 'CAT_YY'的条目处理，日志只会输出符合条件的两行：

> Function is called for DEBT.24 <
> Function is called for DEBT.26 <

方式B：先过滤处理再合并 #

先筛选出满足条件的行处理，再合并回原DataFrame：

import polars as pl

def get_year(entry_id: str) -> int | None:
    arr = entry_id.split('.')
    print(f'> Function is called for {entry_id} <')
    if arr[0] != 'DEBT':
        return None
    else:
        return int(arr[1])

df = pl.DataFrame({
    'Type': ['CAT_YY', 'CAT5', 'CAT_YY', 'CATX'],
    'ID':   ['DEBT.24', 'CASH.EU', 'DEBT.26', 'CASH.US'],
})

# 处理满足条件的行
processed_df = df.filter(pl.col('Type') == 'CAT_YY').with_columns(
    year=pl.col('ID').map_elements(get_year, return_dtype=pl.Int64)
)

# 合并回原数据，填充其他行的year为None
result_df = df.join(processed_df, on=['Type', 'ID'], how='left')
print(result_df)

这种方式下get_year只会被调用两次，完全符合按需执行的要求，输出结果和原代码一致。

内容的提问来源于stack exchange，提问作者Phil-ZXX