Polars 表达式 API 为何出色

本文由Vincent D. Warmerdam客座撰写。他是Rasa的研究倡导者、calmcode.io项目创建者，也是多个开源项目的维护者。他为Polars贡献了 .pipe() 方法，并且最近制作了面向初学者的Polars课程。

我最喜欢的数据集之一，是托管在Kaggle上的《魔兽世界》角色数据集。它包含2008年一个《魔兽世界》服务器的日志：系统每十分钟记录一次正在游戏中的部落阵营玩家。数据大约为644MB，足够小，可以在笔记本电脑上处理；也足够大，在使用Python分析时需要稍微注意效率。

下面是一小段数据：

任务

这份数据可以用于许多有趣的分析。原文提到，2008年是“Frozen Throne”扩展推出的年份，因此可以探索很多与用户流失相关的问题。但本文想探讨的是另一件事：机器人检测。这个游戏十分流行，吸引了许多作弊者，数据集很可能包含非人类玩家，需要把它们找出来。

检测机器人的方法很多，但先从简单的领域规则开始很合适。一个起点是寻找会话时长可疑地过长的用户。如果某个角色被观察到连续玩了36小时而没有休息，可以推测它可能是机器人。最终目标，是把去除机器人作为其他分析之前的预处理步骤。

怎样找到这些用户？在写代码前，先列出数据集需要哪些列：

• 添加表示 session_id 的列。每个ID唯一标识某位用户连续游戏的一次会话。同一个玩家在另一次会话中游戏时，应使用另一个ID。
• 有了 session_id 列，就可以计算每次会话持续多久。
• 有了每次会话的游戏时长，就可以查看每个角色最长的一次会话。如果曾超过某个阈值，例如24小时，就移除这个用户的全部活动。

划分会话

逐步展示如何为数据划分会话，可能更容易理解。假设有下面的数据集：

第一步是排序：先按角色，再按时间戳排序，结果如下。

接着添加两列，判断 timestamp 或 char 列是否发生了足以开始新会话的“跳变”。

再用逻辑或将两个“diff”列合并。

要把这列变成 session_id，只需对 diff 列调用 cumsum。

根据数据的分区方式，也可能有其他更高效的会话划分方法。但下文采用这里的方法。下一步计算会话长度，需要按会话聚合，再把结果附回原数据。

随后再进行一次聚合：这次按角色分组，计算最大会话长度。

这是常见需求

我们将编写一个完成上述任务的查询。但首先应认识到，这类查询非常常见。在Web分析中，先计算会话，再按用户汇总十分常见；它不仅用于机器人检测，也用于一般用户分析。

既然这么常见，应如何实现？令人意外的是，这类查询通常相当难写。可以使用先 .groupby() 再 .join() 的操作，但计算成本相对较高。更不用说，重复写多个这样的组合也很繁琐。

这正是Polars发挥优势的地方。它不仅拥有用Rust编写的快速实现，还提供了一套出色的表达式API。

流水线

回到任务：我们先划分会话，再基于会话计算统计信息。步骤较多，适合用流水线实现。思路是把代码写成职责分开的函数：每个函数接受DataFrame，变换后返回DataFrame，再串联成如下流水线。

(df
 .pipe(set_types)
 .pipe(sessionize, threshold=20 * 60 * 1000)
 .pipe(add_features)
 .pipe(remove_bots, threshold=24))

分离职责是很好的第一步，通常也是容易的部分。现在需要实现这些函数。下面逐一说明。

设置类型

分析之前，应确保数据类型正确。我们将guild列转换为布尔值，并确保时间戳采用正确类型。

def set_types(dataf):
    return (dataf
            .with_columns([
                 pl.col("timestamp").str.strptime(pl.Datetime, fmt="%m/%d/%y %H:%M:%S"),
                 pl.col("guild") != -1,
             ]))

添加会话

接下来添加会话。这涉及排序和添加若干列。

def sessionize(dataf, threshold=20 * 60 * 1_000):
    return (dataf
             .sort(["char", "timestamp"])
             .with_columns([
                 (pl.col("timestamp").diff().cast(pl.Int64) > threshold).fill_null(True).alias("ts_diff"),
                 (pl.col("char").diff() != 0).fill_null(True).alias("char_diff"),
             ])
             .with_columns([
                 (pl.col("ts_diff") | pl.col("char_diff")).alias("new_session_mark")
             ])
             .with_columns([
                 pl.col("new_session_mark").cumsum().alias("session")
             ])
             .drop(['char_diff', 'ts_diff', 'new_session_mark']))

函数增加中间列，以便后续调试；最后删除不需要的列。

添加其他特征

这里是Polars真正出彩的地方。我们不再组合 group_by 和 join 查询，而直接通过表达式API声明，希望在某些分区上计算统计信息。

def add_features(dataf):
    return (dataf
             .with_columns([
                 pl.col("char").count().over("session").alias("session_length"),
                 pl.col("session").n_unique().over("char").alias("n_sessions")
             ]))

计算内容如下：

• .with_columns 中的第一个表达式对角色列计数，实际就是统计行数。表达式链中加入 .over("session")，使计数按会话进行。
• 第二个表达式计算每个角色拥有多少个不同的会话ID。同样，通过在表达式链中加入 .over("char") 实现。

我觉得这非常优雅！一次 .with_columns 调用就能完成很多工作，不必考虑 group_by/join 命令，只需加入所需的计算表达式。

移除机器人

表达式也可以用在其他语句中。需要从数据集移除行时，这尤其方便。

def remove_bots(dataf, max_session_hours=24):
    # We're using some domain knowledge here. The logger of our dataset should
    # log data every 10 minutes. That's what this line is based on.
    n_rows = max_session_hours * 6
    return (dataf
            .filter(pl.col("session_length").max().over("char") < n_rows))

这里再次在表达式链中使用 .over()，计算每个角色的最大 session_length。如果超过允许的最大行数，筛选就会移除属于这个角色的全部行。

锦上添花：灵活的缓存方式

再看一次流水线：

(df
 .pipe(set_types)
 .pipe(sessionize, threshold=20 * 60 * 1000)
 .pipe(add_features)
 .pipe(remove_bots, threshold=24))

Polars的一个优点是，无论DataFrame采用惰性加载还是立即执行模式，这条流水线都可以运行。如果想试验不同的 remove_bots 阈值，可以改写查询，使交互更方便。

df_intermediate = (df
 .pipe(set_types)
 .pipe(sessionize, threshold=20 * 60 * 1000)
 .pipe(add_features))

 df_intermediate.pipe(remove_bots, threshold=24)

这又是一个很好用的API。

结语

本文展示了一种包含会话的查询场景。这类查询通常需要对数据集的不同分区进行聚合。问题不只是查询可能很慢，正确实现它们本身也可能很困难。

def set_types(dataf):
    return (dataf
            .assign(timestamp=lambda d: pd.to_datetime(d['timestamp'], format="%m/%d/%y %H:%M:%S"),
                    guild=lambda d: d['guild'] != -1))

def sessionize(dataf, threshold=60*10):
    return (dataf
             .sort_values(["char", "timestamp"])
             .assign(ts_diff=lambda d: (d['timestamp'] - d['timestamp'].shift()).dt.seconds > threshold,
                     char_diff=lambda d: (d['char'].diff() != 0),
                     new_session_mark=lambda d: d['ts_diff'] | d['char_diff'],
                     session=lambda d: d['new_session_mark'].fillna(0).cumsum())
             .drop(columns=['char_diff', 'ts_diff', 'new_session_mark']))

def add_features(dataf):
    return (dataf
              .assign(session_length=lambda d: d.groupby('session')['char'].transform(lambda d: d.count()))
              .assign(max_sess_len=lambda d: d.groupby('char')['session_length'].transform(lambda d: d.max())))

def remove_bots(dataf, max_session_hours=24):
    n_rows = max_session_hours*6
    return (dataf
            .loc[lambda d: d["max_sess_len"] < n_rows]
            .drop(columns=["max_sess_len"]))

这正是表达式API的价值：让这些常见查询更容易编写。Polars不仅让它们执行更快，也让我们更容易理解查询。

Polars的确很快。在原文的比较中，pandas版本大约运行10分钟，Polars只需7秒，作者称约80倍加速。速度是值得探索Polars的理由，但对我而言，它只是功能的一部分；API同样重要。