Apache DataFusion聚合执行任务无法取消问题分析

Apache DataFusion是一个高性能的查询引擎，但在其聚合执行(AggregateExec)的实现中存在一个关键问题：聚合操作无法被及时取消。这个问题在用户尝试中断长时间运行的聚合查询时尤为明显。

问题本质

DataFusion的AggregateExec在执行聚合操作时，其内部流(stream)实现会在一个循环中持续消费输入数据，而不会主动让出(yield)执行权。当输入源是文件等不会返回Pending状态的流时，Tokio运行时将没有机会中止正在运行的任务。

这个问题通常被查询计划中的CoalesceExec所掩盖，因为CoalesceExec会在单独的任务中运行聚合操作。虽然CoalesceExec使用的RecordBatchReceiverStream确实会返回Pending状态，使得它本身可以被取消，但实际聚合任务的停止仍然需要等待Tokio能够获得执行权。

问题复现

要复现这个问题，可以按照以下步骤操作：

1. 启动datafusion-cli
2. 执行SET datafusion.execution.target_partitions = 1;
3. 在一个较大的表上执行SELECT sum(column) from table;
4. 尝试使用Ctrl-C取消查询

预期行为是查询应该立即停止，但实际行为是查询会继续执行直到完成。

技术分析

问题的核心在于DataFusion的流处理模型与Tokio任务调度的交互方式。在Tokio的异步模型中，任务只能在显式让出执行权时才能被取消。当AggregateExec的实现持续处理数据而不让出执行权时，即使外部触发了取消操作，Tokio也无法立即中断任务。

这个问题在单线程运行时尤为明显，因为此时没有其他任务可以抢占执行权。在多线程运行时，虽然其他任务可以继续执行，但聚合任务本身仍然会消耗CPU资源直到完成。

解决方案

解决这个问题的关键在于在聚合操作的执行流中定期插入让出点(yield point)。目前社区提出了几种解决方案：

1. 在AggregateExec的输入流中包装一个定期让出的适配器，例如每处理64个批次后让出一次
2. 修改CoalescePartitionExec，使其在单分区情况下也使用任务分离的方式执行
3. 在查询执行的顶层添加取消检查点

这些解决方案都需要在保证查询性能不受显著影响的前提下，提供及时取消的能力。性能测试表明，合理的让出频率对查询性能的影响可以控制在可接受范围内。

影响范围

这个问题不仅影响命令行工具的交互体验，还会影响所有通过编程接口使用DataFusion的场景。例如：

• Java应用通过JNI调用DataFusion时无法及时取消查询
• gRPC服务端处理客户端断开连接时无法及时终止查询
• 任何需要设置查询超时的场景

结论

DataFusion的聚合执行取消问题是一个典型的异步任务管理挑战。通过合理地在执行流中插入让出点，可以在保持高性能的同时提供良好的取消支持。这个问题也提醒我们在设计流处理系统时，需要考虑任务管理的各个方面，包括但不限于性能、资源使用和用户交互体验。