Apache Arrow DataFusion 对 Parquet 浮点类型数据剪枝的优化实践

在数据分析领域，数据剪枝（Data Pruning）是一项关键技术，它通过利用数据文件的统计信息（如最大值、最小值等）来减少实际需要读取和处理的数据量。本文将深入探讨 Apache Arrow DataFusion 项目中对 Parquet 格式文件中浮点类型数据剪枝功能的优化过程。

背景与问题发现

在 DataFusion 的查询优化过程中，我们发现一个有趣的现象：当查询条件涉及整数类型或双精度浮点类型（Double）列时，系统能够有效利用 Parquet 文件的统计信息进行数据剪枝；但当查询条件涉及单精度浮点类型（Float）列时，剪枝功能却未能生效。

通过分析查询计划（EXPLAIN），我们发现当查询条件包含类似 float_col > 10.0 这样的表达式时，系统会将 10.0 识别为双精度浮点数（Float64），进而导致 float_col 被强制转换为 Float64 类型。这种类型转换触发了 DataFusion 中一个名为 verify_support_type_for_prune 的验证函数，该函数原本只支持有限的数据类型组合，从而阻止了剪枝优化的应用。

技术深入分析

DataFusion 的物理优化器中包含一个重要的剪枝模块（pruning module），其核心功能是验证查询条件中的数据类型是否支持剪枝操作。原始实现中，这个验证逻辑相对保守，主要支持以下几种情况：

1. 相同类型之间的比较
2. 整数类型之间的比较（如 Int32 与 Int64）
3. 时间戳类型之间的比较

这种保守的设计源于对数据类型转换和比较语义的谨慎考虑，特别是对于浮点类型，由于涉及 NaN（非数字）和无穷大等特殊值的处理，其比较行为比整数类型更为复杂。

解决方案探索

经过深入分析，我们识别出两种潜在的解决方案：

1. 扩展验证函数支持的数据类型：修改 verify_support_type_for_prune 函数，使其支持更广泛的数据类型组合，特别是浮点类型之间的比较。

2. 优化类型强制转换逻辑：在查询解析阶段，将浮点数字面量自动转换为更合适的类型（如将 10.0 识别为 Float32 当与 Float32 列比较时）。

经过讨论和验证，我们最终选择了第一种方案，因为它更直接且能解决更广泛的问题。具体实现中，我们将验证逻辑简化为检查参与比较的类型是否都是数值类型（通过 is_numeric() 方法判断），这样不仅解决了浮点类型的问题，还同时支持了无符号整数等其他数值类型。

实现细节与挑战

在实现过程中，我们遇到了几个技术挑战：

1. NaN 处理：Parquet 文件中的 NaN 值在比较时具有特殊语义（NaN 不等于任何值，包括它自己）。这与 SQL 标准中的比较行为一致，但需要确保剪枝逻辑不会错误地排除包含 NaN 的数据页。

2. 类型转换一致性：需要确保剪枝阶段使用的比较语义与后续实际数据过滤阶段完全一致，避免出现剪枝阶段排除了本应保留的数据，或者保留了本应排除的数据。

3. 性能考量：更宽松的类型验证可能会增加一些运行时开销，但经过测试，这种开销与剪枝带来的性能提升相比可以忽略不计。

实际效果验证

优化后的实现带来了显著的性能改进：

1. 对于包含 float_col > 10.0 条件的查询，现在能够正确利用 Parquet 文件的统计信息进行剪枝，减少了需要读取的数据量。

2. 查询计划中正确显示了 pruning_predicate 信息，表明剪枝优化已被应用。

3. 特殊值（如 NaN）的处理符合预期，没有引入错误的结果。

总结与展望

通过对 DataFusion 剪枝功能的这一优化，我们不仅解决了浮点类型剪枝的问题，还为更广泛的数值类型比较奠定了基础。这一改进特别有利于处理大规模浮点数据集的查询性能。

未来，我们可以进一步探索：

1. 支持更复杂的数据类型和比较操作
2. 优化复合数据类型（如结构体和数组）的剪枝策略
3. 提高剪枝精度，减少假阳性（即被错误保留的不必要数据）

这一优化案例展示了在查询引擎开发中，平衡功能丰富性与正确性的重要性，以及通过深入理解类型系统和查询执行流程来提升系统性能的有效方法。