Polars GPU 加速引擎实战指南：在 NVIDIA GPU 上运行 Lazy API 查询

• 数据分析
• 大数据

导读

本文围绕 Polars 官方用户指南中的 gpu-support.md 展开，系统讲解基于 RAPIDS cuDF 的 GPU 加速执行引擎：从系统要求与安装方式（pip install polars[gpu]）、基础用法（.collect(engine="gpu")）到多 GPU 引擎（RayEngine/DaskEngine/SPMDEngine）、CPU-GPU 互操作与透明回退机制，以及如何判断查询是否真正运行在 GPU 上。读完本文，你将掌握在 Linux/WSL2 + NVIDIA GPU 环境下把 Polars Lazy 查询无缝切换到 GPU 执行、排查回退原因、控制失败策略的完整方案。

一、GPU 支持概览：Open Beta 阶段的能力边界

Polars 为 Python 用户的 Lazy API 提供了一套 GPU 加速执行引擎，底层基于 NVIDIA RAPIDS cuDF。该功能目前处于Open Beta（开放测试）阶段，正在快速迭代，文档中明确指出"currently a single GPU implementation"（lazy/gpu.md），即当前默认实现面向单 GPU 场景，多 GPU 能力需要通过额外的引擎类获得。

核心结论先行：

• GPU 执行仅适用于 Lazy API，Eager DataFrame API 不在支持范围内；
• 触发方式是在.collect或.sink_*调用中传入engine="gpu"；
• 查询若包含不支持的算子，会透明回退到 CPU 引擎，不会报错（除非显式设置raise_on_fail=True）；
• 最终结果始终以普通 CPU 内存中的 Polars DataFrame 返回。

从源码结构看，GPU 引擎是 Polars 统一"执行引擎"抽象体系中的一员。在 engine_config.py 中定义了四个受支持的引擎名："auto"、"in-memory"、"streaming"和"gpu"，而EngineTypeName类型别名与之完全对应（_typing.py）。"gpu"在_engine_from_name中被解析为GPUEngine()实例（engine_config.py），并支持通过Config.set_engine_affinity设为全局默认引擎。

二、系统要求：先检查硬件与驱动

在安装 GPU 后端之前，请确认环境满足以下条件（见 gpu-support.md）：

需要注意：GPU 加速引擎面向的是 Python 用户，Rust 侧的 Polars 本体不包含该后端。详细的系统级要求以 RAPIDS 官方安装指南为准（本项目文档引用了 RAPIDS installation guide 供读者核对最新细节，本文不重复搬运外部内容）。

三、安装：feature flag 与 CUDA 版本匹配

GPU 后端通过 Polars 的 Python 包特性标签（extra）安装，属于常规 安装指南 的一部分：

pip install polars[gpu]

在仓库的 pyproject.toml 中可以看到该 extra 的实际定义：

# GPU Engine gpu = ["cudf-polars-cu12"]

也就是说，polars[gpu]当前被配置为安装cudf-polars-cu12（对应 CUDA 12）。如果系统使用不同版本的 CUDA，需要单独安装与 CUDA 版本后缀匹配的 cudf-polars 库，例如 CUDA 13：

pip install polars cudf-polars-cu13

版本约束：cudf-polars 只支持有限的 Polars 版本范围

文档特别提醒一个安装陷阱：cudf-polars仅支持一定范围内的 Polars 版本。如果 Polars 版本未被固定，包解析器可能会回退选择一个较旧的兼容 Polars 版本。当这种隐式降级不可接受时，应显式固定所需的 Polars 版本，否则不兼容的组合会在依赖解析阶段直接失败。建议做法：

pip install polars==<your-version> cudf-polars-cu12

此外，installation.md 的 GPU 特性标签表把gpu描述为"Run queries on NVIDIA GPUs"，并链接到本文作为详细前提说明。

四、基础用法：从 Lazy 查询到 GPU 执行

GPU 加速引擎的入口非常轻量：照常使用 Lazy API 构建查询，然后在.collect时传入engine="gpu"即可。以下示例来自仓库中的文档代码片段 gpu.py：

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 执行同样适用于sink_*系列调用（如sink_parquet、sink_csv、sink_ipc、sink_ndjson）。从 frame.py 的sink_parquet实现可以看到，引擎对象统一通过_select_engine(engine)解析后分发到对应的 sink 逻辑，GPU 引擎在 engine.py 中同样实现了这些 sink 方法。

engine参数的类型约束

从 _typing.py 看，engine参数的类型为：

EngineTypeName: TypeAlias = Literal["auto", "in-memory", "streaming", "gpu"] EngineType: TypeAlias = Union[EngineTypeName, "Engine"]

即既可以是字符串"gpu"，也可以直接传一个Engine实例（例如GPUEngine(...)对象），后者用于更细粒度的控制。frame.py 的collect文档进一步说明：若所选引擎无法运行查询，Polars 会回退到 in-memory 引擎；GPU 模式被视为不稳定（unstable），并建议用POLARS_VERBOSE=1观察回退原因。

五、GPUEngine对象与多 GPU 执行

engine="gpu"字符串适合单 GPU 场景。需要多 GPU 执行或更多运行时配置时，可以传入GPUEngine对象。Polars 侧的GPUEngine定义在 engine.py，通过 lazyframe/__init__.py 与顶层pl.GPUEngine导出。

GPUEngine的可配置参数（均有关键字默认值）：

注意GPUEngine._post_opt_callback的实现细节（engine.py）：

• 若请求background收集，会发出UserWarning提示"GPU engine does not support background collection"并禁用 GPU 引擎；
• 若处于 eager 模式，则不运行 GPU（且不告警）；
• 其余情况通过import_optional("cudf_polars", ...)惰性导入 cudf-polars，缺失时会给出包含安装指引的明确报错；
• 最终返回cudf_polars.execute_with_cudf(config=self)作为后优化回调，把经过 Polars 优化的查询计划交给 cuDF 执行。

三种多 GPU 引擎子类

自 cudf-polars 26.06 版本起，cudf-polars 库提供 3 个GPUEngine子类（详见原文档引用的 RAPIDS 文档，本文仅转述其用途）：

• RayEngine：借助 Ray 实现多 GPU 执行；
• DaskEngine：借助 Dask 实现多 GPU 执行；
• SPMDEngine：采用"单程序多数据"（Single Program, Multiple Data）模型实现多 GPU 执行。

这 3 个引擎会启动分布式资源，因此建议作为上下文管理器（context manager）使用，退出时自动回收资源。示例（来自文档片段 gpu.py）：

import polars as pl from cudf_polars.engine.ray import RayEngine df = pl.LazyFrame({"a": [1.242, 1.535]}) q = df.select((pl.col("a") ** 4)) with RayEngine() as engine: result = q.collect(engine=engine) print(result)

输出：

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

注意：使用RayEngine或DaskEngine前，必须分别安装 Ray 或 Dask。可通过 cudf-polars 的[ray]、[dask]extra 安装，例如 CUDA 13 环境下：

pip install cudf-polars-cu13[ray] pip install cudf-polars-cu13[dask]

六、工作原理：查询计划如何交给 cuDF

当使用 GPU 加速引擎时，执行流程为（对应 gpu-support.md 的 "How It Works"）：

1. Polars 正常创建并优化查询计划（沿用 Lazy API 的完整优化管线，包括谓词下推、投影下推等）；
2. 优化后的计划被派发到基于 RAPIDS cuDF 的物理执行引擎；
3. cuDF 在 NVIDIA GPU 上完成计算；
4. 最终结果以普通 CPU 内存中的 Polars DataFrame返回。

这一设计在源码中体现为GPUEngine._post_opt_callback把cudf_polars.execute_with_cudf挂接到查询执行的优化后回调阶段，Polars 只负责计划优化，执行权交给 cuDF 后端。

CPU-GPU 互操作性：基于 Arrow 列式内存

CPU 与 GPU 引擎都遵循Apache Arrow 列式内存规范，这是两者之间数据可以快速迁移的根本原因。由此带来两个直接好处：

• CPU 与 GPU 之间的数据搬运开销小、速度快；
• 一个引擎写出的文件（如 Parquet、IPC）可以被另一个引擎直接读取。

同时，GPU 执行只在 Lazy API 中可用，查询执行完毕时物化出的 DataFrame 会驻留在 CPU 内存中，因此不存在"结果留在 GPU 上、后续操作被隐性绑死"的问题，GPU 加速被严格限定在查询计算阶段。

透明回退：查询不会因算子不支持而失败

当传入engine="gpu"时，Polars 会检查优化后的查询计划能否在 GPU 上执行：

• 若能 → 交给 cuDF 在 GPU 上运行；
• 若不能 →透明回退到标准 Polars 引擎，所有查询操作在 CPU 上完成，且不中断执行。

从 frame.py 的文档可以看到官方对这一行为的确认："Not all queries will run successfully on the GPU, however, they should fall back transparently to the default engine if execution is not supported"，并建议开启POLARS_VERBOSE=1观察回退情况。

七、GPU 上支持与不支持的功能清单

GPU 支持目前处于 Open Beta，引擎支持大部分而非全部核心表达式与数据类型。由于表达式是可组合的，官方无法列出完整的表达式矩阵，而是给出了高层次的功能类别清单：

当前支持

• LazyFrame API
• SQL API
• 从 CSV、Parquet、ndjson 以及内存中 CPU DataFrame 的 I/O
• 数值、逻辑、字符串和 datetime 类型的操作
• 字符串处理
• 聚合（含分组聚合与滚动聚合变体）
• 连接（Joins）
• 过滤（Filters）
• 缺失数据处理
• 拼接（Concatenation）

当前不支持

• Eager DataFrame API
• Date、Categorical、Enum、Time、Array、Binary 和 Object 数据类型
• 部分带时区的 Datetime 表达式以及 List 类型表达式
• 时间序列重采样（resampling）
• Folds（折叠）
• 用户自定义函数（UDF）
• Excel 和数据库文件格式

这意味着：如果你的查询涉及 Binary 类型转换、自定义函数或 Excel 输入，GPU 引擎会回退到 CPU。仓库中的回退演示用例正是利用了这一点（gpu.py）：

df = pl.LazyFrame({"value": [1, 2, 3, 4, 5, 6, 7, 8]}) q = df.select(pl.col("value").cast(pl.Binary()).head(1))

八、如何确认查询真的跑在 GPU 上？

由于回退是透明的，一个包含不支持算子的查询会静默回退到 CPU，执行时间看不出任何变化。Open Beta 引擎意味着大部分常见表达式 API 都已支持，但你仍需要以下两种手段来确认查询是否真正使用了 GPU：

方式一：开启 verbose 模式观察PerformanceWarning

在 verbose 模式下，任何无法在 GPU 上执行的查询都会发出PerformanceWarning：

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

对于上面的 Binary 转换示例，输出大致为：

PerformanceWarning: Query execution with GPU not possible: unsupported operations The errors were: - NotImplementedError: dtype=Binary conversion not supported return wrap_df(ldf.collect(engine, callback)) shape: (1, 1) ┌───────┐ │ value │ │ --- │ │ bin │ ╞═══════╡ │ b"1" │ └───────┘

等价地，也可以通过环境变量POLARS_VERBOSE=1开启 verbose 输出（frame.py 中即以此形式给出提示）。

方式二：raise_on_fail=True让回退变成异常

如果希望查询不支持 GPU 时直接失败而不是静默回退，可以传入GPUEngine(raise_on_fail=True)：

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

此时控制台会抛出ComputeError：

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: ('Query execution with GPU not possible: unsupported operations.\nThe errors were:\n- NotImplementedError: dtype=Binary conversion not supported', [NotImplementedError('dtype=Binary conversion not supported')])

当前版本只报告导致失败的最直接原因（proximal cause），官方计划扩展为报告查询中所有不支持的算子。raise_on_fail在GPUEngine.__init__中被写入config并透传给 cudf-polars（见 engine.py）。

九、将 GPU 设为全局默认引擎

除了每次调用时显式传参，还可以通过Config.set_engine_affinity把"gpu"设为默认执行引擎，之后所有.collect()都会默认尝试 GPU（config.py）：

pl.Config.set_engine_affinity("gpu") lf.max().collect()

也可以直接传入GPUEngine对象以同时固化细粒度配置：

pl.Config.set_engine_affinity(pl.GPUEngine(device=1, raise_on_fail=True))

注意语义：设置默认引擎后，查询并不保证一定以该引擎执行（不支持时仍会回退）。配置的持久化通过环境变量POLARS_ENGINE_AFFINITY完成；引擎对象是进程本地的，Config.save不会保存对象形态的 affinity，加载含POLARS_ENGINE_AFFINITY的状态会替换对象 affinity。传入非法引擎名会抛ValueError。

十、测试与质量：Polars 与 NVIDIA RAPIDS 的联合保障

GPU 引擎的质量由 Polars 与 NVIDIA RAPIDS 团队联合维护：

• 完整测试套件：GPU 引擎的每次提交都会运行完整的 Polars 测试套件，确保结果一致性；
• 单元测试通过率：启用 CPU 回退时，GPU 引擎通过 Polars 单元测试的99.2%；禁用回退时通过88.8%；
• 回退下的失败分布：启用回退时约有 100 个失败测试，其中约 40 个因调试输出不一致（debug output mismatch）失败——部分场景 GPU 引擎产生正确结果但使用了不同的数据类型；其余则是未能正确判定查询不支持、导致运行时失败而非回退的用例。

这些数字说明：GPU 引擎对常见算子覆盖良好，但边界情况（类型推断差异、支持判定漏洞）仍在持续打磨中。

十一、什么时候该用 GPU？

根据官方基准测试结论（注意：这是项目文档给出的经验结论，并非本仓库源码可验证的数据）：

• 当工作负载以**分组聚合（grouped aggregations）和连接（joins）**为主时，最有可能观察到 GPU 带来的加速；
• 相反，I/O 密集型查询在 GPU 与 CPU 上的性能通常相近（瓶颈在数据读取而非计算）；
• 根据团队测试，1 TiB 的原始数据集（视工作负载而定）与80 GiB 显存的 GPU 匹配良好。

因此，决策要点是：先分析查询的计算特征，把 GPU 资源留给计算密集、聚合/连接占比高的任务；对 I/O 瓶颈明显的查询，CPU 引擎往往已经足够。

十二、反馈与演进

GPU 支持处于 Open Beta，功能正在快速迭代。遇到缺失功能或 bug，建议在 Polars 官方 issue 跟踪器中提交问题（原文档提供了对应入口）。从仓库结构可以观察到，GPU 引擎的 Python 侧接入点集中在 py-polars/src/polars/lazyframe/（engine.py 的GPUEngine、engine_config.py 的引擎解析与 affinity），而实际计算逻辑全部在 cudf-polars 库内，Polars 仓库本身不包含 cuDF 的实现代码——这也是理解"GPU 支持是 Polars 与 RAPIDS 生态协作产物"的关键视角。

相关参考

• GPU 支持官方指南（本文主体来源）
• Lazy API 入门 与 Lazy 章节的 GPU 简介
• 安装指南与 feature flag 说明
• GPUEngine 源码实现
• 引擎解析与 affinity 配置
• collect 的 engine 参数文档
• 文档代码示例

• 数据分析
• 大数据