Community编程与开发github.com

magmalake/mojo-data-performance

编写高性能 Mojo 数据处理代码,精确定位时间消耗(从管道的每行循环到底层的读取器、写入器和编解码器),并获取可重复验证的基准测试数据。

mojo-data-performance 是什么?

这是一个面向 Mojo 数据管道开发和性能优化的 AI 技能。它帮助开发者识别和处理数据处理中的性能瓶颈,包括数据读取解码、每行处理逻辑和等待开销(存储、锁或线程竞争)三个核心环节。技能提供了从宏观到微观的分层测量方法:先用纯扫描时间作为基准(解码基线),再测量实际工作耗时,通过对比确定瓶颈在哪一层。针对每行循环提出了具体的优化规则(如缓存原始指针、避免分支、使用紧凑数据结构等),每条规则都附有实测数据支撑。它还覆盖了读取器、写入器和编解码器的优化技巧(如避免重复计算、借用代替拷贝、批量处理),以及如何获得可信的基准测试数据(使用中位数而非平均值、预编译基准测试、确保测试环境稳定)。该技能支持 Arrow 记录批次、Parquet、Iceberg 扫描的 Mojo 代码,以及与 pyarrow、parquet-cpp、arrow-rs、Ray 等的对比场景。通常需要与 Modular 的 mojo-syntax 技能配合使用。

兼容平台~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 倍。

本文件的大部分内容涉及消耗扫描批次的代码。对于底层的读取器、写入器或编解码器,以及获取可信的性能数据,请参阅最后两节及其参考资料。

首先,测量时间消耗在哪里

你通常没有第二个实现(Ray、DuckDB、pyarrow)可以对比。相反,使用扫描本身作为基准:

  1. 单独计时扫描。 使用相同的列、过滤条件和 ScanOptions 运行相同的扫描,并丢弃批次。这是纯解码基线:任何管道的速度都不可能快于此。
  2. 单独计时你的工作。 将每批次的调用(fold、add 或其他消耗批次的函数)包裹在 perf_counter_ns() 中,并累加时间。
  3. 比较它们。 在单线程上,总时间在纯解码时间的约 1.5 倍以内意味着工作开销很小;超过此值,工作就是瓶颈,值得优化。在多线程上,还要检查工作是否在单线程上运行,而解码在多线程上运行。

实测示例:对 7700 万条行程记录进行分组聚合在单线程上耗时 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 秒
将所有累积的键侧数据放在一个槽位中(结构体数组),而不是并行数组,这样一次更新只触及一条缓存行。配合下一条规则,时间从 2.09 秒降至 1.40 秒
无分支过滤:计算每一行,将其写入 kept 索引,然后通过 Int(keep) 递增 kept。然后在第二个紧凑循环中添加保留的行。与上一条规则一起测量
当键为小整数时(zone × zone × hour 约 170 万),使用密集数组按索引键而不是使用 Dict。未单独测量;两个版本都使用了密集数组
优化时不要改变浮点数加法的顺序。这样答案必须与旧答案逐位匹配,这提供了一个免费的正确性检查。上述每一步都精确匹配

整个管道的规则

规则实测效果
不要让一个线程消耗多个线程解码的内容。为每个工作线程提供自己的累加器,并在最后合并它们,就像 dista 的 map 任务所做的那样。在 AWS 上,16 个线程仅将单进程分组聚合从 202 秒降至 126 秒
即使只有一个解码线程,也要保持多个读取操作同时进行:iceberg.mojo 中的 ScanOptions.files_at_once 和 overlap_io。从 S3 读取的单线程加速比为 6.6–6.9×
每个进程加载一次共享 C 库并保持句柄,绝不要每次调用都加载。每次调用都打开 zstd 时,扫描花费 52% 的时间在 dyld 上;修复后,时间从 2687 毫秒降至 262 毫秒
缓存任何逐字节重复解析的内容,例如清单中的 Avro 模式。清单规划固定成本下降 3.6×(从 113.7 µs 降至 31.9 µs)
不要按记录分配内存:解码到可重复使用的槽位中。Avro 读取器在其游标之前为每条记录分配了约 60 次内存
了解上限:超过几个核心后,扫描受限于内存带宽,更多线程无济于事。应读取更少的字节。Iceberg 扫描在单核上耗时 35.8 毫秒,在四核上耗时 12.1 毫秒,大约 3× 加速

references/case-studies.md 详细讲述了每个案例,以及是什么问题导致了这些问题。

读取器、写入器和编解码器

当解码本身是缓慢部分时,规则主要是关于不做被浪费的工作。详细内容和数据在 references/readers-and-codecs.md 中。

  • 检查每个快速路径的门条件是否能够被接触它的调用者满足,并通过工具证明其运行。一个向量化路径的门条件无法被任何列表列满足,导致一个 9 KiB 列占用了 43% 的读取时间。
  • 不要计算后续决策通常会丢弃的内容。 构建字典只是为了选择纯编码方案占用了 int64 列写入时间的 67%;将其移除使写入速度提高了 3.3 倍,且输出字节完全相同。
  • 借用而不是复制,并直接解码到最终布局中,例如将定义级别解码到有效性位图中。
  • 在数据具有规律性的地方批量处理: 在 1% null 的情况下,约 250 次 memmove 替代了约 25,000 次逐槽位分支。
  • 了解何时某个阶段已达到内存带宽极限: 一个阶段在 1.2 毫秒内移动了 18.3 MB 数据(约 30 GB/s),只能被跳过。
  • 在增加线程之前找到最慢的项目。 按列块拆分使一次读取的加速比上限仅为 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 的数据集扫描器,而不是 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.

相关技能