设计与实现全解析)
TiDB 非并行 HashAgg 内存溢写Spill to Disk设计与实现全解析【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb本文以 设计文档 docs/design/2021-06-23-spilled-unparallel-hashagg.md 为核心结合当前仓库中 agg_spill.go、agg_hash_executor.go 等源码实现系统讲解 TiDB 聚合执行器在 SQL 内存超限时如何通过分轮次溢写 磁盘回放算法继续完成聚合而非直接 Kill 查询。读完本文你将理解该设计的背景、三步式溢写算法、单线程/并行 HashAgg 各自的触发与回放机制以及相关的测试与风险边界。一、Motivation内存超限为什么要溢写而非杀死在 TiDB 中聚合执行器的计算逻辑分两类并行parallel与非并行unparallel。设计文档明确指出当时的痛点当 SQL 内存使用超过 memory quota 时两种实现都无法利用外部存储来管控内存只能 kill 正在执行的 SQL。对于哈希聚合HashAgg这类算子中间结果被组织为一张 group key 到部分聚合结果的哈希表内存占用随元组tuple插入而单调增长。文档给出的解决思路是引入spilling 算法让聚合在内存不够时把暂时无法归并的新 key 写到磁盘待内存中已能归并的数据处理完后再分批读回磁盘数据继续归并。这样 SQL 便能在内存受限环境下继续运行而非被终止。补充一句设计约束来自文档——本提案只描述非并行unparallel / 单线程HashAgg的溢写并行 HashAgg 的溢写后续才支持。不过在当前的仓库主干中这一后续工作也已完成并行 HashAgg 的溢写由 parallelHashAggSpillHelper 支撑本文一并讲解。二、核心算法三步式溢写流程设计文档给出的是抽象层面的三步算法是所有实现单线程与并行共同遵守的主干当内存占用高于 mem-quota 时将 HashAgg executor 切换为 spill-modeHashAgg 处于 spill-mode 时hashMap 中的元组不再持续增长a. 若处理中的 key 已存在于 Map 中则原地聚合结果并集b. 若处理中的 key 不在 Map 中则把该数据溢写到磁盘所有数据被处理后输出 Map 中的聚合结果并清空 Map再读回磁盘上的溢写数据重复步骤 1~3直到全部数据都被聚合完成。直观理解内存配额就像一个固定容量的缓存容量被占满后每来一个新 key 都先记到磁盘上的等待区等下一轮把已容纳的 key 全部输出、清空缓存后再把等待区数据当作新一轮输入继续处理。这也解释了为什么该算法要求在某一轮内已存在于 Map 的 key 仍需聚合——同一 key 可能在本轮不同 chunk 中反复出现先归并能减少回放轮次与磁盘 IO。三、从设计到代码单线程 HashAgg 的溢写实现设计文档成文于 2021 年当前仓库中对应的实现位于 pkg/executor/aggregate。下面把文档中的抽象算法对应到真实代码逐个环节印证。3.1 溢写状态机与关键阈值在 agg_spill.go 中溢写能力由一组状态与常量支撑状态机spillStatusnoSpill正常→needSpill需要触发→inSpilling正在溢写→spillTriggered已触发maxSpillTimes 10单个算子最多溢写 10 次防止极端场景下无限回放磁盘spilledPartitionNum 256并行实现中把磁盘数据划分为 256 个分区用于支持多轮次分治回放。3.2 AggSpillDiskAction内存追踪器的回调单线程 HashAgg 的溢写触发依赖 TiDB 的内存追踪memory tracker机制。在 agg_spill.go 中定义了AggSpillDiskAction实现memory.ActionOnExceed接口作为内存超限时的回调Action当查询内存配额被超过时若算子尚未处于 spill mode、且满足hasEnoughDataToSpill则把HashAggExec.inSpillMode原子地置为 1并累加memory.QueryForceDisk指标hasEnoughDataToSpill的判断是aggTracker.BytesConsumed() passedInTracker.GetBytesLimit()/5即已处理数据的内存占用至少达到阈值上限的 20% 才允许溢写注释说明这是为避免溢写过于频繁若以上条件不满足例如数据量本身过小不值得溢写则通过GetFallback()触发链上的兜底 Action即原有 OOM 处理路径。触发时 agg_spill.go 会打印日志memory exceeds quota, set aggregate mode to spill-mode并带上当前 consumed / quota / spillTimes 信息。3.3 execute 主循环中的分流HashAggExec 的主循环位于 agg_hash_executor.go。对每个输入 chunk 计算完 group key 后逐行执行以下分流逻辑若 group key已存在于groupSet正常调用UpdatePartialResult更新部分聚合结果Map 不增长若 key不存在则检查是否处于 spill mode处于 spill mode 时inSpillMode 1且 Map 非空不插入哈希表而把该行选入sel集合等待整体落盘未处于 spill mode 时才执行groupSet.Insert正常扩容。chunk 处理完毕后若sel非空调用spillUnprocessedData把被选中的新 key 行写入磁盘。该函数做了两种处理agg_hash_executor.go若整 chunk 都被选中直接把childResult整体交给dataInDisk.Add否则逐行追加进临时 chunktmpChkForSpill满了就dataInDisk.Add一次并 Reset。注意 execute 的 defer 里还会把残留不满一个 chunk 的tmpChkForSpill最后补写入磁盘保证不丢行。3.4 磁盘存储与状态重置HashAggExec 结构中持有dataInDisk *chunk.DataInDiskByChunksagg_hash_executor.go用于按 chunk 方式把溢写行序列化到磁盘并将磁盘跟踪器挂到算子级diskTracker上以统计落盘用量。一轮内所有输入处理完后执行器进入输出与重置阶段resetSpillMode见 agg_hash_executor.go输出并清空groupSet/partialResultMap重置 group key 游标与已用内存统计记录当前已落盘的 chunk 数numOfSpilledChks dataInDisk.NumChunks()清空inSpillMode若还有未读回的落盘数据numOfSpilledChks ! NumChunks()则把磁盘数据重新作为输入继续聚合从而完成设计文档中重复步骤 1~3 直到全部数据聚合完的循环。这就是设计文档第 3 步输出结果 → 清空 Map → 读回磁盘 → 重复在代码中的完整闭环。四、并行 HashAgg 的溢写文档列为 Future Work仓库中已落地设计文档在 Future Work 中列出支持并行 HashAgg 的溢写。当前仓库已将这一项落实实现同样位于 agg_spill.go 与 agg_hash_partial_worker.go、agg_hash_final_worker.go 等文件中。并行实现的核心区别在于需要多 worker 协作因此引入了parallelHashAggSpillHelper它内部用一把锁保护状态waitIfInSpilling条件变量让 partial worker 在正在溢写期间阻塞等待溢写结束后由setSpillTriggered广播唤醒spilledChunksIO按 256 个分区维护每轮产生的磁盘 IO 列表restoreOnePartition从分区序号递减方向取一个分区把该分区内所有落盘文件读回并重建 partial result回放时对每一行调用processRowkey 已存在则用 final worker 的聚合函数MergePartialResult归并key 不存在则重建新的 partial result 槽位。读回的数据用与 partial worker 相同的 partial 聚合函数反序列化aggFuncsForRestoring保证中间结果格式一致。此外 agg_spill_test.go 中的TestGetCorrectResult、TestDistinctAggGetCorrectResult、TestCheckChunkSpill等测试覆盖了正确性 溢写触发 chunk 落盘三类场景TestFallBackAction则验证了无条件溢写时回退到兜底 OOM 处理的行为。五、功能与场景测试设计设计文档对测试提出了明确要求按维度归纳如下可对照 agg_spill_test.go 中的实现逐一验证测试维度文档要求期望结果功能测试使用聚合函数的查询应返回正确结果溢写前后结果一致场景 1单线程 HashAgg 超内存配额溢写生效降低内存占用SQL 成功运行场景 2并行 HashAgg 超配额在实现并行溢写前SQL 先被取消把agg-concurrency设为 1 后可成功运行场景 3数据 ndv不同值数量低、SQL 含 distinct 函数原本会被取消该特性助其成功运行场景 4数据 ndv 高、SQL 含 distinct 函数内存确实不足时仍能正确取消不产生错误结果兼容性测试—N/A不涉及格式/行为兼容变更基准测试非溢写场景不得有明显性能回退性能回退应 2%场景 2 补充说明并行 HashAgg 本身的多 worker 哈希表在溢写实现前无法落盘因此只能依赖把并行度降为 1 退化为单线程或在仓库当前版本中由并行溢写实现接管。场景 3/4 则指向 distinct 聚合——distinct 需要在 Map 中维护已见值集合其内存增长模型特殊是后续工作重点。六、Impacts Risks 与 Future Work6.1 已知影响与风险设计文档坦诚列出一个关键风险点对 distinct 聚合函数而言即使不增加 HashMap 中的新元组数量内存仍可能增长。原因在于 distinct 的去重语义需要在哈希表中保存已出现过的值以判断重复这部分状态与 group key 数量并不完全对应因此溢写触发后 Map 的规模虽被冻结内存却未必同步停止增长。这正是文档把distinct 的友好溢写列入后续工作的原因。6.2 后续工作设计文档给出两项明确的 Future Work截至当前仓库源码其进展如下为 distinct 聚合函数实现更友好的溢写仍属持续优化方向相关工作围绕 distinct 场景内存模型展开支持并行 HashAgg 的溢写已在仓库主干落地实现主体即上文parallelHashAggSpillHelperagg_spill.go。七、小结与阅读路径从一份 2021 年的设计提案出发TiDB 的 HashAgg 已经走过了内存超限即 Kill SQL → 单线程溢写 → 并行溢写的演进路径。对开发者而言最有价值的三个可验证锚点分别是算法主干设计文档中的三步式溢写docs/design/2021-06-23-spilled-unparallel-hashagg.md单线程触发与回放AggSpillDiskAction与execute/resetSpillModeagg_hash_executor.go并行多 worker 协作parallelHashAggSpillHelper与 256 分区回放agg_spill.go以及完整的功能/场景测试agg_spill_test.go。阅读上述文件时可重点观察maxSpillTimes、hasEnoughDataToSpill的 20% 门槛、256 个 spill 分区等常量如何在工程上把抽象算法落地为可控、可观测、可测试的实现——这些细节决定了该特性既能兜底内存超限又不会因频繁落盘造成明显的性能回退。【免费下载链接】tidbTiDB is built for agentic workloads that grow unpredictably, with ACID guarantees and native support for transactions, analytics, and vector search. No data silos. No noisy neighbors. No infrastructure ceiling.项目地址: https://gitcode.com/GitHub_Trending/ti/tidb创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考