andygrove opened a new pull request, #7543: URL: https://github.com/apache/datafusion-comet/pull/7543
## Which issue does this PR close? N/A. This was split out of https://github.com/apache/datafusion-comet/pull/7289. ## Rationale for this change `spark.comet.exec.aggregate.skipPartial.enabled` lets a native partial aggregate stop aggregating once its input looks mostly distinct and send the rest of the task's rows to the shuffle as partial states. It admits only partial aggregates whose functions are all single-argument `COUNT`, so a high-cardinality `GROUP BY` that also computes an `AVG` never benefits. DataFusion builds each row's partial state with `GroupsAccumulator::convert_to_state`, which Comet's `AVG` accumulators did not implement. ## What changes are included in this PR? - `AvgGroupsAccumulator` (Comet's Float64 `AVG`, which also serves integral inputs after a cast to double, and small decimals that Spark rewrites to `avg(UnscaledValue(x))`) now implements `convert_to_state`. Each row becomes its value and a count of one, or a zero sum and a zero count for a null or filtered-out row, which is what `state()` emits for such a group. - `configure_skip_partial_aggregation` admits `AVG` next to `COUNT`, only when the aggregate's result type is `DOUBLE`. A decimal result is built as `AvgDecimal`, whose `convert_to_state` returns a "not implemented" error, so admitting it would fail a query once skipping started. A decimal `AVG` keeps the plan from skipping, as before. Implementing the decimal conversion is left out: its partial state carries overflow in the shared null buffer of the sum and count, which a per-row conversion would have to reproduce exactly. - The `.doc` of `spark.comet.exec.aggregate.skipPartial.enabled` and the Adaptive Partial Aggregation section of the tuning guide describe the admitted aggregates. This remains opt-in behind the existing experimental `spark.comet.exec.aggregate.skipPartial.enabled` setting, which defaults to `false`. Nothing changes for a query that does not enable it. ## How are these changes tested? - Rust: `convert_to_state_matches_one_row_groups` checks that converting rows gives the same states as aggregating each row alone into its own group, including a null row and a filtered-out row. `skip_partial_eligibility_is_fail_closed` covers the gate: `AVG` with a `DOUBLE` result is admitted, alone and next to `COUNT`; `AVG` without an input, without a result type, with a decimal result (alone and next to `COUNT`), or with a `FLOAT` result is not. - Scala: `skip partial aggregation converts AVG rows to partial states` in `CometAggregateSuite` runs high-cardinality `GROUP BY` queries over `avg` of an int, a double, and combinations with `count`, and over `avg` of a `DECIMAL(38, 10)`, with nulls in the values and keys that repeat across tasks. It lowers the probe row threshold to 100 so skipping starts, checks each answer against Spark, and checks the `skipped_aggregation_rows` metric: positive for the int and double cases, zero for the decimal case. Commands run: `cargo test -p datafusion-comet-spark-expr avg`, `cargo test -p datafusion-comet --lib skip_partial`, `cargo clippy --all-targets --workspace -- -D warnings`, `cargo fmt --check`, and `./mvnw test -Dtest=none -Dsuites="org.apache.comet.exec.CometAggregateSuite skip partial aggregation"` on JDK 17. -- 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]
