Apache Arrow DataFusion中Aggregate分组表达式重复问题解析

在Apache Arrow DataFusion项目中发现了一个关于Substrait计划解析的有趣问题，当聚合操作中包含重复的分组表达式时，系统无法正确生成输出列。这个问题虽然看似简单，但涉及到查询计划解析和执行的底层机制，值得我们深入探讨。

问题本质

问题的核心在于当Aggregate操作的分组表达式列表中存在完全相同的表达式时，DataFusion无法正确处理这种情况。根据Substrait规范，Aggregate操作的输出列应该按照分组表达式声明顺序排列，后跟度量表达式。但在实际实现中，当遇到重复的分组表达式时，系统会丢失部分输出列。

问题复现

让我们看一个具体的例子。考虑以下Substrait计划：

{
  "relations": [
    {
      "root": {
        "input": {
          "aggregate": {
            "input": {
              "read": {
                "namedTable": {
                  "names": ["data"]
                }
              }
            },
            "groupings": [
              {
                "groupingExpressions": [
                  {"literal": {"i32": 1}},
                  {"literal": {"i32": 1}}
                ]
              }
            ],
            "measures": []
          }
        },
        "names": ["grouping_col_1", "grouping_col_2"]
      }
    }
  ]
}

这个计划定义了一个简单的聚合操作，对表"data"进行分组，使用两个相同的字面量表达式(值都为1)作为分组键。按照Substrait规范，输出应该包含两列："grouping_col_1"和"grouping_col_2"，但实际上系统无法正确处理这种情况。

技术背景

在SQL查询处理中，聚合操作(Aggregate)是一个核心操作，它根据指定的分组表达式将数据分组，然后对每个组应用聚合函数。Substrait作为一种跨系统的查询计划表示格式，定义了Aggregate操作的标准结构。

在DataFusion的实现中，当解析Substrait计划时，系统需要将逻辑计划转换为物理计划。在这个过程中，分组表达式的处理尤为重要，因为它们决定了数据的组织方式和最终结果的列结构。

问题分析

问题的根源在于DataFusion在构建输出列时，可能使用了某种基于表达式哈希的去重机制，导致相同的表达式被误认为冗余而被移除。虽然在实际查询中，使用完全相同的表达式作为分组键确实没有实际意义，但从规范角度讲，系统应该能够正确处理这种情况。

解决方案探讨

针对这个问题，可以考虑以下几种解决方案：

1. 表达式包装方案：在解析阶段自动为Aggregate操作添加一个Project操作，显式地复制缺失的列。这种方法保持了原始计划的完整性，同时确保输出列的正确性。

2. 规范实施方案：严格遵循Substrait规范，确保所有分组表达式无论是否重复，都能在输出中得到体现。这需要修改表达式处理的内部逻辑。

3. 验证警告方案：在计划解析阶段检测重复的分组表达式，发出警告或错误信息，提示用户优化查询。

从工程实践角度看，第一种方案最为稳健，它既保持了兼容性，又不会影响查询执行的正确性。

实际影响

这个问题虽然看似边界情况，但在某些自动生成的查询计划中可能出现。特别是当使用Substrait-Spark等系统生成计划时，可能会产生这种包含重复表达式的计划。因此，DataFusion作为查询执行引擎，应该具备处理这种计划的能力。

最佳实践建议

对于开发者而言，在处理聚合操作时应注意：

1. 在构建查询计划时尽量避免使用完全相同的分组表达式
2. 如果必须处理来自外部系统的计划，应考虑添加预处理步骤来规范化分组表达式
3. 在实现查询计划解析时，应充分考虑各种边界情况，确保规范的完整支持

总结

这个问题的发现和解决过程展示了查询引擎开发中的一些重要考量。作为DataFusion这样的高性能查询引擎，正确处理各种边界情况对于保证系统的健壮性和兼容性至关重要。通过深入分析这个问题，我们不仅解决了具体的技术挑战，也为类似问题的处理提供了参考模式。