GPU 支持 [公开测试版]

Polars 为使用 Lazy API 的 Python 用户提供了一个基于 NVIDIA GPU 的内存中、GPU 加速的执行引擎，使用的是 RAPIDS cuDF。此功能目前处于公开测试阶段，并且正在快速开发中。

系统要求

NVIDIA Volta™ 或更高版本的 GPU，具有计算能力 7.0+
CUDA 11 或 CUDA 12
Linux 或 Windows Subsystem for Linux 2 (WSL2)

请参阅RAPIDS 安装指南以获取完整详情。

安装

您可以通过功能标志为Polars安装GPU后端，作为正常安装的一部分。

Python

pip install polars[gpu]

注意

如果您有CUDA 11，安装行还需要NVIDIA包索引来获取CUDA 11包。

Python

pip install --extra-index-url=https://pypi.nvidia.com polars cudf-polars-cu11

用法

使用惰性API 如常构建查询后，通过运行.collect(engine="gpu")而不是.collect()来请求启用GPU的执行。

Python

import polars as pl

df = pl.LazyFrame({"a": [1.242, 1.535]})

q = df.select(pl.col("a").round(1))

result = q.collect(engine="gpu")
print(result)

shape: (2, 1)
┌─────┐
│ a   │
│ --- │
│ f64 │
╞═════╡
│ 1.2 │
│ 1.5 │
└─────┘

为了更详细地控制执行过程，例如在多GPU节点上指定使用哪个GPU，我们可以提供一个GPUEngine对象。默认情况下，GPU引擎将使用适用于大多数用例的配置。

Python

q = df.select((pl.col("a") ** 4))

result = q.collect(engine=pl.GPUEngine(device=1))
print(result)

shape: (2, 1)
┌──────────┐
│ a        │
│ ---      │
│ f64      │
╞══════════╡
│ 2.379504 │
│ 5.551796 │
└──────────┘

它是如何工作的

当您使用GPU加速引擎时，Polars会创建并优化查询计划，并将其分派到基于RAPIDS cuDF的物理执行引擎，以在NVIDIA GPU上计算结果。最终结果将作为普通的CPU支持的Polars数据框返回。

GPU支持哪些内容？

GPU支持目前处于公开测试阶段，引擎正在快速发展。引擎目前支持许多核心表达式和数据类型，但并非全部。

由于表达式是可组合的，因此列出GPU上支持的完整表达式矩阵是不可行的。相反，我们提供了当前支持和未支持的高级表达式类别和接口列表。

支持

LazyFrame API
SQL API
从CSV、Parquet、ndjson和内存中的CPU数据帧进行I/O操作。
对数值、逻辑、字符串和日期时间类型的操作
字符串处理
聚合和分组聚合
连接
过滤器
缺失数据
连接

不支持

Eager DataFrame API
流式API
对分类、结构和列表数据类型的操作
滚动聚合
时间序列重采样
时区
折叠
用户自定义函数
JSON、Excel 和数据库文件格式

我的查询是否使用了GPU？

GPU引擎的公开测试版发布意味着我们期望一切都能正常运行，但仍有一些我们正在处理的粗糙边缘。特别是Polars表达式API的完整广度尚未得到支持。通过回退到CPU，您的查询应该能够完成，但您可能不会观察到执行时间上的任何变化。有两种方法可以获取更多关于查询是否在GPU上运行的信息。

在详细模式下运行时，任何无法在GPU上执行的查询都会发出一个PerformanceWarning：

Python

df = pl.LazyFrame(
    {
        "key": [1, 1, 1, 2, 3, 3, 2, 2],
        "value": [1, 2, 3, 4, 5, 6, 7, 8],
    }
)

q = df.select(pl.col("value").sum().over("key"))


with pl.Config() as cfg:
    cfg.set_verbose(True)
    result = q.collect(engine="gpu")

print(result)

PerformanceWarning: Query execution with GPU not supported, reason: 
<class 'NotImplementedError'>: Grouped rolling window not implemented
# some details elided

shape: (8, 1)
┌───────┐
│ value │
│ ---   │
│ i64   │
╞═══════╡
│ 6     │
│ 6     │
│ 6     │
│ 19    │
│ 11    │
│ 11    │
│ 19    │
│ 19    │
└───────┘

要禁用回退，并在查询不受支持时让GPU引擎引发异常，我们可以传递一个适当配置的GPUEngine对象：

Python

q.collect(engine=pl.GPUEngine(raise_on_fail=True))

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/coder/third-party/polars/py-polars/polars/lazyframe/frame.py", line 2035, in collect
    return wrap_df(ldf.collect(callback))
polars.exceptions.ComputeError: 'cuda' conversion failed: NotImplementedError: Grouped rolling window not implemented

目前，仅报告了在GPU上执行失败的近因，我们计划扩展此功能以报告查询中所有不支持的操作。

测试

Polars 和 NVIDIA RAPIDS 团队运行全面的单元和集成测试，以确保 GPU 加速的 Polars 后端能够顺利运行。

每次提交到GPU引擎时，都会运行完整的Polars测试套件，以确保结果的一致性。

GPU引擎目前在启用CPU回退的情况下通过了99.2%的Polars单元测试。如果没有CPU回退，GPU引擎通过了88.8%的Polars单元测试。启用回退后，大约有100个失败的测试：其中大约40个由于调试输出不匹配而失败；有些情况下，GPU引擎产生了正确的结果但使用了不同的数据类型；其余的情况是我们未能正确确定查询不受支持，因此在运行时失败，而不是回退。

何时应该使用GPU？

根据我们的基准测试，当您的工作流主要由分组聚合和连接操作组成时，您最有可能在使用GPU引擎时观察到加速。相比之下，I/O密集型查询通常在GPU和CPU上表现出相似的性能。GPU通常比CPU系统拥有更少的内存，因此非常大的数据集可能会因为内存不足错误而失败。根据我们的测试，50-100 GiB的原始数据集（取决于工作流）适合拥有80GiB内存的GPU。

CPU-GPU 互操作性

CPU和GPU引擎都使用Apache Arrow列式内存规范，这使得在CPU和GPU之间快速移动数据成为可能。此外，一个引擎写入的文件可以被另一个引擎读取。

当使用GPU模式时，如果某些功能不被支持，您的工作流不会失败。当您运行 collect(engine="gpu")时，会检查优化后的查询计划，看它是否可以在 GPU上执行。如果不能，它将透明地回退到标准的Polars引擎并在 CPU上运行。

GPU 执行仅在 Lazy API 中可用，因此当查询执行完成时，物化的 DataFrames 将驻留在 CPU 内存中。

提供反馈

请在Polars的问题跟踪器上报告问题和缺失的功能。