如何在Python Polars中通过单次函数调用为DataFrame添加多列？

现有代码为DataFrame添加count_of_10和count_of_12两列时，每行都会调用两次uses_object函数，效率较低。想优化成每行仅调用一次该函数。

原实现代码：

import polars as pl

df = pl.DataFrame({
    'x': [11, 22],
})

def uses_object(x):
    r = list(range(0, x))
    c10 = r.count(10)
    c12 = r.count(12)
    return c10, c12

df = df.with_columns(
    count_of_10 = pl.col('x').map_elements(lambda x: uses_object(x)[0]),
    count_of_12 = pl.col('x').map_elements(lambda x: uses_object(x)[1]),
)

print(df)

尝试直接多赋值的写法会报错：

df = df.with_columns(
    count_of_10, count_of_12 = uses_object(pl.col('x')),
)

报错信息：

NameError
name 'count_of_10' is not defined.

方案1：返回结构体并展开 #

让uses_object返回字典形式的结构体，通过map_elements一次性生成包含两个字段的结构体列，再展开成独立列。这样每行只会调用一次uses_object：

import polars as pl

df = pl.DataFrame({
    'x': [11, 22],
})

def uses_object(x):
    r = list(range(0, x))
    c10 = r.count(10)
    c12 = r.count(12)
    return {'count_of_10': c10, 'count_of_12': c12}

df = df.with_columns(
    pl.col('x').map_elements(
        uses_object, 
        return_dtype=pl.Struct([
            pl.Field('count_of_10', pl.Int64),
            pl.Field('count_of_12', pl.Int64)
        ])
    ).alias('temp')
).unnest('temp')

print(df)

输出结果和原代码一致，但uses_object每行仅执行一次。

方案2：用元组+结构体简化 #

如果不想修改uses_object的返回格式，保持返回元组，用pl.struct包装后展开，再重命名字段即可：

import polars as pl

df = pl.DataFrame({
    'x': [11, 22],
})

def uses_object(x):
    r = list(range(0, x))
    c10 = r.count(10)
    c12 = r.count(12)
    return c10, c12

df = df.with_columns(
    pl.struct(
        pl.col('x').map_elements(uses_object, return_dtype=pl.Tuple([pl.Int64, pl.Int64]))
    ).alias('temp')
).unnest('temp').rename({'field_0': 'count_of_10', 'field_1': 'count_of_12'})

print(df)

方案3：apply逐行处理（小数据集适用） #

如果数据集不大，也可以用apply方法，每行处理一次并返回多列：

import polars as pl

df = pl.DataFrame({
    'x': [11, 22],
})

def uses_object(row):
    x = row['x']
    r = list(range(0, x))
    c10 = r.count(10)
    c12 = r.count(12)
    return pl.Series([c10, c12], name=['count_of_10', 'count_of_12'])

df = df.apply(uses_object, return_dtype=pl.Struct)

print(df)

注意：apply是逐行处理，性能不如向量化操作，仅适合小数据场景。

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