在 polars 中无法直接对 lazy dataframe 按行执行不可向量化的自定义函数（如 softmax），但可通过 unpivot → group-wise 计算 → pivot 的三步范式高效实现，兼顾惰性求值与语义清晰性。
在 polars 中无法直接对 lazy dataframe 按行执行不可向量化的自定义函数（如 softmax），但可通过 unpivot → group-wise 计算 → pivot 的三步范式高效实现，兼顾惰性求值与语义清晰性。
Polars 的设计哲学强调列式计算与向量化操作，因此原生不支持 map_rows 在 Lazy 模式下的直接使用（尤其在涉及跨列逻辑时）。当你的自定义函数（例如 softmax）需基于每行多个数值联合计算且无法用 pl.Expr 表达时，强行使用 map_batches + map_rows 会破坏惰性执行链、丢失类型推断，并可能导致意外错误或性能退化。
✅ 推荐方案：Unpivot → Group-wise Transform → Pivot
该模式将“按行计算”转化为“按分组聚合/映射”，完全兼容 Lazy API，且能保留所有原始列（如 id）：
import polars as pl
import numpy as np
def softmax(arr):
exp_arr = np.exp(arr)
return exp_arr / exp_arr.sum()
# 构建 Lazy DataFrame
ldf = pl.DataFrame({
'id': [1, 2, 3],
'a': [0.2, 0.1, 0.3],
'b': [0.4, 0.1, 0.3],
'c': [0.4, 0.8, 0.4]
}).lazy()
# ✅ 正确流程：unpivot → map_batches over group → pivot
result = (
ldf
.unpivot(index="id", on=["a", "b", "c"], variable_name="variable", value_name="value")
.with_columns(
pl.col("value").map_batches(lambda s: pl.Series(softmax(s.to_numpy()))).over("id")
)
.collect() # collect 可选；若后续仍需 lazy 操作，可延迟至最终 collect
.pivot(values="value", index="id", columns="variable")
)
print(result)
输出：
shape: (3, 4) ┌─────┬──────────┬──────────┬──────────┐ │ id ┆ a ┆ b ┆ c │ │ --- ┆ --- ┆ --- ┆ --- │ │ i64 ┆ f64 ┆ f64 ┆ f64 │ ╞═════╪══════════╪══════════╪══════════╡ │ 1 ┆ 0.290461 ┆ 0.35477 ┆ 0.35477 │ │ 2 ┆ 0.249143 ┆ 0.249143 ┆ 0.501713 │ │ 3 ┆ 0.322043 ┆ 0.322043 ┆ 0.355913 │ └─────┴──────────┴──────────┴──────────┘
? 关键说明：
-
unpivot(index="id")将宽表转为长表，每行代表原表中某一行某一列的值，同时保留id作为分组标识； -
.map_batches(...).over("id")对每个id组内的value列批量应用 softmax（注意：输入是pl.Series，需转为np.ndarray再处理，返回pl.Series以保持类型安全）； - 最终
pivot()恢复原始宽格式；若需全程 Lazy，可将.collect()移至管道末尾，仅在必要时触发计算。
⚠️ 注意事项：
- 避免在
map_batches中使用map_rows—— 它会强制 materialize 整个 chunk，丧失 Lazy 优势； -
softmax等函数必须保证输入/输出长度一致，否则.over()将报错； - 若原始数据含缺失值，建议在
softmax前添加np.nan_to_num或显式处理，防止np.exp溢出或 NaN 传播； - 对超大规模数据，此模式仍保持内存友好性，因
unpivot和over均支持流式分块处理。
该方法不仅是技术 workaround，更是 Polars “以列为本、以组为界”计算范式的自然延伸——将“行逻辑”重表述为“组内列变换”，既符合底层引擎优化逻辑，也提升了代码可维护性与可测试性。
