使用 Polars 替代 Pandas：性能深度解析

#介绍

在过去十年中，[Pandas](https://pandas.pydata.org/) 一直是 Python 数据工作的基础。对于可以放入内存的数据集，它的速度和熟悉度足以使大多数程序员不会考虑更换库。

然而，一旦开始处理数百万行数据，问题就开始显现：耗时几秒的 groupby 操作，消耗大量 RAM 的中间副本，以及以 Python 级别循环而不是向量化 [C](https://www.cprogramming.com/) 或 [Rust](https://rust-lang.org/) 代码运行的窗口函数。

[Polars](https://pola.rs/) 是一个用 Rust 构建在 [Apache Arrow](https://arrow.apache.org/) 之上的 DataFrame 库。它设计时就将并行性和惰性求值作为核心特性。Pandas 每次操作都会立即顺序执行，而 Polars 可以构建查询计划并在执行前进行优化，大多数操作会自动跨所有可用 CPU 核心并发执行。

在本文中，我们将使用来自 [StrataScratch](https://platform.stratascratch.com/) 编码平台的真实问题来探讨三个实际数据问题。对于每个问题，我们比较两个库的解决方案，并指出性能差异最显著的地方。

#使用 `rank()` 与 `with_row_count()`: 活动排名

在 这个问题 中，目标是根据发送的总邮件数量为每个用户找到其活动排名。发送邮件最多的用户获得排名 1。结果必须按总邮件数量降序排序，使用字母顺序作为平局决胜条件，并且每个排名必须唯一，即使两个用户的邮件数量相同。

#### //数据视图

google_gmail_emails 表存储每封已发送邮件的一行记录，包括发件人 ID (from_user)、收件人 ID (to_user) 和邮件发送日期。以下是该表的预览：

| id | from_user | to_user | day | | --- | --- | --- | --- | | 0 | 6edf0be4b2267df1fa | 75d295377a46f83236 | 10 | | 1 | 6edf0be4b2267df1fa | 32ded68d89443e808 | 6 | | 2 | 6edf0be4b2267df1fa | 55e60cfcc9dc49c17e | 10 | | 3 | 6edf0be4b2267df1fa | e0e0defbb9ec47f6f7 | 6 | | 4 | ... | ... | ... | | 314 | e6088004caf0c8cc51 | e6088004caf0c8cc51 | 5 |

粒度（一行输出的意义）：一个用户，包含其总邮件数量和唯一的活动排名。

#### //常见错误

问题要求即使两个用户的邮件数量相同，排名也必须唯一。一个常见的错误是在 Pandas 中使用 rank(method='dense') 方法，这会导致平局的用户获得相同的排名。正确的方法是 'first'，它通过排序中的位置来打破平局。由于我们在排名之前按 user_id 字母顺序排序，因此最终的排名是唯一且确定的。

Polars 的最优解决方案完全避免了 rank 函数。在按 ["total_emails", "user_id"] 分别降序和升序排序后，.with_row_count("activity_rank", offset=1) 子句从 1 开始分配连续整数。由于排序已经处理了平局，因此不需要额外的平局决胜逻辑。

#### //解决方案

1. Pandas 解决方案

我们将 from_user 重命名为 user_id，按用户分组，统计邮件数量，计算首次排名，并按邮件数量降序排序，使用字母顺序打破平局。

import pandas as pd
import numpy as np
google_gmail_emails = google_gmail_emails.rename(columns={"from_user": "user_id"})
result = google_gmail_emails.groupby(
    ['user_id']).size().to_frame('total_emails').reset_index()
result['activity_rank'] = result['total_emails'].rank(method='first', ascending=False)
result = result.sort_values(by=['total_emails', 'user_id'], ascending=[False, True])

2. Polars 解决方案

我们使用一个 惰性链，在一次遍历中重命名、分组、排序并分配行号。最后调用 .collect() 来生成结果。

import polars as pl
google_gmail_emails = google_gmail_emails.rename({"from_user": "user_id"})
result = (
    google_gmail_emails.lazy()
    .group_by("user_id")
    .agg(total_emails = pl.count())
    .sort(
        by=["total_emails", "user_id"],
        descending=[True, False]
    )
    .with_row_count("activity_rank", offset=1)
    .select([
        pl.col("user_id"),
        "total_emails",
        "activity_rank"
    ])
    .collect()
)

#### //性能对比

Pandas 解决方案在分组后会遍历数据两次：一次用于计算大小，另一次用于分配排名。内部地，rank(method='first') 分配了一个排名数组，通过 argsort 解决平局，并写回结果——对于单个列来说，这比看起来要昂贵得多。Polars 的 group_by 函数将工作负载分配到所有可用的 CPU 核心上，从而显著加快了大型表的聚合速度。由于 .with_row_count() 子句是在排序后的一次 O(n) 顺序遍历，它用最便宜的操作替换了排名函数。在一个包含数百万条电子邮件记录的表中，使用并行聚合而不使用排名函数可以比 Pandas 方法在实际时间上提高 5-10 倍的性能。

以下是代码输出预览：

| user_id | total_emails | activity_rank | | --- | --- | --- | | 32ded68d89443e808 | 19 | 1 | | ef5fe98c6b9f313075 | 19 | 2 | | 5b8754928306a18b68 | 18 | 3 | | 55e60cfcc9dc49c17e | 16 | 4 | | 91f59516cb9dee1e88 | 16 | 5 | | ... | ... | ... | | e6088004caf0c8cc51 | 6 | 25 |

#使用 cumcount() + pivot() vs. over(): 查找用户购买记录

在 这个问题 中，我们需要识别返回活跃用户——具体来说，那些在首次购买后 1 到 7 天内进行了第二次购买的用户。同一天内的购买不应包括在内。结果只是一个符合条件的 user_id 列表。

#### //数据视图

amazon_transactions 表每条购买记录有一行，包含 user_id、item、created_at 日期和 revenue。

以下是表的预览：

| id | user_id | item | created_at | revenue | | --- | --- | --- | --- | --- | | 1 | 109 | milk | 2020-03-03 | 123 | | 2 | 139 | biscuit | 2020-03-18 | 421 | | 3 | 120 | milk | 2020-03-18 | 176 | | ... | ... | ... | ... | ... | | 100 | 117 | bread | 2020-03-10 | 209 |

粒度（一个输出行的含义）：一个在首次购买后 7 天内进行了符合条件的再次购买的用户 ID。

#### //边缘情况

同一天的购买应忽略不计，这意味着第一次和第二次购买之间的间隔必须超过 0 天且最多为 7 天。一个在同一天购买两次的客户不符合条件。

#### //解决方案

两种解决方案都找到了每个用户的最早购买日期，然后过滤出在 1 到 7 天时间窗口内的后续购买。需要注意的一点是：如果 created_at 包含时间戳而不是纯日期，你需要在比较之前截断到日期。否则，同一天不同时间的两次购买会错误地通过严格的不等式。

1. Pandas 解决方案

在 Pandas 中，解决方案涉及隔离每个用户的唯一购买日期，使用 cumcount() 对其进行排名，通过 pivot 将第一个和第二个日期并排放置，并计算天数差异。

import pandas as pd
amazon_transactions["purchase_date"] = pd.to_datetime(amazon_transactions["created_at"]).dt.date
daily = amazon_transactions[["user_id", "purchase_date"]].drop_duplicates()
ranked = daily.sort_values(["user_id", "purchase_date"])
ranked["rn"] = ranked.groupby("user_id").cumcount() + 1
first_two = (ranked[ranked["rn"] <= 2]
             .pivot(index="user_id", columns="rn", values="purchase_date")
             .reset_index()
             .rename(columns={1: "first_date", 2: "second_date"}))
first_two = first_two.dropna(subset=["second_date"])
first_two["diff"] = (pd.to_datetime(first_two["second_date"]) - pd.to_datetime(first_two["first_date"])).dt.days
result = first_two[(first_two["diff"] >= 1) & (first_two["diff"] <= 7)][["user_id"]]

2. Polars 解决方案

Polars 解决方案涉及使用 .over("user_id") 作为窗口表达式计算每个用户的首次购买日期，过滤出符合时间窗口的购买记录，并返回去重的 user_id 列表。

import polars as pl
# 返回活跃用户：第二次购买在首次购买后 1-7 天内（忽略同一天）
returning_users = (
    amazon_transactions
    .lazy()
    # 每个用户的首次购买日期（使用窗口表达式以避免在 LazyFrame 上使用 .groupby）
    .with_columns(
        pl.col("created_at").min().over("user_id").alias("first_purchase_date")
    )
    # 保留严格在首次购买后 1-7 天内的交易
    .filter(
        (pl.col("created_at") > pl.col("first_purchase_date")) &
        (pl.col("created_at") <= pl.col("first_purchase_date") + pl.duration(days=7))
    )
    # 去重用户列表
    .select("user_id")
    .unique()
    .sort("user_id", descending=[False])
)

#### //性能对比

请注意 Pandas 解决方案中的不同 DataFrame 分配数量：去重的每日表、排序的排名表、透视表、dropna 结果和过滤后的输出。这些由五个单独的对象组成，每个对象都会将数据复制到新的内存块中。在大型交易表中，仅透视步骤就可能显著增加内存使用量，因为它将整个数据集重塑为宽格式。

Polars 的懒惰链在 .collect() 之前不会分配任何内存。.over("user_id") 窗口表达式在一次遍历中计算每个用户的最早购买日期，.filter() 立即在同一步骤中应用，而 .unique() 在多个 CPU 核心上并发运行。没有透视操作，没有中间的排序副本，也没有单独的日期转换步骤——Polars 在表达式引擎内部原生处理日期运算。这种方法消耗更少的内存，并且即使在中等规模的数据集上也能更快运行。

以下是代码输出预览：

| user_id | | --- | | 100 | | 103 | | 105 | | ... | | 143 |

#使用 expanding().mean() vs. cum_mean(): 月销售滚动平均值

在这个问题中，我们需要计算 2022 年每月书籍销售的累计平均值。平均值每个月都会增加，使用所有之前的月份数据：2 月平均 1 月和 2 月的数据，3 月平均 1 月、2 月和 3 月的数据，依此类推。输出应包括月份、该月的总销售额以及四舍五入到最接近整数的累计平均值。

#### //数据视图

amazon_books 表包含每本书及其单价的一行数据。book_orders 表包含每个订单的一行数据，将书 ID 链接到数量和订单日期。以下是表的预览：

| book_id | book_title | unit_price | | --- | --- | --- | | B001 | 饥饿游戏 | 25 | | B002 | 外乡人 | 50 | | B003 | 杀死一只知更鸟 | 100 | | ... | ... | ... | | B020 | 地球之柱 | 60 |

book_orders 表包含每个书籍订单的一行数据，将每个订单 ID 链接到订单日期、书 ID 和订购数量：

| order_id | order_date | book_id | quantity | | --- | --- | --- | --- | | 1001 | 2022-01-10 | B001 | 1 | | 1002 | 2022-01-10 | B009 | 1 | | 1003 | 2022-01-15 | B012 | 2 | | ... | ... | ... | ... | | 1084 | 2023-02-01 | B009 | 1 |

粒度（一行输出的意义）：2022 年的一个月，该月的总销售额以及截至该月的所有月销售额的累计平均值。

#### //权衡

使用 Pandas，.expanding().mean() 子句非常方便，但在内部使用 Python 级别的循环处理增长的窗口切片。对于 12 行的月度汇总，这种成本是可以忽略的。但对于大规模的每日或每小时数据（例如，三年的每小时交易），每个扩展窗口切片都会增加逐行累积的开销。

Polars 的 cum_mean() 在 Rust 中运行单次遍历，因此在大规模数据处理上更快。但有一个缺点：问题要求四舍五入到最接近的整数，而 Pandas 默认使用银行家舍入（四舍六入五成双）。Polars 解决方案使用 NumPy 的 cumsum 并显式使用 floor(x + 0.5) 公式来强制执行四舍五入行为。如果需要与预期输出完全匹配，NumPy 方法比两个库中的内置舍入方法更可靠。

#### //解决方案

1. Pandas 解决方案

我们将书籍与订单合并，筛选出 2022 年的数据，按月汇总销售额，并应用 .expanding().mean() 计算累计平均值。

import pandas as pd
import numpy as np
import datetime as dt
merged = pd.merge(book_orders, amazon_books, on="book_id", how="inner")
merged["order_date"] = pd.to_datetime(merged["order_date"])
merged["order_month"] = merged["order_date"].dt.month
merged["year"] = merged["order_date"].dt.year
merged["sales"] = merged["unit_price"] * merged["quantity"]
merged = merged.loc[(merged["year"] == 2022), :]
result = (
    merged.groupby("order_month")["sales"]
    .sum()
    .to_frame("monthly_sales")
    .sort_values(by="order_month")
    .reset_index()
)
result["rolling_average"] = result["monthly_sales"].expanding().mean().round(0)
result

2. Polars：构建惰性管道并收集

我们在惰性链中连接两个表，计算销售额为 unit_price * quantity，筛选出 2022 年的数据，按月汇总，并调用 .collect() 切换到急切模式，以便进行 NumPy 滚动计算。

import polars as pl
import numpy as np
# 步骤 1：准备月度销售额（惰性框架）
monthly_sales_lazy = (
    book_orders.lazy()
    .join(amazon_books.lazy(), on="book_id", how="inner")
    .with_columns([
        (pl.col("unit_price") * pl.col("quantity")).alias("sales"),
        pl.col("order_date").cast(pl.Datetime),
        pl.col("order_date").dt.year().alias("year"),
        pl.col("order_date").dt.month().alias("order_month")
    ])
    .filter(pl.col("year") == 2022)
    .group_by("order_month")
    .agg(pl.col("sales").sum().alias("monthly_sales"))
    .sort("order_month")
)
# 步骤 2：切换到急切模式进行滚动计算
monthly_sales = monthly_sales_lazy.collect()

3. 计算滚动平均值并最终确定

将月度销售额转换为 NumPy 数组后，我们应用四舍五入到最接近整数的舍入规则，将结果重新添加到 Polars DataFrame 中，并选择输出列。

# 步骤 3：带有四舍五入的滚动平均值
sales_np = monthly_sales["monthly_sales"].to_numpy()
cumsum = np.cumsum(sales_np)
rolling_avg = np.floor(cumsum / np.arange(1, len(cumsum)+1) + 0.5).astype(int)
# 步骤 4：重新添加到 Polars DataFrame
monthly_sales = monthly_sales.with_columns([
    pl.Series("rolling_average", rolling_avg)
])
# 步骤 5：最终结果，带有正确的列名
result = monthly_sales.select(["order_month", "monthly_sales", "rolling_average"])

#### //性能比较

这个问题有两个对性能影响最大的操作：连接和累计窗口。在 Pandas 中，pd.merge 会先将两个表中的所有行连接起来，然后再筛选出 2022 年的数据。这意味着在丢弃目标期间之外的行之前，会处理每一年的订单数据。Polars 构建了一个惰性查询计划，并在连接执行前推送 filter(year == 2022) 条件，因此从一开始就连接了较小的数据集。这种谓词下推是自动发生的，无需额外编写代码。

最明显的区别在于滚动平均值的差距。Pandas 的 .expanding().mean() 方法每次增加一行窗口，并在每个段调用 C 代码，但整个过程由 Python 循环控制。而 Polars 的 cum_mean() 方法则在一个 Rust 循环中计算整个列，没有 Python 开销。虽然在月度数据上这种差异可能不易察觉，但如果对三年的每日数据（大约 1,000 行）运行相同的查询，Polars 版本在微秒内完成，而 Pandas 由于扩展窗口的存在显示出可测量的延迟。

以下是代码输出预览：

| order_month | sales | rolling_average | | --- | --- | --- | | 1 | 145 | 145 | | 2 | 250 | 198 | | 3 | 315 | 237 | | ... | ... | ... | | 12 | 710 | 402 |

结论

在所有三个问题中，Polars 的解决方案都遵循相同的模式：构建一个惰性查询计划，尽可能多地将计算推入优化器，并仅在需要具体结果时调用 .collect()。

如果你像大多数分析师一样有多年的 Pandas 使用习惯，语法上需要一些调整，但操作非常接近。.groupby() 变为 .group_by()，.rename() 接受一个普通的字典而不是 columns= 关键字，排名则变为排序后跟 .with_row_count()。

真正的区别在于大规模数据处理。在处理小数据集时，两个库都能快速返回结果，差异不明显。但当行数达到数百万时，Polars 的 Rust 级并行性和单次遍历算法显著优于 Pandas。如果你在使用 Pandas 时遇到性能问题，这三个挑战是一个很好的迁移起点。

[](https://twitter.com/StrataScratch)**[Nate Rosidi](https://twitter.com/StrataScratch)** 是一名数据科学家，从事产品战略工作。他还是教授分析课程的兼职教授，并且是 StrataScratch 的创始人，这是一个帮助数据科学家通过顶级公司的实际面试问题准备面试的平台。Nate 撰写关于职业市场最新趋势的文章，提供面试建议，分享数据科学项目，并涵盖一切与 SQL 相关的内容。