本文介绍如何在 Polars LazyFrame 中以最简、高效且可读的方式实现行级 Softmax 变换，避免冗长的列循环和临时列管理，充分利用 pl.all() 和表达式自动缓存机制。
本文介绍如何在 polars lazyframe 中以最简、高效且可读的方式实现行级 softmax 变换，避免冗长的列循环和临时列管理，充分利用 `pl.all()` 和表达式自动缓存机制。
Softmax 是机器学习中常用的归一化操作，其核心公式为：
$$
\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}
$$
在 Polars 中对多列（如 'a', 'b', 'c'）按行计算 Softmax，关键在于对整行进行指数变换后，再逐列除以该行指数和。传统写法（如使用列表推导式 + with_columns 多次调用）不仅冗长，还易出错且难以维护。
✅ 推荐做法：使用 pl.all() 进行全列批量操作，并借助 Polars 的表达式自动去重与缓存机制（CSER, Common Subexpression Elimination），让底层自动复用 pl.all().exp() 的计算结果：
import polars as pl
df = pl.DataFrame({
'a': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
'b': [5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
'c': [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
}).lazy()
# ✅ 一行完成 Softmax（作用于所有数值列）
result = df.with_columns(
pl.all().exp() / pl.sum_horizontal(pl.all().exp())
).collect()
print(result)
? 原理说明：
pl.all()默认选取所有列（等价于pl.col(pl.datatypes.NUMERIC_DTYPES)在新版中更安全），pl.sum_horizontal(...)沿行求和。Polars 查询优化器会自动识别pl.all().exp()被重复引用两次，将其物化为一个隐藏中间列（如__POLARS_CSER_0x...），避免重复计算——这正是.explain()输出中体现的「Common Subexpression Elimination」优化。
? 若仅需对特定列应用 Softmax（例如仅 'a', 'b', 'c'），推荐显式指定列名以提升可读性与健壮性：
cols = ['a', 'b', 'c']
result = df.with_columns(
pl.col(cols).exp() / pl.sum_horizontal(pl.col(cols).exp())
).collect()
? 调试与验证建议：
- 使用
.explain()查看执行计划，确认是否触发 CSER 优化； - 对比
.collect()前后数据形状与数值合理性（如每行 Softmax 结果之和应 ≈ 1.0）； - 注意：
pl.all()在含非数值列（如字符串、时间戳）的 DataFrame 中可能报错，此时务必显式限定列或使用pl.col(pl.datatypes.FLOAT_DTYPES | pl.datatypes.INTEGER_DTYPES)。
✅ 总结：告别手动循环与临时列管理，善用 pl.all() 或 pl.col(list) 批量表达式 + Polars 内置优化，即可用单行声明式代码安全、高效、清晰地完成 LazyFrame 上的 Softmax 计算。
