資料管線的時間主要花費在三個地方:讀取和解碼資料、對每一列進行的工作,以及等待(等待儲存、等待鎖、或等待忙碌的執行緒)。解碼的部分通常已經在函式庫中調整過了。而每一列上的工作是你寫的,這正是粗心的迴圈會導致 2–3 倍效能損失的地方。
本檔案大部分內容在討論處理掃描 (scan) 產生的批次 (batch) 的程式碼。關於底層的讀取器、寫入器或編解碼器,以及如何取得可信賴的數值,請參閱最後兩節及其參考資料。
首先,測量時間花在哪裡
你很少會有第二個實作(如 Ray、DuckDB、pyarrow)可以作為比較基準。不如直接使用掃描本身作為衡量標尺:
- 單獨計時掃描。 執行相同的掃描,使用相同的欄位、過濾條件和
ScanOptions,然後丟棄這些批次。這是一個純解碼的基準線:沒有任何管線能比它更快。 - 單獨計時你的工作。 將每次批次處理的呼叫(
fold、add,或任何消耗該批次的函式)用perf_counter_ns()包起來,並將時間加總。 - 比較它們。 在單一執行緒上,總時間接近解碼時間的 1.5 倍,表示工作很輕量;超過這個比例,則表示工作是瓶頸,值得優化。在多個執行緒上,還要檢查工作是否只在一個執行緒上運行,而解碼卻在多個執行緒上進行。
量測實例:一個對 7700 萬筆 trips 的 group-by 操作在單一執行緒上耗時 4.6 秒,其中 2.1 秒是逐列工作(佔 45%)。在 8 個執行緒上則耗時 2.6 秒,其中同樣的 2.1 秒工作佔了 82%,原因是只有一個執行緒消耗了所有批次。 references/profiling.md 提供了配方及其陷阱。
逐列迴圈的規則
每條規則都附帶在 dista shuffle 基準測試(7700 萬列,本地表,單一執行緒)中節省的時間。程式碼在 references/hot-loops.md 中。
| 規則 | 量測效果 |
|---|---|
對每個批次只獲取一次每個欄位的原始指標,然後對該指標進行索引。不要在迴圈內對 List 或 Tuple[List, List] 進行索引。 | 整個逐列工作從 1.40 秒降至 0.78 秒 |
| 將一個 key 累積的所有資料並排放置在同一個槽位(結構陣列),而不是平行陣列,這樣一次更新只會觸及一個快取行。 | 與下一條規則一起,從 2.09 秒降至 1.40 秒 |
進行無分支過濾:計算每一列,將結果寫入 kept 索引位置,然後將 kept 增加 Int(keep)。然後在第二個緊湊的迴圈中處理保留的列。 | 與上一條規則一起量測 |
當 key 是小整數(例如 zone × zone × hour 總計 170 萬)時,使用 key 索引一個緊湊陣列,而不是使用 Dict。 | 未單獨量測;兩種版本都使用了緊湊陣列 |
| 在優化過程中,不要改變浮點數加法的順序。這樣結果必須與舊結果逐位元相同,這是一個免費的正確性檢查。 | 上述每一步的結果都完全匹配 |
整個管線的規則
| 規則 | 量測效果 |
|---|---|
| 不要讓一個執行緒消費多個執行緒解碼的結果。讓每個 worker 擁有自己的累加器,最後再合併,就像 dista 的 map 任務所做的那樣。 | 在 AWS 上,16 個執行緒僅將單程序 group-by 從 202 秒降至 126 秒 |
即使只有一個解碼執行緒,也要保持多個讀取操作同時進行:使用 iceberg.mojo 中的 ScanOptions.files_at_once(和 overlap_io)。 | 在從 S3 讀取的單一執行緒上達到 6.6–6.9 倍加速 |
| 每個程序只載入一次共享 C 函式庫並保留其 handle,絕對不要在每次呼叫時載入。 | 當 zstd 在每次呼叫時被打開,掃描有 52% 的時間花在 dyld 上;修正後,從 2,687 毫秒降至 262 毫秒 |
| 快取任何重複的位元組解析結果,例如 manifest 中的 Avro 結構。 | manifest 規劃的固定成本降低了 3.6 倍(從 113.7 µs 降至 31.9 µs) |
| 不要為每筆紀錄分配記憶體:直接解碼到重複使用的 slot 中。 | Avro 讀取器在調整前,每筆紀錄大約分配了 60 次記憶體 |
| 了解天花板:超過一定核心數後,掃描會受到記憶體頻寬限制,更多執行緒也無濟於事。應該改為讀取更少的位元組。 | 一次 Iceberg 掃描在單一核心上耗時 35.8 毫秒,在四個核心上耗時 12.1 毫秒,約為 3 倍 |
references/case-studies.md 詳細說明了每個案例的故事,以及是什麼原因暴露了問題。
讀取器、寫入器與編解碼器
當解碼本身成為瓶頸時,規則主要是關於不做會被浪費掉的工作。詳細內容和數據在 references/readers-and-codecs.md 中。
- 檢查每個快速路徑的閘門條件是否能被到達它的呼叫者滿足,並透過檢測證明它確實被執行。一個向量化路徑的閘門條件沒有任何 list 欄位能滿足,導致一個 9 KiB 的欄位佔了 43% 的讀取時間。
- 不要計算後續決策通常會丟棄的結果。 建立字典只為了之後選擇純文字編碼,這佔了一個 int64 欄位寫入時間的 67%;移除它使寫入快了 3.3 倍,且輸出位元組相同。
- 借用而非複製,並直接解碼到最終的佈局中,例如將定義層級直接解碼到有效性位元圖中。
- 在資料具有連續性的地方成批處理:在 1% 空值率下,約 250 次
memmove取代了約 25,000 次的逐 slot 分支。 - 知道何時一個階段已經達到記憶體頻寬限制:一個移動了 18.3 MB 進出記憶體、耗時 1.2 ms(約 30 GB/s)的階段,只能被跳過。
- 在添加執行緒之前先找出最慢的項目。 按 column chunk 分割導致一次讀取的上限只有 1.39 倍,因為一個大的欄位獨佔了所有時間。
像測試答案一樣測試優化:故意破壞它,然後檢查測試是否失敗(references/testing-optimizations.md)。Mojo 自身的一些陷阱,例如值在其最後一次使用時消亡和 MutUntrackedOrigin,記錄在 references/mojo-hazards.md 中。
取得可信賴的數值
references/measurement.md 包含了完整的規則。以下是在 magmalake 中改變了結論的規則:
- 報告中位數 (p50) 以及 p90,絕不要只報告平均值。 一次緩慢的首次呼叫會使平均值偏移 30%,而 p50 則完全不受影響。
- 預先建立基準測試的可執行檔;絕不要在
mojo run下計時。 JIT 曾使一個基準測試膨脹了 1.9 倍。 - 檢查機器是否處於穩定狀態。 同一個基準測試在空閒時耗時 31 毫秒,而在 16 個忙碌執行緒旁邊則耗時 89 毫秒。
- 指明參考資料的 API 和執行緒數量。
pq.read_table是 pyarrow 的 dataset scanner,而不是 parquet-cpp 的讀取路徑,並且 pyarrow 的執行緒設定會使一個檔案的讀取時間變化 2.9 倍。 - 當兩個數值不一致時,先檢查它們是否以相同方式產生,然後再解釋差異。
在宣稱加速之前
- 比較前後答案,在可能的情況下進行逐位元比較(見上一條迴圈規則)。一個更快但錯誤的答案比一個緩慢但正確的答案更糟。
- 報告兩個數值以及設定:列數、執行緒數、本地磁碟或 S3,以及快取是否為熱的。
- 說明還有哪些待辦事項。 在 shuffle 修復之後,777 毫秒中仍有 325 毫秒花在複製欄位上,而且在多執行緒情況下,工作仍然集中在一個執行緒上。