如何对两个Polars LazyFrame执行按列对应相乘操作？

我定义了两个结构完全一致的Polars LazyFrame：

import polars as pl
import numpy as np

n = 5  # 补充示例用的行数
df1 = pl.LazyFrame(data={
    'foo': np.random.uniform(0,127, size= n).astype(np.float64),
    'bar': np.random.uniform(1e3,32767, size= n).astype(np.float64),
    'baz': np.random.uniform(1e6,2147483, size= n).astype(np.float64)
})

df2 = pl.LazyFrame(data={
    'foo': np.random.uniform(0,127, size= n).astype(np.float64),
    'bar': np.random.uniform(1e3,32767, size= n).astype(np.float64),
    'baz': np.random.uniform(1e6,2147483, size= n).astype(np.float64)
})

我需要将df1和df2中对应列的每一行元素相乘。如果先把LazyFrame转为非延迟的DataFrame，可以直接用*操作符实现：

df1.collect() * df2.collect()

得到的结果类似：

foo         bar         baz
f64         f64         f64
3831.295563 6.4637e6    3.3669e12
164.194271  2.9691e8    2.2696e12
3655.918761 1.9444e7    2.3625e12
7191.48868  3.7044e7    3.1687e12
9559.505277 2.6864e8    2.5426e12

但直接对LazyFrame执行df1 * df2会抛出异常：

TypeError: unsupported operand type(s) for *: 'LazyFrame' and 'LazyFrame'

请问该如何在LazyFrame上实现对应列的逐元素相乘？

解决方案 #

Polars的LazyFrame不支持直接用*操作符做整体运算，需要显式指定每一列的相乘逻辑，以下是两种简洁的实现方式：

方法1：利用pl.col().from()引用另一个LazyFrame的列 #

通过from()方法可以直接从另一个LazyFrame中获取同名列，然后逐列定义相乘表达式：

# 遍历所有列，生成对应列的相乘表达式
result = df1.select([
    (pl.col(col) * pl.col(col).from(df2)).alias(col) 
    for col in df1.columns
])

# 执行计算并获取结果
result.collect()

方法2：先横向合并LazyFrame再计算 #

先将两个LazyFrame横向合并（给第二个LazyFrame的列加后缀区分），再对每对原列和后缀列做相乘：

# 横向合并，给df2的列加后缀"_df2"
merged_df = df1.hstack(df2, suffix="_df2")

# 遍历原列名，计算对应列的乘积并重命名回原列名
result = merged_df.select([
    (pl.col(col) * pl.col(f"{col}_df2")).alias(col)
    for col in df1.columns
])

result.collect()

两种方法都能保留LazyFrame的延迟计算特性，直到调用collect()时才实际执行运算，适合处理大数据集。

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