本教程旨在指导用户如何在polars lazyframes中实现列级乘法操作，尤其是在需要排除一个共同索引列的情况下。文章将详细介绍如何利用polars的`struct`表达式将非索引列封装起来，并通过`join`操作对齐两个lazyframes，进而执行结构体之间的元素级乘法，最后使用`unnest`将结果展开，从而提供一个高效且符合polars范式的解决方案。
在数据处理领域，尤其是在从Pandas等工具迁移到Polars时，用户经常会遇到需要对两个结构相同、具有共同索引的数据集进行元素级运算的场景。例如，对两个DataFrame中除了时间索引列之外的所有对应数值列执行乘法。虽然在Pandas中，这通常可以通过简单的df1 * df2实现，但在Polars的LazyFrame环境中，由于其优化和表达式驱动的特性，需要采用一种更具Polars风格的方法来达到相同的效果。
场景描述与Polars中的挑战
假设我们有两个Polars LazyFrame，它们都包含一个time列作为时间索引，以及若干数值列（例如foo, bar, baz）。我们的目标是让df1中的foo列乘以df2中的foo列，df1中的bar列乘以df2中的bar列，依此类推，同时保留time列作为结果的索引。
首先，我们创建示例数据：
import polars as pl
import numpy as np
import pandas as pd
n = 5 # 示例数据行数
# 创建第一个LazyFrame
df1 = 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()
# 创建第二个LazyFrame
df2 = 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("df1 示例数据:")
print(df1.collect())
print("\ndf2 示例数据:")
print(df2.collect())在Polars中，直接对LazyFrames执行df1 * df2并不能像Pandas那样自动进行基于列名的对齐和乘法。尝试使用pl.concat([df1, df2]).group_by('time').agg(pl.col("*").mul(pl.col("*")))等聚合操作，往往会因为聚合函数对列表的处理方式，导致结果不是我们期望的列级乘积，而是生成包含列表的复杂结构。
Polars 解决方案：结合 Structs 和 Join
Polars提供了一种强大且灵活的方式来处理这类跨DataFrame的列级操作，即利用结构体（struct）表达式和连接（join）操作。
整个解决方案可以分解为以下几个关键步骤：
1. 将非索引列封装为结构体
首先，我们需要将每个LazyFrame中除了time列之外的所有数值列封装到一个名为cols的结构体列中。这样做的好处是，可以将所有需要进行操作的列作为一个整体来处理，简化后续的乘法运算。
# 将df1的非时间列封装到结构体中
df1_struct = df1.select("time", cols=pl.struct(pl.exclude("time")))
# 将df2的非时间列封装到结构体中
df2_struct = df2.select("time", cols=pl.struct(pl.exclude("time")))
print("\ndf1 封装后的结构体示例:")
print(df1_struct.collect())通过pl.struct(pl.exclude("time"))表达式，我们指示Polars创建一个新的结构体列cols，其中包含除time列之外的所有其他列。
2. 通过索引列对齐两个结构体LazyFrame
接下来，我们需要将这两个包含结构体的LazyFrame基于共同的time列进行连接。这将确保来自df1和df2的对应行能够被正确地匹配起来。
# 基于time列进行左连接，对齐两个结构体LazyFrame
joined_df = df1_struct.join(
df2_struct,
on="time",
how="left"
)
print("\n连接后的LazyFrame示例:")
print(joined_df.collect())连接操作会生成一个包含time列、cols列（来自df1）和cols_right列（来自df2）的LazyFrame。此时，cols和cols_right是两个结构体列，它们在相同的time值上是相互对应的。
3. 执行结构体之间的元素级乘法
Polars允许直接对具有相同内部字段的结构体列执行元素级运算。这意味着我们可以直接将cols结构体列乘以cols_right结构体列，Polars会自动将其内部的对应字段进行乘法运算。
# 对结构体列执行元素级乘法
multiplied_struct_df = joined_df.select(
"time",
pl.col("cols") * pl.col("cols_right")
)
print("\n结构体乘法后的LazyFrame示例:")
print(multiplied_struct_df.collect())这一步会生成一个新的结构体列，其中每个字段的值都是原结构体中对应字段的乘积。
4. 展开结构体列以恢复原始列结构
最后一步是使用unnest()方法将乘法结果的结构体列展开，恢复成独立的列，从而得到与原始Pandas操作结果相似的扁平化DataFrame结构。
# 展开结构体列，得到最终结果
final_result = multiplied_struct_df.unnest("cols")
print("\n最终结果 LazyFrame:")
print(final_result.collect())完整代码示例
将上述步骤整合起来，我们得到以下完整的Polars LazyFrame列级乘法解决方案：
import polars as pl
import numpy as np
import pandas as pd
n = 5 # 示例数据行数
# 创建第一个LazyFrame
df1 = 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()
# 创建第二个LazyFrame
df2 = 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 LazyFrame列级乘法操作
result_df = (
df1.select("time", cols=pl.struct(pl.exclude("time"))) # 封装df1的非时间列
.join(
df2.select("time", cols=pl.struct(pl.exclude("time"))), # 封装df2的非时间列并连接
on = "time",
how = "left"
)
.select("time", pl.col("cols") * pl.col("cols_right")) # 对结构体列执行乘法
.unnest("cols") # 展开结果结构体
)
print("\n最终的Polars LazyFrame列级乘法结果:")
print(result_df.collect())注意事项与总结
- Polars范式: 这种利用struct和join的方法是Polars处理跨DataFrame复杂列级操作的典型范式。它在LazyFrame模式下表现出色，能够利用Polars的查询优化器进行高效计算。
- 灵活性: 这种方法不仅适用于乘法，也适用于其他元素级运算（加法、减法、除法等），只需将pl.col("cols") * pl.col("cols_right")替换为相应的运算符即可。
- 列名匹配: 确保两个LazyFrame中需要进行操作的列具有相同的名称，这样pl.exclude("time")才能正确地选择并封装对应的列。结构体之间的运算依赖于内部字段的名称匹配。
- 性能优势: 对于大型数据集，LazyFrame结合struct和join的策略能够避免不必要的中间计算和内存分配，从而提供显著的性能优势。
通过掌握这种技术，用户可以更自信地在Polars环境中实现复杂的跨DataFrame操作，充分发挥其高性能数据处理的能力。
