如何在Polars流引擎中便捷实现矩阵乘法？

1. 核心思路 #

Polars的流引擎天然支持逐行运算，我们可以通过定义矩阵/向量的列映射规则，自动生成Polars表达式，彻底替代手动编写乘加项的繁琐操作。核心是把线性代数运算的矩阵索引，直接映射到数据表的列名，再自动展开为对应的逐行运算表达式。

2. 基础实现：矩阵-向量乘法自动生成 #

下面的工具函数可以输入矩阵对应的列组、向量对应的列名，自动生成Polars的输出列表达式：

import polars as pl

def mat_vec_expr(matrix_col_groups: list[list[str]], vec_cols: list[str]) -> list[pl.Expr]:
    """
    生成矩阵-向量乘法的Polars表达式列表，每个元素对应输出向量的一个分量
    """
    output_exprs = []
    # 遍历矩阵的每一行，生成对应分量的点积表达式
    for row_cols in matrix_col_groups:
        dot_product = sum(pl.col(col) * pl.col(vec_col) for col, vec_col in zip(row_cols, vec_cols))
        output_exprs.append(dot_product.alias(f"v{len(output_exprs)+1}"))
    return output_exprs

# 针对你的场景调用示例
matrix_cols = [["c1","c2","c3"], ["c4","c5","c6"], ["c7","c8","c9"]]
vec_cols = ["c48","c49","c50"]
result_exprs = mat_vec_expr(matrix_cols, vec_cols)

# 用流模式处理超大规模数据
df = pl.scan_csv("your_large_data.csv")  # 扫描模式不加载全量数据
result_df = df.with_columns(result_exprs)
result_df.sink_parquet("output.parquet")  # 直接写入，无需手动分批

这个函数会自动生成v1/v2/v3三个表达式，完全等价于手动编写的乘加逻辑，同时适配Polars流引擎的内存友好特性。

3. 扩展：链式矩阵乘法（行内） #

如果需要链式运算（比如M2*(M1*v)），可以基于基础函数嵌套实现自动展开：

def chain_mat_vec_expr(matrix_chain: list[list[list[str]]], vec_cols: list[str]) -> list[pl.Expr]:
    """
    处理链式矩阵-向量乘法，按输入顺序依次运算
    matrix_chain: 按运算顺序排列的矩阵列组列表，例如[M1_cols, M2_cols]
    """
    current_vec_exprs = [pl.col(col) for col in vec_cols]
    for matrix_cols in matrix_chain:
        next_vec_exprs = []
        for row_cols in matrix_cols:
            dot_product = sum(pl.col(col) * vec_expr for col, vec_expr in zip(row_cols, current_vec_exprs))
            next_vec_exprs.append(dot_product)
        current_vec_exprs = next_vec_exprs
    # 给最终结果命名
    return [expr.alias(f"res{idx+1}") for idx, expr in enumerate(current_vec_exprs)]

# 示例：先执行M1*v，再用M2乘中间结果
M1_cols = [["c1","c2"], ["c3","c4"]]  # 2x2矩阵
M2_cols = [["c5","c6"], ["c7","c8"]]  # 2x2矩阵
vec_cols = ["c49","c50"]
chain_exprs = chain_mat_vec_expr([M1_cols, M2_cols], vec_cols)

4. 关键注意事项 #

• 所有运算都是逐行独立的，完全适配Polars流引擎，无论数据集多大，都不会一次性加载全量数据到内存。
• 即使矩阵规模很大（如100x100），生成的表达式会自动被Polars查询优化器处理，不会影响性能。
• 需保证矩阵与向量的维度匹配：矩阵的列数必须等于对应向量的长度，否则会抛出运算错误。

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