Apache Arrow DataFusion 排序执行中的 Tokio 阻塞线程优化

在 Apache Arrow DataFusion 项目中，当执行内存密集型排序操作时，系统会将中间结果溢出（spill）到磁盘文件以避免内存不足。然而，在处理大量溢出文件时，特别是在合并（merge）阶段，系统会为每个溢出文件创建一个 Tokio 阻塞线程，这可能导致线程资源耗尽，进而引发查询"挂起"的问题。

问题背景

在 Comet（DataFusion 的一个实现）中，当查询在最小内存配置下运行时，如果生成了大量溢出文件（例如183个），系统会为每个文件创建一个独立的 Tokio 阻塞线程来读取数据。由于 Comet 默认配置的阻塞线程数较少（如10个），这种设计会导致线程资源迅速耗尽，查询无法继续执行。

技术挑战

核心问题在于当前的实现方式：

1. 每个溢出文件在合并阶段都会创建一个独立的 Tokio 阻塞线程
2. 线程数量与溢出文件数量呈线性关系
3. 系统无法动态感知和适应可用的线程资源

这种设计在溢出文件数量较少时表现良好，但当内存压力大、溢出文件数量激增时，就会成为系统瓶颈。

解决方案探索

社区提出了几种改进思路：

1. 多阶段合并策略：

   • 将大量溢出文件分批合并
   • 例如183个文件可先分成约18组，每组10个文件合并为1个中间文件
   • 最后再将18个中间文件合并为最终结果
   • 这种方法会增加约一倍的I/O操作，但能有效控制线程使用

2. 统一合并流优化：

   • 修改 SortPreservingMergeStream 的实现
   • 使其支持多级合并而不仅限于单次合并
   • 这种方案也可应用于使用相同机制的聚合操作

3. 线程资源感知：

   • 理想情况下系统应能感知可用线程资源
   • 但目前 Tokio 不提供查询最大阻塞线程数的API
   • 需要设计合理的默认值和动态调整策略

实现考量

在实际实现中，需要考虑以下因素：

1. I/O并行度与系统负载的平衡：

   • 过多的并行I/O可能导致磁盘争用
   • 需要找到最佳并行度以最大化吞吐

2. 内存使用优化：

   • 多阶段合并会增加临时文件
   • 需要合理控制内存缓冲区大小

3. 错误处理：

   • 确保在资源不足时能优雅降级
   • 提供明确的错误信息而非直接挂起

未来方向

这一优化不仅限于排序操作，还可推广到：

1. 哈希聚合操作中的溢出处理
2. 其他需要大量磁盘I/O的操作
3. 构建通用的溢出文件管理框架

通过系统级的溢出管理策略，可以更合理地利用系统资源，提高大数据量查询的稳定性和性能。

这一优化案例展示了在分布式查询引擎中，资源管理策略对系统稳定性的关键影响，也为类似系统的设计提供了有价值的参考。