如何在Polars中使用map_batches将动态列数的struct传入@guvectorize装饰的自定义函数，并将结果展开为指定后缀的新列

嘿，这个问题我帮你捋清楚！你需要处理动态数量的列，用@guvectorize加速计算，还要把结果拆成带「_result」后缀的新列，咱们一步步来实现：

1. 定义适配动态列的@guvectorize函数 #

首先，因为我们要处理的是动态数量的列，这些列会被合并成一个二维数组（每行对应原数据的一行，每列对应选中的原列）。所以@guvectorize的签名要写成"(m,n)->(m,n)"，其中m是行数，n是动态的列数——这样不管你选多少列，函数都能适配。

举个例子，假设你的复杂操作是给每个元素乘2，函数可以这么写：

import polars as pl
import numpy as np
from numba import guvectorize

@guvectorize([(float64[:,:], float64[:,:])], "(m,n)->(m,n)")
def process_array(arr, result):
    # 这里替换成你的实际复杂数组操作
    result[:] = arr * 2

注意：result[:] = ...这种写法是guvectorize要求的，必须直接修改传入的result数组，不能返回新数组。

2. 在Polars中处理动态列并调用函数 #

接下来，我们用pl.concat_arr把所有目标列（比如所有浮点列）合并成一个数组列，再通过map_batches调用刚才的函数：

# 示例DataFrame，这里故意加了3列来测试动态列数
df = pl.DataFrame({
    "values1": [1.0, 2.0, 3.0],
    "values2": [10.0, 20.0, 30.0],
    "values3": [100.0, 200.0, 300.0]
})

# 动态选中所有浮点列，存下来方便后续重命名
target_cols = df.select(pl.cs.float()).columns

# 生成结果数组列
df_with_result = df.with_columns(
    pl.concat_arr(target_cols)
    .map_batches(process_array)
    .alias("result_array")
)

3. 将结果数组展开为带后缀的新列 #

现在我们有了一个result_array列，每个元素是对应行的计算结果数组。接下来要把它拆成带「_result」后缀的新列：

final_df = (
    df_with_result
    # 把数组列转成struct，字段名用原列名
    .with_columns(
        pl.col("result_array").list.to_struct(fields=target_cols).alias("result_struct")
    )
    # 展开struct成单独的列
    .unnest("result_struct")
    # 给每个展开的列重命名，加上_result后缀
    .rename({col: f"{col}_result" for col in target_cols})
    # 删掉中间的临时列（可选）
    .drop("result_array")
)

print(final_df)

运行后你会得到这样的输出：

shape: (3, 6)
┌─────────┬─────────┬─────────┬──────────────┬──────────────┬──────────────┐
│ values1 ┆ values2 ┆ values3 ┆ values1_result ┆ values2_result ┆ values3_result │
│ ---     ┆ ---     ┆ ---     ┆ ---          ┆ ---          ┆ ---          │
│ f64     ┆ f64     ┆ f64     ┆ f64          ┆ f64          ┆ f64          │
╞═════════╪═════════╪═════════╪══════════════╪══════════════╪══════════════╡
│ 1.0     ┆ 10.0    ┆ 100.0   ┆ 2.0          ┆ 20.0         ┆ 200.0        │
│ 2.0     ┆ 20.0    ┆ 200.0   ┆ 4.0          ┆ 40.0         ┆ 400.0        │
│ 3.0     ┆ 30.0    ┆ 300.0   ┆ 6.0          ┆ 60.0         ┆ 600.0        │
└─────────┴─────────┴─────────┴──────────────┴──────────────┴──────────────┘

一些额外说明 #

• 如果你的自定义操作需要输出不同形状的结果（比如每行输出一个汇总值），只要调整guvectorize的签名就行，比如"(m,n)->(m)"，同时修改函数里的操作保证结果形状匹配。
• @guvectorize会利用Numba把函数编译成机器码，处理大数据量时比普通Python函数快很多，非常适合你的需求。

备注：内容来源于stack exchange，提问作者Andi