Polars中如何通过依赖多列的函数同时更新指定列并保留原列与数据类型？

嘿，这个问题我太懂了，之前处理Polars多列依赖更新的时候也踩过类似的坑！你之前用struct+map_elements的方法确实有点繁琐，还丢列、改类型，我给你两个更优雅的解决方案，完美解决你提到的三个问题！

方案一：用map_rows直接处理每行（推荐，Polars 0.19.0+可用） #

Polars在0.19.0版本之后新增了map_rows方法，它可以直接把每行数据转换成字典让你操作，修改完指定列后返回整个行，自动保留所有原列，简直是为这种场景量身定做的！

先看代码：

def update_row(row):
    # 计算总和s，用原行的数值直接运算
    s = row["a"] + row["b"] + row["c"] + row["i"]
    
    # 更新a、b、c，这里按照你的需求转回原Int8类型（根据业务调整转换逻辑）
    # 比如把比例值缩放至Int8的范围（-128到127）后取整
    row["a"] = int(round((row["a"] / s) * 127))
    row["b"] = int(round((row["b"] / s) * 127))
    row["c"] = int(round((row["c"] / s) * 127))
    
    # i和o列完全不动，直接保留原数值和类型
    return row

# 应用函数并收集结果
result = lf.map_rows(update_row).collect()

这个方案的优势简直拉满：

• 自动保留所有原列：不用手动写select/unnest，o列直接留在结果里
• 类型完全可控：i列还是原来的Int8，a、b、c也能转回原类型（转换逻辑你可以根据业务改）
• 代码超直观：直接操作行字典，逻辑一目了然，再也不用写嵌套的struct操作了

方案二：旧版本Polars兼容方案（用with_columns+map_elements） #

如果你还在使用低于0.19.0的Polars版本，那可以改进你之前的写法，用with_columns替代select来保留原列，同时指定返回类型来避免类型变化：

def update_struct(s):
    s_val = s["a"] + s["b"] + s["c"] + s["i"]
    # 计算更新后的a、b、c并转回Int8
    new_a = int(round((s["a"] / s_val) * 127))
    new_b = int(round((s["b"] / s_val) * 127))
    new_c = int(round((s["c"] / s_val) * 127))
    
    # 返回包含更新后字段的struct，明确指定每个字段的类型
    return pl.struct(
        a=new_a,
        b=new_b,
        c=new_c,
        i=s["i"]  # 保留原i的数值和类型
    )

result = (
    lf
    # 用with_columns而不是select，保留所有原列
    .with_columns(
        pl.struct("a", "b", "c", "i")
        .map_elements(
            update_struct,
            # 指定返回的struct类型，确保字段类型和原数据一致
            return_dtype=pl.Struct([
                pl.Field("a", pl.Int8),
                pl.Field("b", pl.Int8),
                pl.Field("c", pl.Int8),
                pl.Field("i", pl.Int8)
            ])
        )
        .alias("updated")
    )
    .unnest("updated")
    .collect()
)

这个方案解决了之前的三个问题：

• 用with_columns替代select，o列不会丢失
• 手动指定返回的struct类型，i列保持Int8不变
• 逻辑更清晰，明确控制每个字段的处理方式

关于数据类型的补充说明 #

因为原a、b、c是Int8类型，而a/s得到的是浮点数，所以需要根据你的业务需求把结果转换回Int8。上面的例子用了“缩放至Int8范围后取整”的逻辑，你也可以改成直接取整、截断，或者如果业务允许a、b、c变成浮点数，那就去掉类型转换的部分，在return_dtype里指定对应的Float类型就行啦。

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