DataFusion项目参数类型推断测试的优化实践

在Apache DataFusion项目中，参数类型推断是一个重要功能，它允许SQL查询使用参数化查询（如$1、$2等占位符），并在执行前推断这些参数的类型。然而，当前的测试实现存在大量重复代码，使得测试维护和理解变得困难。

问题现状

当前测试代码中，每个参数类型推断测试都包含以下重复部分：

1. 创建SQL查询字符串
2. 生成逻辑计划
3. 验证参数类型推断结果
4. 替换参数值为实际值
5. 验证最终逻辑计划

这种模式导致测试代码冗长且难以一目了然地看出测试覆盖了哪些场景。更重要的是，当需要修改测试逻辑时，需要在多个地方进行相同的更改，增加了维护成本。

解决方案设计

为了解决这个问题，我们可以引入一个专门的测试结构体ParameterTest，它将封装所有测试相关的信息：

struct ParameterTest {
    sql: &'static str,                // 测试SQL语句
    expected_types: Vec<(&'static str, Option<DataType>)>, // 预期参数类型
    param_values: Vec<ScalarValue>,   // 参数值
}

这个结构体将提供一个run方法，该方法会：

1. 根据SQL创建逻辑计划
2. 验证参数类型推断结果是否符合预期
3. 使用提供的参数值替换占位符
4. 返回包含初始和最终逻辑计划的格式化字符串

优化后的测试示例

优化后的测试代码将更加简洁明了：

#[test]
fn test_infer_types_from_predicate() {
    let test = ParameterTest {
      sql: "SELECT id, age FROM person WHERE age = $1",
      expected_types: vec![("$1", Some(DataType::Int32))],
      param_values: vec![ScalarValue::Int32(Some(10))],
    };
 
    assert_snapshot!(test.run(), @r#"
 ** Initial Plan:
    Projection: person.id, person.age
      Filter: person.age = $1
        TableScan: person
 ** Final Plan:
    Projection: person.id, person.age
      Filter: person.age = Int32(10)
        TableScan: person
    "#);
}

技术优势

1. 代码复用：将公共测试逻辑封装在ParameterTest中，避免重复
2. 可读性提升：测试意图更加清晰，一眼就能看出测试场景
3. 维护简化：修改测试逻辑只需在一个地方进行
4. 覆盖可视化：更容易看出哪些场景已被覆盖，哪些尚未测试
5. 快照测试：结合insta的快照测试功能，可以轻松验证逻辑计划的正确性

实现细节

在实际实现中，ParameterTest的run方法将执行以下步骤：

1. 使用DataFusion的SQL解析器将SQL转换为逻辑计划
2. 调用get_parameter_types方法获取参数类型推断结果
3. 验证推断结果与预期一致
4. 使用with_param_values方法替换参数占位符
5. 格式化初始和最终逻辑计划为字符串
6. 返回包含两个计划的可读性强的字符串

这种设计不仅适用于简单的参数推断测试，还可以轻松扩展到更复杂的场景，如多参数、嵌套查询、不同类型参数等测试用例。

总结

通过引入ParameterTest结构体，我们显著提升了DataFusion项目中参数类型推断测试的可维护性和可读性。这种模式不仅适用于当前场景，也可以作为其他类似测试的参考模板，体现了良好的测试代码设计原则。