将NVIDIA加速计算引入Polars

Polars 近日发布了一款由 RAPIDS cuDF 驱动的全新 GPU 引擎，该引擎可将 NVIDIA GPU 上的 Polars 工作流速度最高提速 13 倍，使数据科学家仅在一台机器上就能实现在数秒内处理数亿行数据。

日新月异的数据挑战

Pandas 等传统数据处理库均为单线程，当处理数据超过数百万行时就会显得“力不从心”。分布式数据处理系统虽然可以处理数十亿行数据，但又会增加处理中小型数据集的复杂程度和经费开支。

在高效处理数千万至数亿行数据的工具之间一直存在着差距。这类工作负载常见于金融、零售、制造等行业的模型开发、需求预测和物流中。

在面向数据科学家和工程师的 Python 库中，Polars 的增长速度位居前列，其设计初衷就是为了应对这些挑战。该程序库使用了先进的查询优化技术，以减少不必要的数据移动和处理，使数据科学家可以仅用一台机器就能流畅地处理数亿行规模的工作负载。Polars 弥补了单线程解决方案速度过慢和分布式系统会增加非必要复杂性的缺陷，提供了一个极具吸引力的“中等规模”数据处理解决方案。

将 NVIDIA 加速计算引入 Polars

与其他仅使用 CPU 的数据处理工具相比，Polars 利用多线程执行、高级内存优化和惰性求值（lazy evaluation）功能，可显著提升开箱即用的加速性能。

但由于各行各业的企业面临日益增长的数据处理需求，例如需要分析数十亿笔金融交易、管理复杂的库存系统等，都需要更加强大的性能。这时就该轮到加速计算发挥作用了：

cuDF 是 NVIDIA CUDA-X 加速库中的 RAPIDS 套件的一部分，它是一个由 GPU 提供加速的 DataFrame 程序库，能够利用 GPU 的大规模并行处理能力显著提高数据处理性能。

Polars 团队与 NVIDIA 一起将 cuDF 的速度与 Polars 的效率相结合，使性能最高提速至在 CPU 上 Polars 的 13 倍。如此，即便用户的数据处理工作量增长到数亿甚至数十亿行数据，依然能够保持交互。

图 1.图中所示的是 PDS-H 基准测试 22 个查询中加速幅度最大的 4 个查询。在运行包含众多复杂分组和连接操作的查询时，RAPIDS cuDF 驱动的 Polars GPU 引擎的速度提速至 CPU 上 Polars 的 13 倍。

PDS-H基准测试规模系数 80 | GPU：NVIDIA H100 | CPU：英特尔 Xeon W9-3495X（Sapphire Rapids） | 存储：本地NVMe。备注：PDS-H 源自 TPC-H，但这些结果与 TPC-H 的结果不具有可比性。

由于 Polars GPU 引擎直接内置在 Polars Lazy API 中，用户只需通过 pip 安装 polars[gpu] 并将 [engine=”gpu”] 发送至 collect 操作，即可将 GPU 加速应用于他们的工作流。在后台，Polars 会尝试首先在 GPU 上执行操作，必要时再返回 CPU。这种方法可确保：

通过使用 Polars 的查询优化器，实现高效执行和最低内存使用

用户无须修改现有 Polars 代码，即可访问 GPU 引擎

完全兼容 Polars 不断发展的数据可视化、I/O 和机器学习库生态系统

pip install polars[gpu] --extra-index-url=https://pypi.nvidia.com
 
import polars as pl
 
(transactions
 .group_by("CUST_ID")
 .agg(pl.col("AMOUNT").sum())
 .sort(by="AMOUNT", descending=True)
 .head()
 .collect(engine="gpu"))

结语

RAPIDS cuDF 驱动的 Polars GPU 引擎现已发布公测版，为各行各业的数据科学家和工程师提供了一种适用于中等规模数据处理的强大工具。该引擎最高能够将 NVIDIA GPU 上的 Polars 工作流速度提速 13 倍，可以在不产生分布式系统开销的情况下，高效处理数亿行规模的数据集。Polars GPU 引擎直接内置在 Polars API 中，使所有用户都能轻松访问。