Apache DataFusion中SortMergeJoinStream的结构优化实践

背景与问题分析

在Apache DataFusion项目的物理计划执行模块中，SortMergeJoinStream是实现排序合并连接(Sort Merge Join)的核心数据结构。随着功能迭代，该结构体逐渐膨胀为包含大量字段的复杂实现，导致代码可读性和维护性下降。这种现象在长期演进的开源项目中十分常见——初始设计时未充分考虑结构划分，随着功能叠加逐渐形成"上帝对象"。

结构体现状分析

原始的SortMergeJoinStream结构体主要存在三个典型问题：

1. 字段混杂：流处理状态、缓冲管理、输出控制等不同维度的字段混杂在一起
2. 职责模糊：单个结构体同时承担状态维护、过程控制和结果生成等多重职责
3. 扩展困难：新增功能时难以快速定位相关字段，容易造成代码冲突

优化方案设计

经过社区讨论，确定采用渐进式重构策略：

第一阶段：逻辑分组

首先通过注释分隔将字段按功能划分为：

• 流处理状态组：包含流式输入的状态跟踪、当前处理批次等信息
• 缓冲管理组：管理缓冲数据的加载状态、连接标记等
• 输出控制组：处理输出批次的构建、大小控制等
• 配置参数组：保存连接类型、批大小等不变参数

这种分组方式不改变原有代码结构，仅通过注释提高可读性，风险最小且立即可见效果。

第二阶段：结构分解

在后续功能开发时，逐步将逻辑分组转化为实际子结构体，例如：

struct StreamState {
    current_batch: Option<RecordBatch>,
    row_index: usize,
    join_markers: Vec<bool>,
    // ...其他流相关字段
}

struct BufferManager {
    batches: Vec<RecordBatch>,
    current_index: usize,
    // ...其他缓冲相关字段
}

技术价值

这种重构方案体现了几个重要的工程实践原则：

1. 渐进式改进：通过注释分组先行，降低重构风险
2. 关注点分离：将连接算法的不同维度解耦
3. 可测试性：独立的状态结构更便于单元测试
4. 可扩展性：为未来支持spill等功能预留结构空间

经验总结

对于复杂算子的重构，DataFusion社区给出了一个经典范例：

1. 先通过非侵入式改进提高可读性
2. 在后续功能开发时自然演进结构
3. 保持每个变更小而可控
4. 充分利用社区协作进行设计验证

这种模式特别适合在保证系统稳定性的前提下，持续改善代码质量。对于其他面临类似问题的分布式计算项目具有很好的参考价值。