DataFrame单行除法技巧详解

本文深入探讨了在Polars中高效执行DataFrame行与单行DataFrame除法操作的技巧。传统方法通过复制单行DataFrame以匹配行数，造成了显著的性能瓶颈。本文提出一种更优方案，利用Polars的`with_columns`方法，结合列迭代和表达式，实现了高性能、内存友好的元素级除法。该方法避免了创建大型中间DataFrame，显著提升了数据处理效率。通过实例演示，对比了低效的复制方法和高效的逐列操作方法，详细阐述了Polars表达式系统和广播机制在优化数据处理中的关键作用。掌握这些技巧，能帮助读者在Polars中编写出更高效、更节省资源的数据处理代码，尤其是在处理大规模数据集时，优势更为明显。

问题背景：DataFrame与单行DataFrame的除法

在数据处理中，我们经常需要对整个DataFrame的每一行或每一列应用一组特定的操作。当这组操作涉及与一个固定单行值进行元素级运算时（例如，除法），如何在Polars中高效实现是一个常见需求。与Pandas中可能通过df.divide实现类似功能不同，Polars在DataFrame级别的隐式广播行为有所不同，直接对两个形状不匹配的DataFrame执行数学运算通常会导致错误或需要显式对齐。

考虑以下场景：我们有一个包含多行数据的Polars DataFrame，以及一个包含对应列除数的单行Polars DataFrame。目标是将主DataFrame的每一列与其对应的除数进行元素级除法。

import polars as pl
from itertools import repeat

# 原始数据DataFrame
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 │
└─────┴─────┴──────┴──────┘

低效的解决方案：复制单行DataFrame

一种直观但效率低下的方法是，将单行divisors DataFrame复制多次，使其行数与主df DataFrame相同，然后直接执行除法。

# 低效方法：复制单行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这个大型中间DataFrame会消耗大量的内存和计算时间，效率极低。

高效的解决方案：逐列操作与with_columns

Polars的设计哲学鼓励利用其表达式系统和延迟计算（Lazy Evaluation）来优化性能。对于这种场景，最佳实践是避免创建大型中间DataFrame，而是通过迭代列并利用with_columns方法进行逐列除法。

with_columns方法允许我们一次性添加或修改多个列。我们可以构建一个字典，其中键是目标列名（即df的列名），值是对应的Polars表达式，表示该列除以divisors DataFrame中对应的单值。

# 高效方法：使用 with_columns 逐列除法
# 注意：这里假设 df 的列名与 divisors 的列名可以通过索引或某种映射关系对应。
# 在本例中，df 的列是 'a', 'b', 'c', 'd'，divisors 的列是 'd1', 'd2', 'd3', 'd4'。
# 我们通过 enumerate(df.columns) 和 f"d{i+1}" 来建立这种对应关系。
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. 避免中间DataFrame： 核心优势在于divisors[f"d{i+1}"]直接从单行divisors DataFrame中提取出单个Series（或Scalar值），Polars能够将其自动广播到pl.col(col)对应的整个列进行运算。这避免了创建与原始DataFrame等大的中间DataFrame，显著节省了内存和计算资源。
2. Polars表达式系统： pl.col(col)创建了一个指向指定列的表达式，而divisors[f"d{i+1}"]则直接获取了该列的Series。Polars的优化器能够高效地处理这种列与Series之间的运算，尤其是在Lazy模式下，可以进一步优化执行计划。
3. 列名映射： 示例中通过enumerate(df.columns)和f-string f"d{i+1}"来动态构建divisors DataFrame中的列名，这要求df和divisors的列顺序或命名规则能够建立起明确的对应关系。如果列名不规律，可能需要一个显式的映射字典来指定df的哪些列应该除以divisors的哪些列。
4. 数据类型： 除法操作通常会导致结果变为浮点数类型（f64），即使原始数据是整数类型。Polars会自动处理这种类型转换。

总结

在Polars中对DataFrame进行行与单行DataFrame的元素级除法时，最推荐且高效的方法是利用with_columns结合列迭代和Polars表达式。这种方法通过避免创建大型冗余的中间DataFrame，充分利用了Polars的内存优化和查询优化能力，从而实现高性能的数据处理。理解Polars的表达式系统和广播机制对于编写高效的数据处理代码至关重要。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《DataFrame单行除法技巧详解》文章吧，也可关注golang学习网公众号了解相关技术文章。