Polars列式除法优化技巧分享

本文深入探讨了Polars中DataFrame列式除法的高效操作技巧，着重解决了使用单行DataFrame对另一DataFrame进行列式除法的性能瓶颈。传统方法通过重复构建大型DataFrame进行除法，效率低下且消耗大量内存。文章详细介绍了利用`with_columns`结合字典推导式和列表达式的优化方案，该方案能够避免不必要的内存复制，显著提升性能。通过实例代码演示，展示了如何简洁高效地实现列式除法，是Polars数据处理中一项重要的最佳实践，尤其适用于处理大规模数据集，提升数据处理速度并降低资源消耗。

Polars中DataFrame列式除法概述

在数据处理中，我们经常需要对DataFrame的每一列或每一行应用特定的操作。当需要将一个DataFrame的行（或列）除以一组特定的数值时，如果这些数值来源于一个单行（或单列）的DataFrame，如何高效地实现这一操作就成为了一个常见问题。特别是在处理大型数据集时，性能和内存效率是至关重要的考量因素。

低效的实现方式：重复与拼接

在Polars中，如果直接尝试将一个DataFrame与一个单行DataFrame进行除法运算，Polars的广播机制默认不会直接将单行DataFrame的每列值广播到目标DataFrame的对应整列。一种直观但效率低下的方法是手动将单行DataFrame重复多次，使其行数与目标DataFrame相同，然后再进行元素级的除法。

考虑以下场景：我们有一个包含多行数据的DataFrame df，以及一个包含除数信息的单行DataFrame divisors。

import polars as pl
from itertools import repeat

# 示例数据
data = {'a': [i for i in range(1, 5)],
        'b': [i for i in range(1, 5)],
        'c': [i for i in range(1, 5)],
        'd': [i for i in range(1, 5)]}
df = pl.DataFrame(data)

# 单行除数DataFrame
divisors = pl.DataFrame({'d1': 1, 'd2': 10, 'd3': 100, 'd4': 1000})

print("原始DataFrame (df):")
print(df)
print("\n除数DataFrame (divisors):")
print(divisors)

输出：

原始DataFrame (df):
shape: (4, 4)
┌─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 1   ┆ 1   │
│ 2   ┆ 2   ┆ 2   ┆ 2   │
│ 3   ┆ 3   ┆ 3   ┆ 3   │
│ 4   ┆ 4   ┆ 4   ┆ 4   │
└─────┴─────┴─────┴─────┘

除数DataFrame (divisors):
shape: (1, 4)
┌─────┬─────┬─────┬──────┐
│ d1  ┆ d2  ┆ d3  ┆ d4   │
│ --- ┆ --- ┆ --- ┆ ---  │
│ i64 ┆ i64 ┆ i64 ┆ i64  │
╞═════╪═════╪═════╪══════╡
│ 1   ┆ 10  ┆ 100 ┆ 1000 │
└─────┴─────┴─────┴──────┘

为了使 divisors DataFrame的行数与 df 匹配，我们可以手动复制 divisors 并进行拼接：

# 低效方法：重复并拼接除数DataFrame
divisors_as_big_as_df = pl.concat([item for item in repeat(divisors, len(df))])
divided_df_inefficient = df / divisors_as_big_as_df

print("\n重复后的除数DataFrame (divisors_as_big_as_df):")
print(divisors_as_big_as_df)
print("\n低效方法得到的除法结果 (divided_df_inefficient):")
print(divided_df_inefficient)

输出：

重复后的除数DataFrame (divisors_as_big_as_df):
shape: (4, 4)
┌─────┬─────┬─────┬──────┐
│ d1  ┆ d2  ┆ d3  ┆ d4   │
│ --- ┆ --- ┆ --- ┆ ---  │
│ i64 ┆ i64 ┆ i64 ┆ i64  │
╞═════╪═════╪═════╪══════╡
│ 1   ┆ 10  ┆ 100 ┆ 1000 │
│ 1   ┆ 10  ┆ 100 ┆ 1000 │
│ 1   ┆ 10  ┆ 100 ┆ 1000 │
│ 1   ┆ 10  ┆ 100 ┆ 1000 │
└─────┴─────┴──────┴──────┘

低效方法得到的除法结果 (divided_df_inefficient):
shape: (4, 4)
┌─────┬─────┬──────┬───────┐
│ a   ┆ b   ┆ c    ┆ d     │
│ --- ┆ --- ┆ ---  ┆ ---   │
│ f64 ┆ f64 ┆ f64  ┆ f64   │
╞═════╪═════╪══════╪═══════╡
│ 1.0 ┆ 0.1 ┆ 0.01 ┆ 0.001 │
│ 2.0 ┆ 0.2 ┆ 0.02 ┆ 0.002 │
│ 3.0 ┆ 0.3 ┆ 0.03 ┆ 0.003 │
│ 4.0 ┆ 0.4 ┆ 0.04 ┆ 0.004 │
└─────┴─────┴──────┴───────┘

这种方法虽然能得到正确的结果，但其缺点显而易见：当 df 包含大量行时，divisors_as_big_as_df 会占用大量的内存，并且 pl.concat 操作本身也可能非常耗时，严重影响性能。

高效的解决方案：利用with_columns进行列式操作

Polars提供了更高效的机制来处理这类问题，即通过with_columns结合列表达式。我们可以遍历目标DataFrame的每一列，然后将该列与divisors DataFrame中对应列的单个值进行除法运算。Polars的表达式引擎能够智能地将这个单值广播到整列。

# 高效方法：使用with_columns进行列式除法
divided_df_efficient = df.with_columns(
    # 使用字典推导式为每一列生成新的表达式
    **{col: pl.col(col) / divisors[f"d{i+1}"]
       for (i, col) in enumerate(df.columns)}
)

print("\n高效方法得到的除法结果 (divided_df_efficient):")
print(divided_df_efficient)

输出：

高效方法得到的除法结果 (divided_df_efficient):
shape: (4, 4)
┌─────┬─────┬──────┬───────┐
│ a   ┆ b   ┆ c    ┆ d     │
│ --- ┆ --- ┆ ---  ┆ ---   │
│ f64 ┆ f64 ┆ f64  ┆ f64   │
╞═════╪═════╪══════╪═══════╡
│ 1.0 ┆ 0.1 ┆ 0.01 ┆ 0.001 │
│ 2.0 ┆ 0.2 ┆ 0.02 ┆ 0.002 │
│ 3.0 ┆ 0.3 ┆ 0.03 ┆ 0.003 │
│ 4.0 ┆ 0.4 ┆ 0.04 ┆ 0.004 │
└─────┴─────┴──────┴───────┘

代码解析与优势

1. df.with_columns(...): 这是Polars中用于添加或修改DataFrame列的核心方法。它接受一系列表达式作为参数，每个表达式定义了新列的计算逻辑。
2. 字典推导式 {col: expression for ...}: 我们使用字典推导式来动态地为 df 中的每一列生成一个修改表达式。字典的键是原始列的名称 (col)，值是对应列的计算表达式。** 操作符用于将字典解包为关键字参数传递给 with_columns。
3. pl.col(col): 这代表了当前正在处理的 df 中的那一列。
4. divisors[f"d{i+1}"]: 这是实现高效广播的关键。
   • enumerate(df.columns) 遍历 df 的列名及其索引。
   • f"d{i+1}" 根据索引动态构建 divisors DataFrame中的列名（例如，当 i 为0时，对应 d1；当 i 为1时，对应 d2，以此类推）。
   • divisors[...] 表达式会从 divisors DataFrame中提取指定列的数据。由于 divisors 是一个单行DataFrame，divisors[f"d{i+1}"] 实际上会返回一个包含单个值的Polars Series（或在表达式上下文中是一个表示该单值的表达式）。
   • 当一个 pl.col() 表达式与一个包含单个值的Series（或对应的表达式）进行算术运算时，Polars的查询优化器会智能地将这个单值广播到 pl.col() 所代表的整个列上，而无需在内存中实际复制该值。这极大地减少了内存消耗和计算开销。

这种方法的优势包括：

• 高性能: 避免了显式地创建和操作大型中间DataFrame，Polars的内部优化器能够高效地执行列式操作。
• 内存效率: 无需为除数创建与目标DataFrame相同大小的副本，显著减少了内存占用。
• 代码简洁: 使用字典推导式使得代码更加紧凑和易读。

注意事项

• 列名匹配: 上述解决方案假设 divisors DataFrame的列名（如 d1, d2）与 df 的列顺序相对应。如果列名不匹配或对应关系更复杂，您可能需要调整 f"d{i+1}" 的逻辑，或者使用一个映射字典来明确指定 df 列与 divisors 列的对应关系。
• 数据类型: 除法运算通常会导致整数类型转换为浮点数类型，以保留小数部分。Polars会自动处理这种类型提升。
• 零除错误: 如果 divisors 中包含零值，将导致除以零的错误。在实际应用中，需要考虑如何处理这种情况，例如使用 pl.Expr.fill_nan() 或 pl.Expr.replace() 进行错误处理。

总结

在Polars中对DataFrame进行列式除法，尤其是当除数来源于一个单行DataFrame时，最推荐的方法是利用 with_columns 结合字典推导式和Polars的表达式系统。这种方法不仅能够提供卓越的性能和内存效率，还能使代码更加清晰和易于维护。通过避免不必要的DataFrame复制和拼接操作，我们可以充分发挥Polars在处理大规模数据时的强大能力。

到这里，我们也就讲完了《Polars列式除法优化技巧分享》的内容了。个人认为，基础知识的学习和巩固，是为了更好的将其运用到项目中，欢迎关注golang学习网公众号，带你了解更多关于的知识点！