GPU 支持 [公开测试版]

Polars 为使用 Lazy API 的 Python 用户在 NVIDIA GPU 上提供了基于内存的 GPU 加速执行引擎，该引擎使用 RAPIDS cuDF。此功能目前处于开放测试版，并正在快速开发中。

系统要求

NVIDIA Volta™ 或更高版本的 GPU，具有计算能力 7.0+
CUDA 11 或 CUDA 12
Linux 或适用于 Linux 的 Windows 子系统 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

用法

在使用 Lazy 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 DataFrames 进行 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 系统具有更少的 RAM，因此非常大的数据集可能会因内存不足错误而失败。根据我们的测试，50-100 GiB 的原始数据集（取决于工作流程）在具有 80GiB 内存的 GPU 上运行良好。

CPU-GPU 互操作性

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

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

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

提供反馈

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