在 Apache Arrow 中使用 jemalloc 實現更快速、可擴展的記憶體配置

已發布 2018年7月20日
作者 Uwe Korn (uwe)

隨著 Apache Arrow 0.9 版本的發布，我們已將陣列緩衝區的預設分配器從系統分配器切換為 OSX 和 Linux 上的 jemalloc。這適用於 Arrow 的 C++/GLib/Python 實作。在大多數情況下，更換預設分配器通常是為了避免許多小型、頻繁的（取消）分配所發生的問題。相反地，在 Arrow 中，我們通常處理大型的記憶體內資料集。雖然 jemalloc 為避免記憶體頁面（4kb）以下分配的 RAM 碎片提供了良好的策略，但它也提供了改進跨越多個記憶體頁面分配效能的功能。

在 Apache Arrow 之外，jemalloc 為 Facebook 的基礎架構提供動力（這也是其大部分開發發生的地方）。它也被用作 Rust 中的預設分配器，並且它還有助於 Redis 減少 Linux 上的記憶體碎片（“Allocator”）。

Arrow 中我們需要的一項記憶體分配特殊性是記憶體應為 64 位元組對齊。這是為了讓我們能夠從像 AVX 這樣的 SIMD 指令集中獲得最佳效能。雖然最現代的 SIMD 指令也適用於未對齊的記憶體，但它們在對齊的記憶體上的效能要好得多。為了讓我們的分析應用程式獲得最佳效能，我們希望分配所有記憶體，以便最大限度地提高 SIMD 效能。

對於對齊的分配，POSIX API 僅提供 aligned_alloc(void** ptr, size_t alignment, size_t size) 函數來分配對齊的記憶體。還有 posix_memalign(void **ptr, size_t alignment, size_t size) 來修改分配以達到首選的對齊方式。但它們都不能滿足分配的擴展。雖然 realloc 函數通常可以在不實際移動分配的情況下擴展分配，但它不能確保在移動分配的情況下保持對齊。

在 Arrow 建置時未啟用 jemalloc 的情況下，這會導致在每次新的分配擴展時複製資料。為了減少記憶體複製的次數，我們使用 jemalloc 的 *allocx()-API 來建立、修改和釋放對齊的分配。這給我們帶來顯著加速的典型任務之一是在 Arrow 表格的增量建構上，該表格由多個列組成。我們通常事先不知道表格的大小，並且需要在載入資料時擴展我們的分配。

為了使用 2 倍的記憶體擴展來增量建構向量，我們將使用以下 C 程式碼與標準 POSIX API

size_t size = 128 * 1024;
void* ptr = aligned_alloc(64, size);
for (int i = 0; i < 10; i++) {
  size_t new_size = size * 2;
  void* ptr2 = aligned_alloc(64, new_size);
  memcpy(ptr2, ptr, size);
  free(ptr);
  ptr = ptr2;
  size = new_size;
}
free(ptr);

使用 jemalloc 的特殊 API，我們能夠省略對 memcpy 的顯式呼叫。在無法就地完成記憶體擴展的情況下，仍然由分配器呼叫它，但並非在所有情況下都需要。這將我們的使用者程式碼簡化為

size_t size = 128 * 1024;
void* ptr = mallocx(size, MALLOCX_ALIGN(64));
for (int i = 0; i < 10; i++) {
  size *= 2;
  ptr = rallocx(ptr, size, MALLOCX_ALIGN(64));
}
dallocx(ptr, MALLOCX_ALIGN(64));

為了了解使用 jemalloc 的實際好處，我們查看 Arrow C++ 中的基準測試。在那裡，我們模擬了一個典型的使用案例，即增量建立原始值陣列。對於陣列的建立，我們不知道最終陣列中元素的數量，因此我們需要不斷擴展儲存資料的記憶體區域。此基準測試的程式碼是 Arrow C++ 原始碼中 builder-benchmark 的一部分，名為 BuildPrimitiveArrayNoNulls。

沒有 jemalloc 的執行時間

BM_BuildPrimitiveArrayNoNulls/repeats:3                 636726 us   804.114MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 621345 us   824.019MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 625008 us    819.19MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_mean            627693 us   815.774MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_median          625008 us    819.19MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_stddev            8034 us   10.3829MB/s

有 jemalloc 的執行時間

BM_BuildPrimitiveArrayNoNulls/repeats:3                 630881 us   811.563MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 352891 us   1.41687GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 351039 us   1.42434GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_mean            444937 us   1.21125GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_median          352891 us   1.41687GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_stddev          161035 us   371.335MB/s

基準測試針對每種配置運行了三次，以查看效能差異。每種配置的第一次運行產生了相同的效能，但在所有後續運行中，使用 jemalloc 的版本速度大約快兩倍。在這些情況下，用於建構陣列的記憶體區域可以在不移動資料的情況下就地擴展。這是可能的，因為有分配給進程的記憶體頁面未使用，但未被作業系統回收。如果沒有 jemalloc，我們就無法利用它們，原因很簡單，預設分配器沒有提供對齊重新分配的 API。