博客

技术分享

高效扩展 Polars 的 GPU Parquet 读取器

2025.04.11 29分钟阅读

高效扩展 Polars 的 GPU Parquet 读取器

在处理大型数据集时，数据处理工具的性能变得至关重要。Polars 是一个以速度和效率闻名的开源数据操作库，提供由 cuDF 驱动的 GPU 加速后端，可以显著提高性能。

“但是，为了充分利用 Polars GPU 后端的强大功能，必须优化数据加载过程并有效管理工作流程所需的内存。随着 GPU 后端开发的不断推进，在使用 GPU Parquet 阅读器时，随着数据集大小的增加，我们还可以使用一些其他技术来保持高性能。现有的 Polars GPU Parquet 读取器 (到版本 24.10) 无法针对更高的数据集大小进行扩展。”

本文将探讨分块 Parquet Reader 与 Unified Virtual Memory (UVM) 相结合后，如何在性能上优于非分块阅读器和基于 CPU 的方法。

规模因素和非分块读取器带来的挑战

随着规模系数 (SF) 的增加，非分块 GPU Polars Reader (24.10) 往往难以实现。超过 SF200 后，性能会显著下降。在某些情况下，例如 Query 9，非分块 GPU 读取器会出现故障，甚至在达到 SF50 之前。这种限制是由于在 GPU 的内存中加载大型 Parquet 文件时存在内存限制而产生的。非分块 Parquet Reader 图中的缺失数据突出显示了在更高比例因子下遇到的 out-of-memory (OOM) 错误。

图 1.Query 13 执行可靠性，24.10 至 24.12 Parquet Reader 对比

通过分块 Parquet 读取改善 IO 和峰值内存

为了克服这些内存限制，分块的 Parquet Reader 变得至关重要。通过以较小的数据块读取 Parquet 文件，显存占用减少，使 Polars GPU 能够处理更大的数据集。与非分块阅读器相比，使用具有 16 GB 通过读取限制的分块 Parquet Reader 可为任何给定查询执行更多的比例系数。对于 Query 9，必须使用 16 GB 或 32 GB 的分块 Parquet 读取数据，才能执行并获得更好的吞吐量。