feat: support max_by and min_by aggregate expressions - #4817
Conversation
Add native support for Spark's max_by(x, y) aggregate, which returns the value of x associated with the maximum value of y. The expression is implemented as a native DataFusion aggregate (MaxMinBy) that compares the ordering column via Arrow's row format. Null orderings are ignored, the value paired with the maximum ordering is returned (and may itself be null), and an all-null-ordering group yields null. Both the value and ordering must be fixed-length types: a variable-length or nested type forces Spark's SortAggregate, which Comet does not run, so those cases fall back to Spark.
Add min_by alongside max_by, both served by the shared native MaxMinBy aggregate. Refactor the Scala serde into a shared CometMaxMinBy base. Add a vectorized GroupsAccumulator that keeps each group's best (value, ordering) pair as Arrow row-format bytes, so grouped max_by/min_by avoids the per-group ScalarValue work of the generic GroupsAccumulatorAdapter. Extend the aggregate microbenchmark and SQL file tests to cover both max_by and min_by.
andygrove
changed the title
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @andygrove!
| /// `max_by` (largest ordering wins), descending for `min_by` (smallest ordering wins). Nulls | ||
| /// sort first (smallest) so they are never selected as the extremum; null orderings are also | ||
| /// skipped explicitly. | ||
| fn extremum_sort_options(is_max: bool) -> SortOptions { |
There was a problem hiding this comment.
The Arrow row encoding does not reproduce Spark's ordering for -0.0 and 0.0, so max_by/min_by can pick a different row than Spark when the ordering column is a float containing signed zeros.
Spark orders doubles with SQLOrderingUtil.compareDoubles, which is if (x == y) 0 else java.lang.Double.compare(x, y) and documents "3. -0.0 == 0.0" (SQLOrderingUtil.scala:27-31, wired in via PhysicalDoubleType.ordering, PhysicalDataType.scala:183-184). So Spark treats -0.0 and 0.0 as a tie, and the tie is resolved by row order: the update chain keeps the existing value when predicate(old, new) holds, so for max_by an equal ordering leaves the earlier row's value in place.
Arrow's row format instead encodes floats by flipping bits off the sign (arrow-row/src/fixed.rs:149-157), which is a total order that places -0.0 strictly below 0.0. So rows.row(i) >= rows.row(b) at :196 and candidate >= self.best_ordering[...] at :340 see a strict inequality where Spark sees equality.
Concretely, for max_by(v, ord) over (1, 0.0), (2, -0.0): Spark ties, keeps the first, returns 1. Comet ranks 0.0 > -0.0, so the second row does not displace the first and it also returns 1. But reverse the input to (1, -0.0), (2, 0.0) and Spark still ties and returns 1, while Comet sees 0.0 > -0.0, takes the second, and returns 2. min_by diverges symmetrically.
Neither fixture covers this. max_by.sql has NaN (:122) and both infinities (:198) but no signed zero, and the Rust tests only cover NaN (max_min_by.rs:474). Please add a signed-zero case to both SQL fixtures and a Rust test, and normalize -0.0 to 0.0 on the ordering column before row conversion so ties fall through to the row-order tie-break.
Note this is the mirror image of the mode finding in #4782: there the fix is to stop collapsing -0.0, because mode keys on OpenHashSet.equals, which distinguishes them. Here the fix is to start collapsing, because max_by compares with SQLOrderingUtil, which does not. The two aggregates genuinely differ, so it is worth a comment at each site recording which Spark comparison path applies.
| } else { | ||
| // Compare the batch's extremum ordering against the running extremum using the | ||
| // same row encoding. Build a two-row array [running, candidate] and compare. | ||
| let pair = ScalarValue::iter_to_array(vec![ |
There was a problem hiding this comment.
The scalar accumulator allocates a RowConverter per call to update_from (:181) and then, for every batch that has a non-null ordering, builds a two-element array and converts it just to compare the batch extremum against the running one (:211-216).
The grouped path already avoids this by holding ordering_converter as a field (:260) and keeping the running best as OwnedRow bytes, so comparison is a byte compare with no allocation (:340). The scalar accumulator can do the same: hold the converter and store ordering as an OwnedRow alongside the ScalarValue needed by state(). As written, the global (non-grouped) path pays two RowConverter constructions and an extra array build per batch, which is the path the benchmark table shows at 24ms and therefore the one where fixed per-batch overhead is most visible.
| )])?; | ||
| let rows = converter.convert_columns(&[Arc::clone(ordering_arr)])?; | ||
|
|
||
| // Find the index of the extremum ordering in this batch (last one wins on a tie, |
There was a problem hiding this comment.
The comment says "last one wins on a tie, matching Spark's sequential row processing", which is backwards for the value that gets kept.
Spark's update keeps the existing value on a tie: If(predicate(extremumOrdering, orderingExpr), valueWithExtremumOrdering, valueExpr) where predicate for max_by is oldExpr > newExpr (MaxByAndMinBy.scala:69-77, :113-114). On a tie the predicate is false, so the expression evaluates to valueExpr, the new row. So last-wins is correct for max_by here, but the reasoning is not "sequential processing", it is the specific strictness of the predicate, and the same strictness is what makes the -0.0 case above diverge. Please state the actual rule so the next reader can check it against MaxByAndMinBy.scala rather than re-deriving it.
2 participants
Which issue does this PR close?
Closes #3841.
Rationale for this change
max_by(x, y)returns the value ofxassociated with the maximum value ofy, andmin_by(x, y)returns the value associated with the minimum. Both were previously unsupported, so any query using them fell back to Spark. There is nomax_by/min_byaggregate in DataFusion or thedatafusion-sparkcrate, so this adds a native implementation.What changes are included in this PR?
MaxMinBy(native/spark-expr/src/agg_funcs/max_min_by.rs) serving bothmax_byandmin_by(parameterized by anis_maxflag). It compares the ordering column using Arrow's row format, so any orderable fixed-length type is handled with a single Spark-consistent total order (NaN sorts largest, matching Spark). Null orderings are ignored, the value paired with the extremum ordering is returned (and may itself be null), and an all-null-ordering group yields null.GroupsAccumulatorfor grouped aggregation. Each group keeps its best(value, ordering)pair as Arrow row-format bytes, so selecting the extremum for a batch is one row conversion plus per-row byte comparisons. This avoids the per-groupScalarValuework of DataFusion's genericGroupsAccumulatorAdapter. The scalarAccumulatoris retained for the non-grouped path.MaxByandMinByprotobuf messages and wiring inplanner.rs.CometMaxBy/CometMinByserdes sharing aCometMaxMinBy[T]base, registered inQueryPlanSerde.getSupportLevelrestricts both the value and the ordering to fixed-length types: a variable-length or nested type (string, binary, struct) forces Spark'sSortAggregate, which Comet does not accelerate, so those cases fall back to Spark.max_by/min_bygroups inCometAggregateExpressionBenchmark.The 2-argument implementations are identical across Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1, so no version shim is needed. The 3-argument top-k forms
max_by(x, y, k)/min_by(x, y, k)exist only on Spark master and are out of scope.This implementation was scaffolded with the
implement-comet-expressionproject skill, which also ran theaudit-comet-expressionskill to drive the test-coverage pass.Benchmark
CometAggregateExpressionBenchmarkon an Apple M3 Ultra, 1M rows (best time in ms; lower is better).max_by/min_byrun fully natively and are at parity with Spark on this scan-dominated workload, with a small edge at high group cardinality where the vectorizedGroupsAccumulatormatters most.max_by(int, long)max_by(double, int)max_by(decimal, int)max_by(int, long)max_by(int, long)min_by(int, long)min_by(double, int)min_by(int, long)How are these changes tested?
max_min_by.rs(14 total) covering the scalar accumulator and the grouped accumulator: basic max/min selection, null-ordering handling, null value at the extremum, all-null orderings, NaN ordering, empty groups, multi-group aggregation, filter handling, and partial/final merge equivalence.spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sqlandmin_by.sqlcovering global and grouped aggregation, null handling, literal arguments, negative and boundary orderings, NaN/±Infinity, decimal/date/timestamp/boolean types, multiple aggregates in one query, and mixingmax_bywithmin_by. Each query runs through both Spark and Comet and asserts native execution.