Community程式設計與開發github.com

magmalake/mojo-data-performance

撰寫高效的 Mojo 資料處理程式碼,精確定位時間花費,從管線的逐列迴圈到底層的讀取器、寫入器與編解碼器,並取得經得起反覆驗證的基準測試數據。

mojo-data-performance 是什麼?

本技能提供一套系統性方法,協助開發者在使用 Mojo 語言處理資料(特別是 Arrow Record Batch、Parquet 或 Iceberg 掃描)時,分析並優化效能。它不僅涵蓋逐列迴圈的微觀優化技巧(如指標索引、分支預測、快取佈局),還深入探討整個管線的瓶頸,包括多執行緒佇列、I/O 並行度與共享資源管理。技能內含大量來自 magmalake 專案的真實量測數據,讓每條規則都有明確的效能改善幅度。支援的平台包括 Modular 的 Mojo 生態系,並與 mojo-syntax 技能互補。透過先測量掃描本身作為基準,再比較工作階段的耗時,開發者能快速判斷瓶頸所在,避免盲目優化。此外,本技能也提供如何取得可信賴基準測試數據的指南,指出常見陷阱(如 JIT 編譯膨脹、系統干擾),並要求比較前後結果時必須保持一致的方法。

相容平台~Claude Code~Codex CLI✓Cursor
npx skills add https://github.com/magmalake/skills/tree/HEAD/mojo-data-performance

Installed? Explore more 程式設計與開發 skills: steipete/bluebubbles, steipete/eightctl, steipete/blucli · View all 6 →

在你喜歡的 AI 中提問

開啟一個已預先載入此 Agent Skill 的新對話。

說明文件

資料管線的時間主要花費在三個地方:讀取和解碼資料、對每一列進行的工作,以及等待(等待儲存、等待鎖、或等待忙碌的執行緒)。解碼的部分通常已經在函式庫中調整過了。而每一列上的工作是你寫的,這正是粗心的迴圈會導致 2–3 倍效能損失的地方。

本檔案大部分內容在討論處理掃描 (scan) 產生的批次 (batch) 的程式碼。關於底層的讀取器、寫入器或編解碼器,以及如何取得可信賴的數值,請參閱最後兩節及其參考資料。

首先,測量時間花在哪裡

你很少會有第二個實作(如 Ray、DuckDB、pyarrow)可以作為比較基準。不如直接使用掃描本身作為衡量標尺:

  1. 單獨計時掃描。 執行相同的掃描,使用相同的欄位、過濾條件和 ScanOptions,然後丟棄這些批次。這是一個純解碼的基準線:沒有任何管線能比它更快。
  2. 單獨計時你的工作。 將每次批次處理的呼叫(fold、add,或任何消耗該批次的函式)用 perf_counter_ns() 包起來,並將時間加總。
  3. 比較它們。 在單一執行緒上,總時間接近解碼時間的 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 毫秒花在複製欄位上,而且在多執行緒情況下,工作仍然集中在一個執行緒上。

Individual skills in this repo

This repo contains 2 individual skills — each has its own dedicated page.

magmalake/dista

Write and run distributed Mojo programs with dista, magmalake's Ray-style task framework, where one compiled binary is the driver, the head and every worker. Use when writing or reviewing Mojo code that imports dista or mentions distributed Mojo, a remote task, `Task`/`Task2`/`TaskMany`/`TaskWith`, `ObjectRef`, `rt.remote`, `rt.get`, `Seq`, an actor with methods or checkpoints, a cluster, a head, workers, `DISTA_ROLE`/`DISTA_HEAD`, the `dista` command (status, drain, scale, timeline, stop), draining, autoscaling, or a dista runtime inside an orca pipeline step; and when a dista job hangs, is refused with a registry mismatch, or fails with a lost result. Pairs with Modular's mojo-syntax skill and magmalake's mojo-data-performance skill.

magmalake/magmalake-apps

Build a data application on magmalake as one Mojo project, with pipelines run by orca, a lake of Iceberg tables over S3, and a long-running flare web app reading them, then deploy it to AWS with `orca aws`. Use when starting a new magmalake project or laying one out (pixi, mojoshelf, `bin/`, `pipelines/`, `src/`); when a pipeline step writes Iceberg tables that an app reads; when writing a flare web app, console or dashboard over the lake; when running the project locally with Postgres and a local warehouse; or when deploying to AWS (`orca aws up`, `push`, `register`, `app`, `status`, `down`), setting up its hostnames, secrets, public bucket or costs. Pipelines in depth (partitions, keyed, sensors, dista tasks) are the `orca` skill's.

相關技能