如何用Polars内置函数替代循环实现按ID生成可变长度结构体？

问题背景 #

你当前通过for循环实现了根据id生成不同字段结构体的需求，但希望改用Polars内置函数替代循环，提升代码简洁性和性能。

优化方案1：分组处理（推荐，性能优先） #

利用Polars的group_by+apply机制，按id分组后批量生成对应结构体，避免循环修改DataFrame：

import polars as pl
pl.Config.set_fmt_str_lengths(1000)
import warnings
warnings.filterwarnings("ignore")

df = pl.DataFrame({
    "id": [1, 1, 2, 3],
    "dog": ["labrador", "labrador", "boxer", "airedale"],
    "age": [2, 2, 3, 4],
    "owner": ["Jim", "Paul", "Kim", "Lynne"],
    "spots": [None, None, True, None],
    "food": ["Raw", "Kibble", None, "Raw"],
    "leash": [True, False, None, True]
})

id_struct_dict = {
    1: ["owner", "food"],
    2: ["owner", "spots", "food"],
    3: ["owner", "food", "leash"]
}

def process_group(group_df):
    current_id = group_df["id"][0]
    target_fields = id_struct_dict[current_id]
    return group_df.with_columns(
        pl.struct(target_fields).map_elements(str).alias("struct")
    )

# 分组处理并合并结果
df_optimized = df.group_by("id", maintain_order=True).apply(process_group)
print(df_optimized)

优化方案2：动态字段映射（纯表达式实现） #

通过map_elements动态获取当前id对应的字段列表，结合list.eval构造结构体，全程使用Polars表达式链：

df_optimized2 = df.with_columns(
    pl.struct(
        # 根据id获取目标字段列表，再映射到对应列
        pl.col("id").map_elements(lambda x: id_struct_dict[x], return_dtype=pl.List(pl.String))
        .list.eval(pl.col().map(lambda col_name: pl.col(col_name)))
    ).map_elements(str).alias("struct")
)
print(df_optimized2)

方案优势说明 #

• 避免了原循环中多次修改DataFrame的冗余操作，代码更简洁易维护
• 分组处理方案符合Polars向量化计算的设计思想，数据量越大性能优势越明显
• 纯表达式方案无需额外定义函数，适合追求代码紧凑的场景

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