从67.25s到4.36s——从pandas到polars的重写之路

1天下苦pandas久矣，pandas不仅不支持多线程和惰性执行、还存在内存放大效应和索引复杂等问题。

当我的o2o实践项目进行至一半时，特征工程部分的时间消耗已经到了不可接受的地步，其花费的时间与模型相当。我萌生了重写的想法，但在当时我对pandas和polars都不熟悉，而LLM显然对pandas更熟悉，为了稳妥起见，我还是继续使用pandas，这也导致最终项目积重难返，特征工程部分达到2000+行Python和67+s的恐怖屎山规模，在ddl前我仍在尝试特征工程提高测试集AUC表现，也错失了最后重构的机会。

正式进入寒假，我终于有时间对项目进行一些大刀阔斧的重构。经过一番简单的调研，我选择使用polars重写，彻底放弃pandas。

Figure 1: TPC-H benchmark

我选择polars的原因主要有三点

  • 快。polars的速度极快，在各类易用的PyPI库中名列前茅，如Figure 1 所示。默认提供众多优化，如多线程、惰性执行、内存管理等。提高速度的同时毫无心智负担。
  • 跨语言。pandas只提供一套Python的API，而polars则是支持Node.js、Python、Rust、R等多种语言，且API类似，跨语言迁移没有困难。
  • 易用的API。polars的API严谨稳定。pandas的特色是高度灵活的API，但带来的是低效率和高度的抽象。而polars则提供大量“预制”接口，非必要不使用Python函数，保证极高的性能。

Polars重写的Commit

不过令人啼笑皆非的是，使用polars重写后，虽然特征工程产生的数据整体没有变化，但相比于pandas的结果发生了shuffle，这继而导致了模型的执行结果发生了变动，非常可惜的是测试集上的AUC下降，经过一番暴力调参仍然没有达到之前的高度。

  1. 1存在标题党的嫌疑