comphead commented on code in PR #7543:
URL: https://github.com/apache/datafusion-comet/pull/7543#discussion_r4532980346
##########
docs/source/user-guide/latest/tuning/operators.md:
##########
@@ -172,7 +172,7 @@ pruning does not require that setting.
Set `spark.comet.exec.aggregate.skipPartial.enabled=true` to let Comet bypass
partial hash
aggregation for high-cardinality grouping when it is not reducing the number
of rows enough. This
experimental optimization is disabled by default. It currently applies only to
fused native
-shuffle-writer plans whose partial aggregates are grouping-only or
single-argument `COUNT`.
+shuffle-writer plans whose partial aggregates are grouping-only,
single-argument `COUNT`, or `AVG` over non-decimal inputs.
Review Comment:
The gate checks for a `DOUBLE` result, not a non-decimal input. Spark
rewrites `avg` over a decimal with precision 11 or lower into
`avg(UnscaledValue(x))`, which has a `DOUBLE` result and is admitted, so `avg`
over `DECIMAL(10, 2)` does skip. Could this line, the matching one in the
Eligibility paragraph below, and the `.doc` in `CometConf` say `AVG` except
over decimals with precision above 11? Both lines here also run past the wrap
the rest of the section uses.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -3431,6 +3431,59 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("skip partial aggregation converts AVG rows to partial states") {
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "input").toUri.toString
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ // Mostly distinct keys within each task, with some keys repeating
across tasks so the
+ // final aggregate merges converted states with each other. Values
include nulls.
+ spark
+ .range(0L, 16384L, 1L, 8)
+ .selectExpr(
+ "id % 12000 AS k",
+ "CASE WHEN id % 5 = 0 THEN NULL ELSE CAST(id % 17 AS INT) END AS
i",
+ "CASE WHEN id % 7 = 0 THEN NULL ELSE CAST(id % 11 AS DOUBLE) / 3
END AS d",
+ "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id % 13 AS DECIMAL(10,
2)) / 7 END AS small",
Review Comment:
`CAST(... AS DECIMAL(10, 2)) / 7` is written as `DECIMAL(14, 6)`
(`Divide.resultDecimalType` with the literal as `DECIMAL(1, 0)`).
`DecimalAggregates` only rewrites `avg` to `avg(UnscaledValue(x))` when
`precision + 4 <= 15`, so `avg(small)` plans as `AvgDecimal` like `avg(big)`
and never skips. The small-decimal path the description mentions is not
exercised. Would dropping the `/ 7` (keeping `DECIMAL(10, 2)`) and asserting
`skipped(avgQuery("avg(small)")) > 0L` cover it? I have not run it.
##########
native/spark-expr/src/agg_funcs/avg.rs:
##########
@@ -365,4 +386,32 @@ mod tests {
assert_eq!(partial.evaluate()?, ScalarValue::Float64(None));
Ok(())
}
+
+ /// Converting rows straight to partial states, as a partial aggregate
that stops aggregating
+ /// does, gives each row the state that aggregating it alone into a group
would.
+ #[test]
+ fn convert_to_state_matches_one_row_groups() -> Result<()> {
+ use arrow::array::BooleanArray;
+ let values: ArrayRef = Arc::new(Float64Array::from(vec![
+ Some(1.5),
+ None,
+ Some(-2.0),
+ Some(4.0),
+ ]));
+ let filter = BooleanArray::from(vec![Some(true), Some(true),
Some(true), Some(false)]);
Review Comment:
The filter here has no null slot, so the `filter.is_valid(idx)` branch is
not covered, and none of the Scala queries has a `FILTER` clause, so
`opt_filter` is only reached from this test. Could one entry here be `None`,
and could the Scala test add something like `avg(i) FILTER (WHERE d > 1)`,
which yields a null predicate wherever `d` is null? A sliced input such as
`values.slice(1, 3)` would also cover the offset handling.
##########
native/spark-expr/src/agg_funcs/avg.rs:
##########
@@ -334,12 +334,33 @@ where
])
}
+ /// Each row's own partial state, as `state` emits it: the value and a
count of one, or a
+ /// zero sum and count for a null or filtered-out row.
fn convert_to_state(
&self,
- _values: &[ArrayRef],
- _opt_filter: Option<&arrow::array::BooleanArray>,
+ values: &[ArrayRef],
+ opt_filter: Option<&arrow::array::BooleanArray>,
) -> Result<Vec<ArrayRef>> {
- not_impl_err!("Input batch conversion to state not implemented")
+ let values = values[0].as_primitive::<T>();
+ let mut sums = Vec::with_capacity(values.len());
+ let mut counts = Vec::with_capacity(values.len());
+ for idx in 0..values.len() {
Review Comment:
Once skipping starts, every remaining row of the task goes through this
loop, one `is_valid` and `push` per row. With no nulls and no filter, the
common case, the input buffer can be reused as the sums with no copy and the
counts filled with ones, as DataFusion's
`AvgGroupsAccumulator::convert_to_state` does. Otherwise one
`NullBuffer::union` of the filter (via `filter_to_nulls` in `collect.rs`) and
the input nulls gives the mask, and the sums and counts come from one zip over
it. Since Comet's `merge_batch` reads `.values()`, masked slots still need a
zero sum, which this keeps. I have not compiled it.
--
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]
