Polars LazyFrame多列乘法：跳过索引列的高效策略-Python教程-PHP中文网

Polars LazyFrame多列乘法：跳过索引列的高效策略

花韻仙語

发布： 2025-11-20 13:34:29

原创

714人浏览过

Polars LazyFrame多列乘法：跳过索引列的高效策略

本文详细介绍了在polars中对两个lazyframe进行列式乘法运算的高效方法，尤其是在需要排除特定索引列时。通过利用`pl.struct`将非索引列封装成结构体、使用`join`操作对齐数据，以及直接对结构体进行乘法运算，最后通过`unnest`展开结果，实现了类似于pandas的直观操作，同时保持了polars的惰性计算优势。

在数据处理中，经常需要对两个具有相同结构的数据集进行元素级别的运算，例如逐列相乘。对于Pandas用户而言，这通常通过简单的算术运算符即可实现，例如df1 * df2，其中索引列会自动对齐并排除在乘法之外。然而，在Polars的LazyFrame环境中执行类似操作，尤其是在需要排除某个“索引”列（如时间戳）时，需要一种更具Polars特色的策略。

问题场景与Polars的挑战

假设我们有两个Polars LazyFrame，它们都包含一个时间戳列（time）和多个数值列（foo, bar, baz），结构如下：

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

# 为了重现性，设置随机种子
np.random.seed(42)
n = 5 # 示例数据行数

# 创建第一个Polars 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()

# 创建第二个Polars 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 (LazyFrame):")
print(df1.collect())
print("\ndf2 (LazyFrame):")
print(df2.collect())

登录后复制

直接尝试使用concat和group_by进行聚合操作，例如pl.concat([df1, df2]).group_by('time').sum()对于求和是有效的。然而，Polars的product聚合函数通常用于计算单个列的乘积，并且直接将agg(pl.col("*").mul(pl.col("*")))应用于group_by会产生包含列表的列，这并非我们期望的列式乘法结果。

# 尝试使用concat和group_by进行乘法聚合，但结果不符合预期
# result_agg_mul = pl.concat([df1, df2]).group_by('time').agg(pl.col("*").mul(pl.col("*"))).sort('time').collect()
# print("\n尝试聚合乘法 (不符合预期):")
# print(result_agg_mul)

登录后复制

上述尝试会生成列表列，因为group_by后agg中的pl.col("*")会分别代表每个组内的所有列，mul操作会作用于这些列的序列，而非跨DataFrame的对应列。

Polars LazyFrame多列乘法的高效策略

为了在Polars LazyFrame中实现列式乘法，同时排除特定的“索引”列（如time），我们可以利用pl.struct、join和unnest的组合。

步骤一：将非索引列封装为结构体（Struct）

首先，将每个LazyFrame中除了“索引”列之外的所有数值列封装到一个新的struct列中。pl.struct(pl.exclude("time"))表达式可以方便地选择除time之外的所有列并将其组合成一个结构体。

# 将df1的非时间列封装到名为'data_struct'的结构体中
df1_struct = df1.select("time", data_struct=pl.struct(pl.exclude("time")))
print("\ndf1_struct (LazyFrame):")
print(df1_struct.collect())

# 将df2的非时间列封装到名为'data_struct'的结构体中
df2_struct = df2.select("time", data_struct=pl.struct(pl.exclude("time")))
print("\ndf2_struct (LazyFrame):")
print(df2_struct.collect())

登录后复制

通过这一步，我们将每个LazyFrame的多个数值列逻辑上合并成了一个单一的struct类型列，其内部包含了原始的foo, bar, baz字段。

步骤二：通过time列进行连接（Join）

接下来，使用time列作为键，对这两个包含结构体的LazyFrame进行左连接（left join）。这将确保两个LazyFrame中对应时间戳的数据行能够对齐。

Logome

AI驱动的Logo生成工具

133

查看详情

# 基于'time'列进行左连接
joined_df = df1_struct.join(
    df2_struct,
    on="time",
    how="left",
    suffix="_right" # 为右侧DataFrame的结构体列添加后缀以区分
)
print("\nJoined LazyFrame (after struct creation and join):")
print(joined_df.collect())

登录后复制

连接后的LazyFrame将包含time列、来自df1_struct的data_struct列，以及来自df2_struct的data_struct_right列。

步骤三：对结构体列执行乘法运算

Polars允许直接对结构体列进行算术运算。当两个结构体列相乘时，Polars会自动执行元素级别的乘法，即结构体中同名字段的对应值相乘。

# 对结构体列执行乘法，并选择结果
multiplied_struct_df = joined_df.select(
    "time",
    product_data=pl.col("data_struct") * pl.col("data_struct_right")
)
print("\nLazyFrame after multiplying struct columns:")
print(multiplied_struct_df.collect())

登录后复制

这一步产生了新的product_data结构体列，其中包含了foo, bar, baz字段的乘积结果。

步骤四：展开结构体（Unnest）

最后，使用unnest()方法将包含乘积结果的product_data结构体列展开，恢复成独立的列，从而得到最终的、与Pandas操作结果类似的LazyFrame。

# 展开结构体列，得到最终结果
final_result_df = multiplied_struct_df.unnest("product_data")
print("\nFinal result (after unnesting):")
print(final_result_df.collect())

登录后复制

完整代码示例

将上述所有步骤整合起来，实现Polars LazyFrame的列式乘法操作：

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

# 为了重现性，设置随机种子
np.random.seed(42)
n = 5 # 示例数据行数

# 创建第一个Polars 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()

# 创建第二个Polars 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()

# 核心操作链
result_lazyframe = (
    df1.select("time", data_struct=pl.struct(pl.exclude("time"))) # df1非时间列结构化
    .join(
        df2.select("time", data_struct=pl.struct(pl.exclude("time"))), # df2非时间列结构化
        on="time",
        how="left",
        suffix="_right" # 为右侧的结构体列添加后缀
    )
    .select(
        "time",
        product_data=pl.col("data_struct") * pl.col("data_struct_right") # 结构体相乘
    )
    .unnest("product_data") # 展开结果结构体
)

# 收集并打印最终结果
final_df = result_lazyframe.collect()
print("\n最终的Polars LazyFrame列式乘法结果:")
print(final_df)

登录后复制

注意事项与总结

pl.struct的灵活性： pl.struct是一个非常强大的工具，它允许我们将多列视为一个逻辑单元进行操作。这在处理复杂的数据转换或聚合时尤其有用，例如在本例中，它使得对非索引列进行整体操作成为可能。
join的重要性： 在对多个LazyFrame进行操作时，确保数据正确对齐至关重要。join操作是实现这一目标的核心机制，它根据指定的键（time列）将来自不同LazyFrame的数据行关联起来。
惰性计算优势： 整个操作链都是在LazyFrame上执行的，这意味着Polars会构建一个查询计划，并在collect()被调用时才实际执行计算。这对于处理大规模数据集时提供了显著的性能优势和内存效率。
列名管理： 在join操作中，使用suffix参数可以避免连接后出现同名列的冲突。这里我们为右侧的结构体列添加了_right后缀，以便在后续操作中明确引用。
可读性： 尽管这种方法比Pandas的直接乘法看起来更复杂，但它在Polars的表达力框架内是清晰且高效的。一旦理解了struct和unnest的用途，这种模式就会变得非常直观。