高效地?cái)U(kuò)展Polars GPU Parquet讀取器

在處理大型數(shù)據(jù)集時(shí)，數(shù)據(jù)處理工具的性能至關(guān)重要。Polars 作為一個(gè)以速度和效率著稱的開源數(shù)據(jù)處理庫，它提供了由 cuDF 驅(qū)動(dòng)的 GPU 加速后端，能夠顯著提升性能。

然而，要想充分發(fā)揮 Polars GPU 后端的優(yōu)勢，優(yōu)化數(shù)據(jù)加載過程并有效地管理工作流所需的內(nèi)存至關(guān)重要。隨著 GPU 后端開發(fā)的持續(xù)進(jìn)展，在使用 GPU Parquet 讀取器處理不斷增大的數(shù)據(jù)集時(shí)，可以使用一些其他技術(shù)來保持高性能?，F(xiàn)有的 Polars GPU Parquet 讀取器（包括至 24.10 版本）無法針對更大的數(shù)據(jù)集進(jìn)行擴(kuò)展。

本文探討了分塊 Parquet 讀取器與統(tǒng)一虛擬內(nèi)存（UVM）相結(jié)合，如何能夠超越非分塊讀取器和基于 CPU 的方法。

標(biāo)度因數(shù)和非分塊讀取器帶來的挑戰(zhàn)

隨著標(biāo)度因數(shù)（SF）的增加，非分塊 GPU Polars 讀取器（24.10 版）常常會遇到困難。當(dāng)標(biāo)度因數(shù)超過 200 時(shí)，性能會顯著下降。在 Query 9 等特定場景下，非分塊 GPU 讀取器甚至在標(biāo)度因數(shù)達(dá)到 50 之前就會出問題。這種限制源于將大型 Parquet 文件加載到 GPU 內(nèi)存時(shí)的內(nèi)存限制。非分塊 Parquet 讀取器圖表會缺失數(shù)據(jù)，這凸顯了在標(biāo)度因數(shù)較高時(shí)遇到的內(nèi)存溢出（OOM）錯(cuò)誤。

圖 1. Query 13 執(zhí)行的可靠性，對比 24.10 版和 24.1 2版 Parquet 讀取器

通過分塊 Parquet 讀取提升 I/O 和峰值內(nèi)存性能

為了突破這些內(nèi)存限制，分塊 Parquet 讀取器就變得至關(guān)重要。通過將 Parquet 文件切分為較小數(shù)據(jù)塊進(jìn)行讀取，可以減少內(nèi)存占用，使 Polars GPU 能夠處理更大的數(shù)據(jù)集。對于任何給定的查詢，與非分塊讀取器相比，使用單次讀取限制為 16GB 的分塊 Parquet 讀取器能夠執(zhí)行更多的標(biāo)度因數(shù)。對于 Query 9，必須使用 16GB 或 32GB 的分塊 Parquet 讀取方式才能執(zhí)行并實(shí)現(xiàn)更高的吞吐量。

圖 2. 針對 Query 9，通過改變標(biāo)度因數(shù)不同分塊大?。╬ass_read_limit）進(jìn)行吞吐量對比

借助統(tǒng)一虛擬內(nèi)存讀取更大的數(shù)據(jù)集

雖然分塊讀取優(yōu)化了內(nèi)存管理，但統(tǒng)一虛擬內(nèi)存將性能提升到了新的高度。統(tǒng)一虛擬內(nèi)存使 GPU 能夠直接訪問系統(tǒng)內(nèi)存，進(jìn)一步緩解了內(nèi)存限制并提高了數(shù)據(jù)傳輸效率。

相比之下，未配備統(tǒng)一虛擬內(nèi)存的分塊讀取器在標(biāo)度因數(shù)達(dá)到 100 之前就會觸發(fā)內(nèi)存溢出錯(cuò)誤。分塊讀取器與統(tǒng)一虛擬內(nèi)存相結(jié)合，可以在更高的標(biāo)度因數(shù)下成功地執(zhí)行查詢，不過吞吐量會受到影響。

圖 3 顯示了這一顯著優(yōu)勢。與非分塊 Parquet 讀取器相比，配備統(tǒng)一虛擬內(nèi)存的分塊 Parquet 讀取器能夠成功地執(zhí)行標(biāo)度因數(shù)更高的查詢。

圖 3. 配備統(tǒng)一虛擬內(nèi)存的分塊讀取器、CPU 以及未配備統(tǒng)一虛擬內(nèi)存時(shí)的 Query 13 吞吐量（數(shù)值越高性能表現(xiàn)越好）

穩(wěn)定性和吞吐量

在選擇最佳的pass_read_limit時(shí)，需要重點(diǎn)考慮穩(wěn)定性和吞吐量之間的平衡。圖 1-圖 3 表明，16GB 或 32GB 的pass_read_limit實(shí)現(xiàn)了兩者最優(yōu)的平衡。

32GBpass_read_limit：除了 Query 9 和 Query 19 因內(nèi)存溢出異常遭遇失敗外，所有其他查詢均成功完成。

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

分塊 GPU 與 CPU 的比較

觀察結(jié)果顯示，每次查詢的吞吐量通常高于 CPU Polars，這使得許多在未分塊時(shí)無法完成的查詢得以完成。建議將pass_read_limit參數(shù)設(shè)置為 16GB 或 32GB，該閾值設(shè)定具有合理性。與非分塊 Parquet 讀取器相比，16GB 或 32GB 的pass_read_limit能夠在更高的標(biāo)度因數(shù)下成功地執(zhí)行查詢。

結(jié)論

對于 Polars GPU 而言，配備統(tǒng)一虛擬內(nèi)存的分塊 Parquet 讀取器通常優(yōu)于 Polars CPU 和非分塊Parquet 讀取器，尤其是在處理大規(guī)模數(shù)據(jù)集和高標(biāo)度因數(shù)場景時(shí)。通過優(yōu)化數(shù)據(jù)加載過程，可以充分發(fā)揮 Polars GPU 的潛力，顯著提升性能。作為最新的cudf-polars（24.12 版及更高版本）的一部分，分塊 Parquet 讀取器和統(tǒng)一虛擬內(nèi)存是讀取 Parquet 文件的默認(rèn)方式。這使得所有查詢和標(biāo)度因數(shù)都實(shí)現(xiàn)了像上述的性能提升。