andygrove opened a new pull request, #5216: URL: https://github.com/apache/datafusion-comet/pull/5216
## Which issue does this PR close? Part of #5199 (item 4: `DecimalPrecision.promote` is a second full traversal of every expression tree). ## Rationale for this change `QueryPlanSerde.exprToProto` calls `DecimalPrecision.promote` on every expression it converts, and `promote` is a `TreeNode.transformUp`. `transformUp` walks the whole tree and calls `mapChildren` at every node, which allocates through `mapProductIterator` regardless of whether any case matches. The result is a full extra traversal of every expression tree before the serde walks it again, paid on the driver for every query and again for every query stage under AQE. The rule only ever rewrites arithmetic that produces a decimal, so on the overwhelming majority of expressions the traversal produces the identical tree back. Measured with a throwaway probe over decimal-free expression trees of varying size (Spark 4.1 / JDK 17, 20k iterations per measurement, best of 5 after warm-up): | tree nodes | promote before | promote after | share of `exprToProto` before | after | | ---------- | -------------- | ------------- | ----------------------------- | ----- | | 22 | 0.95 us | 0.33 us | 4.6% | 1.3% | | 67 | 2.88 us | 0.97 us | 7.1% | 2.6% | | 247 | 11.36 us | 3.53 us | 8.4% | 2.8% | | 967 | 45.39 us | 18.91 us | 8.7% | 3.9% | Trees that do contain decimal arithmetic pay the guard traversal on top of the rewrite, but that traversal short-circuits at the first match and does not allocate. ## What changes are included in this PR? Adds a `containsDecimalArithmetic` guard and only runs the `transformUp` when it holds. The guard is deliberately a superset of the rule: it matches `Add` / `Subtract` / `Multiply` / `Divide` / `Remainder` whose `dataType` is a `DecimalType`, ignoring the operands that the rule's own patterns check. It is sound to decide this up front because the rewrite only wraps nodes in `CheckOverflow`, and `CheckOverflow`'s data type is its child's, so no node's data type changes part way through the `transformUp` and a tree without decimal arithmetic cannot grow any. Behavior is unchanged: `transformUp` already returns the input tree by identity when no case matches, so the guard removes work rather than changing the result. ## How are these changes tested? Existing coverage: `CometExpressionSuite` (140 tests) and the Spark 4.1 `CometDecimalArithmeticViewSuite` regression tests for #4124 and #5075 all pass. Added `DecimalPrecisionSuite`, which pins the rule's behavior directly and is version-agnostic (the existing unit coverage was Spark 4.1 only). It asserts that trees without decimal arithmetic come back by identity, that each of the five arithmetic operators is wrapped in `CheckOverflow` with the operator's own `dataType` as the target, that a decimal operator buried under nodes the rule does not rewrite is still promoted, and that a mixed tree has only its decimal arithmetic rewritten. These pass both with and without this change, which is the point: they document that the guard is not observable. The probe used for the numbers above was throwaway and is not included. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
