Polars apply函数未传递列名，迁移Pandas代码遇计算难题求助

针对你从Pandas迁移到Polars时遇到的动态列名+行级计算的问题，这里提供两种解决方案，优先推荐向量化实现（符合Polars的设计理念，性能最优）：

方案1：向量化计算（推荐） #

Polars的核心优势是向量化操作，完全不需要依赖apply就能实现需求，性能比行级操作高几个数量级。利用polars.col()可以动态引用列名，直接完成计算：

import polars as pl

data = {'A': [1, 2, 3], 'B': [4, 5, 6]}
df = pl.DataFrame(data)

var1 = 'A'
var2 = 'B'

# 生成包含tuple的列表列
df = df.with_columns(
    pl.concat_list(
        pl.struct(pl.col(var1) + pl.col(var2), pl.col(var1) * pl.col(var2))
    ).alias('calc_data')
)

如果需要更灵活的结构，也可以用map_elements配合Struct类型：

df = df.with_columns(
    pl.struct([pl.col(var1), pl.col(var2)])
    .map_elements(
        lambda x: list(zip([x[var1]+x[var2]], [x[var1]*x[var2]])),
        return_dtype=pl.List(pl.Tuple(pl.Int64, pl.Int64))
    )
    .alias('calc_data')
)

方案2：行级操作（仅用于复杂逻辑） #

如果实际计算逻辑比示例复杂，必须用行级处理，Polars的map_elements可以结合参数传递动态列名，避免依赖全局变量：

import polars as pl

data = {'A': [1, 2, 3], 'B': [4, 5, 6]}
df = pl.DataFrame(data)

var1 = 'A'
var2 = 'B'

def some_calculation(row, col1, col2):
    cal1 = row[col1] + row[col2]
    cal2 = row[col1] * row[col2]
    return list(zip([cal1], [cal2]))

# 通过闭包传递列名参数给自定义函数
df = df.with_columns(
    pl.struct(var1, var2)
    .map_elements(lambda row: some_calculation(row, var1, var2), return_dtype=pl.List(pl.Tuple(pl.Int64, pl.Int64)))
    .alias('calc_data')
)

关键说明： #

• Polars的apply（DataFrame级）是按列操作的，而非Pandas的行级逻辑，行级处理要对应map_elements（作用于Series/Struct）或row方法。
• 尽量避免行级操作，大数据集下向量化操作的性能碾压行级循环。
• 动态列名通过pl.col(col_name)或直接传入字符串引用，通过参数传递列名比依赖外部变量更清晰可控。

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