高效地扩展Polars GPU Parquet读取器

在处理大型数据集时，数据处理工具的性能至关重要。Polars 作为一个以速度和效率著称的开源数据处理库，它提供了由 cuDF 驱动的 GPU 加速后端，能够显著提升性能。

然而，要想充分发挥 Polars GPU 后端的优势，优化数据加载过程并有效地管理工作流所需的内存至关重要。随着 GPU 后端开发的持续进展，在使用 GPU Parquet 读取器处理不断增大的数据集时，可以使用一些其他技术来保持高性能。现有的 Polars GPU Parquet 读取器（包括至 24.10 版本）无法针对更大的数据集进行扩展。

本文探讨了分块 Parquet 读取器与统一虚拟内存（UVM）相结合，如何能够超越非分块读取器和基于 CPU 的方法。

标度因数和非分块读取器带来的挑战

随着标度因数（SF）的增加，非分块 GPU Polars 读取器（24.10 版）常常会遇到困难。当标度因数超过 200 时，性能会显著下降。在 Query 9 等特定场景下，非分块 GPU 读取器甚至在标度因数达到 50 之前就会出问题。这种限制源于将大型 Parquet 文件加载到 GPU 内存时的内存限制。非分块 Parquet 读取器图表会缺失数据，这凸显了在标度因数较高时遇到的内存溢出（OOM）错误。

图 1. Query 13 执行的可靠性，对比 24.10 版和 24.1 2版 Parquet 读取器

通过分块 Parquet 读取提升 I/O 和峰值内存性能

为了突破这些内存限制，分块 Parquet 读取器就变得至关重要。通过将 Parquet 文件切分为较小数据块进行读取，可以减少内存占用，使 Polars GPU 能够处理更大的数据集。对于任何给定的查询，与非分块读取器相比，使用单次读取限制为 16GB 的分块 Parquet 读取器能够执行更多的标度因数。对于 Query 9，必须使用 16GB 或 32GB 的分块 Parquet 读取方式才能执行并实现更高的吞吐量。

图 2. 针对 Query 9，通过改变标度因数不同分块大小（pass_read_limit）进行吞吐量对比

借助统一虚拟内存读取更大的数据集

虽然分块读取优化了内存管理，但统一虚拟内存将性能提升到了新的高度。统一虚拟内存使 GPU 能够直接访问系统内存，进一步缓解了内存限制并提高了数据传输效率。

相比之下，未配备统一虚拟内存的分块读取器在标度因数达到 100 之前就会触发内存溢出错误。分块读取器与统一虚拟内存相结合，可以在更高的标度因数下成功地执行查询，不过吞吐量会受到影响。

图 3 显示了这一显著优势。与非分块 Parquet 读取器相比，配备统一虚拟内存的分块 Parquet 读取器能够成功地执行标度因数更高的查询。

图 3. 配备统一虚拟内存的分块读取器、CPU 以及未配备统一虚拟内存时的 Query 13 吞吐量（数值越高性能表现越好）

稳定性和吞吐量

在选择最佳的pass_read_limit时，需要重点考虑稳定性和吞吐量之间的平衡。图 1-图 3 表明，16GB 或 32GB 的pass_read_limit实现了两者最优的平衡。

32GBpass_read_limit：除了 Query 9 和 Query 19 因内存溢出异常遭遇失败外，所有其他查询均成功完成。

16GBpass_read_limit：所有查询均成功完成。

分块 GPU 与 CPU 的比较

观察结果显示，每次查询的吞吐量通常高于 CPU Polars，这使得许多在未分块时无法完成的查询得以完成。建议将pass_read_limit参数设置为 16GB 或 32GB，该阈值设定具有合理性。与非分块 Parquet 读取器相比，16GB 或 32GB 的pass_read_limit能够在更高的标度因数下成功地执行查询。

结论

对于 Polars GPU 而言，配备统一虚拟内存的分块 Parquet 读取器通常优于 Polars CPU 和非分块Parquet 读取器，尤其是在处理大规模数据集和高标度因数场景时。通过优化数据加载过程，可以充分发挥 Polars GPU 的潜力，显著提升性能。作为最新的cudf-polars（24.12 版及更高版本）的一部分，分块 Parquet 读取器和统一虚拟内存是读取 Parquet 文件的默认方式。这使得所有查询和标度因数都实现了像上述的性能提升。