LazyFrame多列乘积计算方法

本文详细介绍了如何在Polars LazyFrame中，对两个具有相同结构和时间索引的数据帧进行列式乘法运算，同时排除作为索引的时间列。与Pandas的隐式操作不同，Polars通过利用`pl.struct`将非索引列打包、执行基于时间列的`join`操作，然后对结构体进行元素级乘法，最后使用`unnest`展开结果，从而实现高效且明确的多列乘积计算，尤其适用于大规模数据集的惰性计算场景。

Polars LazyFrame 多列乘积运算：排除特定列的结构化方法

在数据处理和分析中，经常需要对两个结构相似的数据集执行逐列操作，例如乘法。对于熟悉Pandas的用户来说，直接使用df1 * df2即可实现基于索引的列式乘积。然而，当迁移到Polars，特别是使用其强大的LazyFrame功能时，实现类似操作需要更具Polars特色的方法。本文将详细阐述如何在Polars LazyFrame中，高效地对两个DataFrame（排除作为索引的列）执行列式乘积运算。

1. 问题背景与Pandas对比

在Pandas中，如果两个DataFrame具有相同的索引和列名，可以直接通过乘法运算符进行逐元素相乘：

import pandas as pd
import numpy as np
import polars as pl

n = 5 # 示例数据行数

# 创建示例Pandas DataFrame
df_pd1 = pd.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).set_index('time')

df_pd2 = pd.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).set_index('time')

# Pandas中的直接乘法
result_pd = df_pd1 * df_pd2
print("Pandas 乘法结果:")
print(result_pd)

输出结果会是两个DataFrame对应列的乘积，且索引保持不变。

在Polars中，直接对两个LazyFrame执行类似操作（如df1 * df2）并不直接支持，因为Polars的运算符重载和表达式系统设计有所不同。尝试通过concat后group_by并聚合乘积，例如pl.concat([df1, df2]).group_by('time').agg(pl.col("*").mul(pl.col("*")))，通常会导致每个分组返回一个列表，而不是期望的单个乘积值，因为mul操作在聚合上下文中对列表中的每个元素执行，而非聚合整个组的乘积。

2. Polars LazyFrame 的解决方案：结构体与连接

Polars的强大之处在于其表达式系统和对结构化数据的灵活处理。解决此问题的核心思路是：将需要进行乘法运算的列（即除索引列外的所有列）打包成一个结构体（struct），然后通过索引列进行连接（join），最后对连接后的结构体执行乘法，并将结果结构体展开。

首先，创建两个Polars LazyFrame：

# 创建示例Polars LazyFrame
df_pl1 = pl.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).lazy()

df_pl2 = pl.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).lazy()

print("\n原始 Polars LazyFrame 1:")
print(df_pl1.collect())
print("\n原始 Polars LazyFrame 2:")
print(df_pl2.collect())

接下来，我们分步实现列式乘法：

步骤 1: 将非索引列打包成结构体

使用pl.struct表达式将除“time”列之外的所有列打包成一个名为“cols”的结构体。这使得我们可以将所有相关列作为一个整体进行处理。

# 将df1的非时间列打包成结构体
df_pl1_struct = df_pl1.select("time", cols=pl.struct(pl.exclude("time")))
print("\nLazyFrame 1 结构体化:")
print(df_pl1_struct.collect())

# 将df2的非时间列打包成结构体
df_pl2_struct = df_pl2.select("time", cols=pl.struct(pl.exclude("time")))
print("\nLazyFrame 2 结构体化:")
print(df_pl2_struct.collect())

pl.exclude("time")是一个非常实用的表达式，它会自动选择除“time”列之外的所有列。结果会是一个包含time列和一个名为cols的结构体列的LazyFrame。

步骤 2: 基于索引列进行连接

通过time列对两个结构体化的LazyFrame进行左连接（left join）。这将确保两个DataFrame中对应时间戳的结构体行被对齐。连接后，我们将得到一个包含time、cols（来自df_pl1）和cols_right（来自df_pl2）的LazyFrame。

# 连接两个结构体化的LazyFrame
joined_df = (
    df_pl1_struct
    .join(
       df_pl2_struct,
       on = "time",
       how = "left"
    )
)
print("\n连接后的 LazyFrame:")
print(joined_df.collect())

步骤 3: 对结构体执行元素级乘法

Polars允许直接对结构体列执行元素级运算，前提是结构体中的字段兼容。因此，我们可以直接将cols结构体乘以cols_right结构体。结果将是一个新的结构体，其中包含每个原始字段的乘积。

# 对连接后的结构体执行乘法
multiplied_struct_df = (
    joined_df
    .select("time", pl.col("cols") * pl.col("cols_right"))
)
print("\n结构体乘法后的 LazyFrame:")
print(multiplied_struct_df.collect())

步骤 4: 展开结果结构体

最后一步是使用unnest()方法将乘法结果结构体（这里仍然名为cols，因为它是pl.col("cols") * pl.col("cols_right")的结果，并被赋予了默认名cols）展开回独立的列。这将恢复原始的DataFrame结构，但其中的值已经是乘积结果。

# 展开结构体以获取最终结果
final_result_df = (
    multiplied_struct_df
    .unnest("cols")
)
print("\n最终结果 LazyFrame:")
print(final_result_df.collect())

3. 完整代码示例

将上述步骤整合到一起，形成完整的Polars解决方案：

import pandas as pd
import numpy as np
import polars as pl

n = 5 # 示例数据行数

# 创建示例Polars LazyFrame
df_pl1 = pl.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).lazy()

df_pl2 = pl.DataFrame(data={
    'time': pd.date_range('2023-01-01', periods=n, freq='1 min'),
    '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)
}).lazy()

# Polars 解决方案
result_pl = (
    df_pl1.select("time", cols=pl.struct(pl.exclude("time"))) # 将df1的非时间列打包
    .join(
       df_pl2.select("time", cols=pl.struct(pl.exclude("time"))), # 将df2的非时间列打包并连接
       on = "time",
       how = "left"
    )
    .select("time", pl.col("cols") * pl.col("cols_right")) # 对结构体执行乘法
    .unnest("cols") # 展开结果结构体
    .collect() # 收集结果
)

print("\nPolars LazyFrame 乘法结果:")
print(result_pl)

4. 注意事项与总结

• 惰性计算的优势： 这种方法充分利用了Polars的惰性计算特性。在.collect()被调用之前，所有的操作都只是构建查询计划，这对于处理大规模数据集非常高效。
• 明确性与控制： 相比Pandas的隐式行为，Polars要求更明确地定义操作。通过pl.struct和join，我们清晰地表达了“对齐这些数据，然后对特定部分执行操作”的意图。
• 灵活性： 这种结构体-连接-展开的模式不仅适用于乘法，也适用于其他需要对多列进行并行、元素级操作的场景。
• 列名冲突： 在join操作中，如果两个DataFrame有除连接键之外的同名列，Polars会自动为右侧DataFrame的列添加后缀（如_right）。本例中，由于我们将非索引列打包成结构体，并分别命名为cols和cols_right，避免了直接的列名冲突。
• 性能： 对于大型数据集，这种基于表达式和结构体的方法通常比迭代或使用apply函数更具性能优势，因为它允许Polars优化底层的计算。

通过理解和运用pl.struct、join和unnest这些核心概念，Polars用户可以高效且灵活地处理复杂的跨DataFrame操作，充分发挥其在数据处理方面的强大能力。