数据管道的时间消耗分布在三个地方:数据读取和解码、对每一行执行的工作以及等待时间(等待存储、锁或繁忙的线程)。解码通常在库中已经调优。对每一行的工作是你编写的部分,草率的循环会导致性能下降 2–3 倍。
本文件的大部分内容涉及消耗扫描批次的代码。对于底层的读取器、写入器或编解码器,以及获取可信的性能数据,请参阅最后两节及其参考资料。
首先,测量时间消耗在哪里
你通常没有第二个实现(Ray、DuckDB、pyarrow)可以对比。相反,使用扫描本身作为基准:
- 单独计时扫描。 使用相同的列、过滤条件和
ScanOptions运行相同的扫描,并丢弃批次。这是纯解码基线:任何管道的速度都不可能快于此。 - 单独计时你的工作。 将每批次的调用(
fold、add或其他消耗批次的函数)包裹在perf_counter_ns()中,并累加时间。 - 比较它们。 在单线程上,总时间在纯解码时间的约 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 毫秒仍然用于复制列,并且在多线程上,工作仍然在单线程上执行。