本文介绍如何使用 Polars 表达式重写传统 NumPy 风格的加权/非加权相关系数函数，并在 group_by().agg() 中高效调用，避免 map_batches 或命名空间注册等复杂操作。
本文介绍如何使用 polars 表达式重写传统 numpy 风格的加权/非加权相关系数函数，并在 `group_by().agg()` 中高效调用，避免 `map_batches` 或命名空间注册等复杂操作。
在 Polars 中实现自定义相关系数（如无需去均值的 Pearson 变体），关键在于完全基于表达式（Expression）编程，而非依赖 NumPy 数组或 Python 标量运算。Polars 的 agg() 上下文天然支持表达式链式计算，因此只需将原始 NumPy 函数逻辑迁移为 Polars 表达式语法即可——无需 map_elements（逐元素低效）、不必 map_batches（需手动处理 chunk）、更无需注册自定义命名空间。
以下是推荐的实现方式：
import polars as pl
def rcor(x: pl.Expr, y: pl.Expr, w: pl.Expr | None = None) -> pl.Expr:
"""
计算无需去均值的相关系数：r = Σ(w·x·y) / √[Σ(w·x²) · Σ(w·y²)]
支持加权（w 非 None）与非加权（w 为 None）两种模式。
"""
if w is not None:
sxx = (w * x * x).sum()
syy = (w * y * y).sum()
sxy = (w * x * y).sum()
else:
sxx = (x * x).sum()
syy = (y * y).sum()
sxy = (x * y).sum()
return sxy / (sxx * syy).sqrt()
使用时直接传入列表达式，无缝集成进聚合流程：
result = (
df.group_by(["UNIVERSE", "datetime"])
.agg(
corr_xy = rcor(pl.col("ratio_wsm"), pl.col("y1d_nn_r")),
# 如需加权版本（假设存在权重列 'weight'）：
# corr_wt = rcor(pl.col("ratio_wsm"), pl.col("y1d_nn_r"), pl.col("weight"))
)
)
✅ 优势说明：
- 零拷贝 & 向量化：全程运行于 Polars 引擎内部，利用 Arrow 内存模型和查询优化器，性能远超 map_batches；
- 类型安全 & 延迟执行：返回 pl.Expr，可与其他表达式组合（如 .filter()、.with_columns()）；
- 自动分组对齐：group_by().agg() 自动按分组键广播计算，无需手动切片或索引；
- 支持空值处理：默认遵循 Polars 的 null 语义（sum() 自动跳过 null），如需特殊处理可添加 .drop_nulls() 链式调用。
⚠️ 注意事项：
- 输入 x, y, w 必须为 pl.Expr 类型（如 pl.col("col_name")），不可传入 Series 或 numpy.ndarray；
- 若输入列含大量 null 值，建议提前清洗（如 df.drop_nulls(["ratio_wsm", "y1d_nn_r"])），否则 sqrt(0) 可能导致 NaN；
- 该公式本质是“余弦相似度”形式，不具统计学上的中心化性质，请确保业务逻辑明确要求此定义。
总结：Polars 的强大之处在于将“自定义聚合”简化为表达式组合。只要坚持用 pl.Expr 构建逻辑，就能写出简洁、高效、可维护的分组计算代码——无需绕路，直击核心。
