This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 149a1dbb7d feat: support max_by and min_by aggregate expressions
(#4817)
149a1dbb7d is described below
commit 149a1dbb7dc53c89f901c1eb50793c6015b3a4bf
Author: Andy Grove <[email protected]>
AuthorDate: Thu Sep 10 21:49:05 2026 -0600
feat: support max_by and min_by aggregate expressions (#4817)
---
.../expression-audits/agg_funcs.md | 14 +
docs/source/user-guide/latest/expressions.md | 4 +-
native/core/src/execution/planner.rs | 20 +-
native/proto/src/proto/expr.proto | 16 +
native/spark-expr/src/agg_funcs/max_min_by.rs | 895 +++++++++++++++++++++
native/spark-expr/src/agg_funcs/mod.rs | 2 +
.../org/apache/comet/serde/QueryPlanSerde.scala | 2 +
.../scala/org/apache/comet/serde/aggregates.scala | 85 +-
.../sql-tests/expressions/aggregate/max_by.sql | 284 +++++++
.../sql-tests/expressions/aggregate/min_by.sql | 266 ++++++
.../CometAggregateExpressionBenchmark.scala | 24 +-
11 files changed, 1605 insertions(+), 7 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md
b/docs/source/contributor-guide/expression-audits/agg_funcs.md
index 164ca2a72c..b2d2efae6b 100644
--- a/docs/source/contributor-guide/expression-audits/agg_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md
@@ -70,6 +70,13 @@
- Comet implementation: the native side delegates to
`datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps
`DistinctArrayAggAccumulator` with `ignore_nulls = true` in a
`NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty
array. The `containsNull` mismatch against Spark's declared output type, and
its rationale, are identical to [collect_list](#collect-list).
- `CometCollectSet` reports `Incompatible` for float and double input when
`spark.comet.exec.strictFloatingPoint=true`, because the native distinct
comparison treats `NaN == NaN` and collapses repeated `NaN`s into a single
element while Spark keeps each one. The native path for floating-point input is
then opt-in via `spark.comet.expression.CollectSet.allowIncompatible=true`. All
other input types are `Compatible`.
+## max_by
+
+- Spark 3.4.3 (2026-07-03): `MaxBy` is a 2-argument `DeclarativeAggregate`
registered as `expression[MaxBy]("max_by")`. Buffer is
`(valueWithExtremumOrdering, extremumOrdering)`; 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. Comet implements a native
`max_by` aggregate. Only fixed-length value and ordering types are accelerated:
a variable-length or nested type (string, binary, stru [...]
+- Spark 3.5.8 (2026-07-03): aggregate logic identical to 3.4.3.
+- Spark 4.0.1 (2026-07-03): aggregate logic identical to 3.4.3; only the
`@ExpressionDescription` example and note text differ.
+- Spark 4.1.1 (2026-07-03): aggregate logic identical to 3.4.3. The 3-argument
top-k form `max_by(x, y, k)` arrived in Spark 4.2 (`MaxMinByK.scala`, absent on
`branch-4.1`), so it is absent from 3.4 through 4.1 and present on the 4.2
profile this repo builds. `MaxByBuilder.build` still returns a plain `MaxBy`
for the 2-argument call, and the 3-argument call becomes `MaxMinByK`, a
different class with no serde registration, so Comet handles only the
2-argument form and the top-k form falls back.
+
## median
- Spark 3.4.3 (audited 2026-06-24): `Median(child)` is a
`RuntimeReplaceableAggregate` with `replacement = Percentile(child,
Literal(0.5))`. Catalyst rewrites `median(x)` to `percentile(x, 0.5)` before
Comet sees the plan, so it is served by `CometPercentile`.
@@ -77,6 +84,13 @@
- Spark 4.0.1 (audited 2026-06-24): `replacement` becomes `lazy val`;
semantics unchanged.
- Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1.
+## min_by
+
+- Spark 3.4.3 (2026-07-03): `MinBy` shares the abstract `MaxMinBy`
`DeclarativeAggregate` with `MaxBy`, differing only in the comparison direction
(`least` / `<` instead of `greatest` / `>`). Registered as
`expression[MinBy]("min_by")`. Null orderings are ignored, the value paired
with the minimum ordering is returned (and may itself be null), and an
all-null-ordering group yields null. Comet serves it through the same native
`MaxMinBy` aggregate as `max_by`, with the same fixed-length v [...]
+- Spark 3.5.8 (2026-07-03): aggregate logic identical to 3.4.3.
+- Spark 4.0.1 (2026-07-03): aggregate logic identical to 3.4.3; only the
`@ExpressionDescription` example and note text differ.
+- Spark 4.1.1 (2026-07-03): aggregate logic identical to 3.4.3. The 3-argument
top-k form `min_by(x, y, k)` arrived in Spark 4.2 alongside `max_by(x, y, k)`;
`MinByBuilder.build` returns a plain `MinBy` for the 2-argument call, so Comet
handles only the 2-argument form and the top-k form falls back.
+
## percentile
- Spark 3.4.3 (audited 2026-06-24): `Percentile(child, percentageExpression,
frequencyExpression, ..., reverse)` over `PercentileBase`. Exact percentile
using `index = p * (n - 1)` linear interpolation, NULL inputs skipped,
empty/all-null group returns NULL. `CometPercentile` maps the
single-literal-percentage, default-frequency, numeric-input, ascending form to
DataFusion's `percentile_cont` (same interpolation). Array-of-percentages, a
non-default frequency argument, descending order, [...]
diff --git a/docs/source/user-guide/latest/expressions.md
b/docs/source/user-guide/latest/expressions.md
index 65e73ec3bc..208a5f3f12 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -109,11 +109,11 @@ The tables below list every Spark built-in expression
with its current status.
| `last_value` | ✅ | Native | |
| `listagg` | 🔜 | — | String aggregation |
| `max` | ✅ | Native | |
-| `max_by` | 🔜 | — |
[#3841](https://github.com/apache/datafusion-comet/issues/3841) |
+| `max_by` | ✅ | Native | Value and ordering must be fixed-length types |
| `mean` | ✅ | Native | |
| `median` | ✅ | — | Rewrites to `percentile(col, 0.5)` and runs natively for
supported percentile inputs |
| `min` | ✅ | Native | |
-| `min_by` | 🔜 | — |
[#3841](https://github.com/apache/datafusion-comet/issues/3841) |
+| `min_by` | ✅ | Native | Value and ordering must be fixed-length types |
| `mode` | ✅ | Native | `mode(col)` only; Spark breaks ties
non-deterministically, so Comet returns the smallest tied value and falls back
by default, opt-in via allowIncompatible
([#3970](https://github.com/apache/datafusion-comet/issues/3970)) |
| `percentile` | ✅ | Native | Single literal percentage on numeric input runs
natively; array of percentages and a frequency argument fall back to Spark |
| `percentile_cont` | ✅ | — | Spark 4.0+ `WITHIN GROUP (ORDER BY ...)`;
ascending only runs natively, `DESC` falls back to Spark |
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index 1cd84d314f..eeaee34478 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -149,9 +149,9 @@ use datafusion_comet_proto::{
use datafusion_comet_spark_expr::{
jvm_udf::JvmScalarUdfExpr, ApproxPercentile, ArrayInsert, Avg, AvgDecimal,
Cast, CheckOverflow,
Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow,
GetArrayStructFields,
- GetStructField, HllPlusPlus, IfExpr, ListExtract, Mode,
NormalizeNaNAndZero, Regr, RegrType,
- SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn, Variance,
WideDecimalBinaryExpr,
- WideDecimalOp,
+ GetStructField, HllPlusPlus, IfExpr, ListExtract, MaxMinBy, Mode,
NormalizeNaNAndZero, Regr,
+ RegrType, SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn,
Variance,
+ WideDecimalBinaryExpr, WideDecimalOp,
};
use itertools::Itertools;
use jni::objects::{Global, JObject};
@@ -3188,6 +3188,20 @@ impl PhysicalPlanner {
let func =
AggregateUDF::new_from_impl(HllPlusPlus::new(expr.precision));
Self::create_aggr_func_expr("approx_count_distinct", schema,
vec![child], func)
}
+ AggExprStruct::MaxBy(expr) => {
+ let value = self.create_expr(expr.value.as_ref().unwrap(),
Arc::clone(&schema))?;
+ let ordering =
+ self.create_expr(expr.ordering.as_ref().unwrap(),
Arc::clone(&schema))?;
+ let func = AggregateUDF::new_from_impl(MaxMinBy::new_max_by());
+ Self::create_aggr_func_expr("max_by", schema, vec![value,
ordering], func)
+ }
+ AggExprStruct::MinBy(expr) => {
+ let value = self.create_expr(expr.value.as_ref().unwrap(),
Arc::clone(&schema))?;
+ let ordering =
+ self.create_expr(expr.ordering.as_ref().unwrap(),
Arc::clone(&schema))?;
+ let func = AggregateUDF::new_from_impl(MaxMinBy::new_min_by());
+ Self::create_aggr_func_expr("min_by", schema, vec![value,
ordering], func)
+ }
AggExprStruct::Mode(expr) => {
let child = self.create_expr(expr.child.as_ref().unwrap(),
Arc::clone(&schema))?;
let datatype =
to_arrow_datatype(expr.datatype.as_ref().unwrap());
diff --git a/native/proto/src/proto/expr.proto
b/native/proto/src/proto/expr.proto
index 06fdf19925..34d502daac 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -164,6 +164,8 @@ message AggExpr {
CollectList collectList = 21;
Regr regr = 22;
Mode mode = 23;
+ MaxBy maxBy = 24;
+ MinBy minBy = 25;
}
// Optional filter expression for SQL FILTER (WHERE ...) clause.
@@ -353,6 +355,20 @@ message HllPlusPlus {
int32 precision = 2;
}
+message MaxBy {
+ // The value returned by the aggregate (associated with the maximum
ordering).
+ Expr value = 1;
+ // The ordering expression whose maximum selects the returned value.
+ Expr ordering = 2;
+}
+
+message MinBy {
+ // The value returned by the aggregate (associated with the minimum
ordering).
+ Expr value = 1;
+ // The ordering expression whose minimum selects the returned value.
+ Expr ordering = 2;
+}
+
message Mode {
Expr child = 1;
DataType datatype = 2;
diff --git a/native/spark-expr/src/agg_funcs/max_min_by.rs
b/native/spark-expr/src/agg_funcs/max_min_by.rs
new file mode 100644
index 0000000000..147f1eb92b
--- /dev/null
+++ b/native/spark-expr/src/agg_funcs/max_min_by.rs
@@ -0,0 +1,895 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{new_null_array, Array, ArrayRef, AsArray, BooleanArray};
+use arrow::compute::SortOptions;
+use arrow::datatypes::{DataType, Field, FieldRef, Float32Type, Float64Type};
+use arrow::row::{OwnedRow, RowConverter, SortField};
+use datafusion::common::{not_impl_err, Result, ScalarValue};
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::{
+ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, Signature,
Volatility,
+};
+use datafusion::physical_expr::expressions::format_state_name;
+use std::mem::size_of_val;
+use std::sync::Arc;
+
+/// Spark-compatible `max_by(value, ordering)` / `min_by(value, ordering)`
aggregate.
+///
+/// Returns the `value` associated with the maximum (`max_by`) or minimum
(`min_by`)
+/// non-null `ordering`. Rows with a null `ordering` are ignored. The returned
value
+/// may itself be null when it is the value paired with the extremum ordering.
If every
+/// `ordering` in the group is null, the result is null.
+///
+/// Spark's `MaxBy`/`MinBy` are `DeclarativeAggregate`s that keep a `(value,
ordering)`
+/// buffer and, on a tie in the ordering, the later row wins. Because ties
across
+/// partitions are processed in an unspecified order, Spark documents the
function as
+/// non-deterministic when several rows share the extremum ordering.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MaxMinBy {
+ name: String,
+ signature: Signature,
+ /// `true` for `max_by`, `false` for `min_by`.
+ is_max: bool,
+}
+
+impl std::hash::Hash for MaxMinBy {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.name.hash(state);
+ self.signature.hash(state);
+ self.is_max.hash(state);
+ }
+}
+
+impl MaxMinBy {
+ /// Create a `max_by` aggregate.
+ pub fn new_max_by() -> Self {
+ Self {
+ name: "max_by".to_string(),
+ signature: Signature::any(2, Volatility::Immutable),
+ is_max: true,
+ }
+ }
+
+ /// Create a `min_by` aggregate.
+ pub fn new_min_by() -> Self {
+ Self {
+ name: "min_by".to_string(),
+ signature: Signature::any(2, Volatility::Immutable),
+ is_max: false,
+ }
+ }
+}
+
+impl AggregateUDFImpl for MaxMinBy {
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ // The result has the same type as the `value` argument.
+ Ok(arg_types[0].clone())
+ }
+
+ fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn
Accumulator>> {
+ let value_type = acc_args.exprs[0].data_type(acc_args.schema)?;
+ let ordering_type = acc_args.exprs[1].data_type(acc_args.schema)?;
+ Ok(Box::new(MaxMinByAccumulator::try_new(
+ value_type,
+ ordering_type,
+ self.is_max,
+ )?))
+ }
+
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+ let value_type = args.input_fields[0].data_type().clone();
+ let ordering_type = args.input_fields[1].data_type().clone();
+ Ok(vec![
+ Arc::new(Field::new(
+ format_state_name(&self.name, "value"),
+ value_type,
+ true,
+ )),
+ Arc::new(Field::new(
+ format_state_name(&self.name, "ordering"),
+ ordering_type,
+ true,
+ )),
+ ])
+ }
+
+ fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+ true
+ }
+
+ fn create_groups_accumulator(
+ &self,
+ args: AccumulatorArgs,
+ ) -> Result<Box<dyn GroupsAccumulator>> {
+ let value_type = args.exprs[0].data_type(args.schema)?;
+ let ordering_type = args.exprs[1].data_type(args.schema)?;
+ Ok(Box::new(MaxMinByGroupsAccumulator::try_new(
+ value_type,
+ ordering_type,
+ self.is_max,
+ )?))
+ }
+}
+
+/// Sort options that make the wanted extremum encode to the largest row
bytes: ascending for
+/// `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 {
+ SortOptions {
+ descending: !is_max,
+ nulls_first: true,
+ }
+}
+
+/// Canonicalize a floating-point ordering column so that Arrow's row-format
byte order reproduces
+/// Spark's comparison for this aggregate.
+///
+/// Spark compares the ordering with
`SQLOrderingUtil.compareDoubles`/`compareFloats`, wired in via
+/// `PhysicalDoubleType.ordering`/`PhysicalFloatType.ordering`. That is
+/// `if (x == y) 0 else java.lang.Double.compare(x, y)`, which has two
consequences Arrow's row
+/// format does not share:
+///
+/// * `-0.0` and `0.0` tie, because the `x == y` short-circuit is IEEE
equality. Arrow encodes
+/// floats by flipping the bits off the sign, a total order placing `-0.0`
strictly below `0.0`.
+/// * every `NaN` is one value and sorts above `+Infinity`, because
`Double.compare` goes through
+/// `doubleToLongBits`. Arrow uses the raw bits, so a sign-bit-set `NaN`
would sort below
+/// `-Infinity` instead.
+///
+/// Folding `-0.0` into `0.0` and every `NaN` into the canonical `NaN` makes
the row bytes agree
+/// with `compareDoubles` on both counts. This is verified identical on Spark
3.4 through master.
+///
+/// Note this is the opposite of what `mode` needs: `mode` keys a hash map via
+/// `OpenHashSet`'s `equals` (`java.lang.Double.equals`), which distinguishes
`-0.0` from `0.0`, so
+/// it must *not* fold them. Same two input values, different Spark comparison
path, opposite
+/// correct behaviour.
+fn canonicalize_float_ordering(array: &ArrayRef) -> ArrayRef {
+ match array.data_type() {
+ DataType::Float32 =>
Arc::new(array.as_primitive::<Float32Type>().unary::<_, Float32Type>(
+ |v| {
+ if v.is_nan() {
+ f32::NAN
+ } else if v == 0.0 {
+ // `-0.0 == 0.0` in IEEE 754, so this catches negative
zero only.
+ 0.0
+ } else {
+ v
+ }
+ },
+ )),
+ DataType::Float64 =>
Arc::new(array.as_primitive::<Float64Type>().unary::<_, Float64Type>(
+ |v| {
+ if v.is_nan() {
+ f64::NAN
+ } else if v == 0.0 {
+ 0.0
+ } else {
+ v
+ }
+ },
+ )),
+ _ => Arc::clone(array),
+ }
+}
+
+/// Accumulator that tracks the running `(value, ordering)` pair for the
extremum ordering.
+#[derive(Debug)]
+struct MaxMinByAccumulator {
+ /// Converts the ordering column into Arrow's byte-comparable row format.
Held as a field so a
+ /// batch update does not pay a `RowConverter` construction, matching the
grouped accumulator.
+ ordering_converter: RowConverter,
+ /// Ordering type, needed to produce a null ordering scalar in `state`
before any row is seen.
+ ordering_type: DataType,
+ /// The value paired with the current extremum ordering. May be null.
+ value: ScalarValue,
+ /// Row bytes of the current extremum ordering. `None` until a non-null
ordering is seen, so
+ /// comparing against the running extremum is a byte compare with no
allocation.
+ best_ordering: Option<OwnedRow>,
+}
+
+impl MaxMinByAccumulator {
+ fn try_new(value_type: DataType, ordering_type: DataType, is_max: bool) ->
Result<Self> {
+ let ordering_converter =
RowConverter::new(vec![SortField::new_with_options(
+ ordering_type.clone(),
+ extremum_sort_options(is_max),
+ )])?;
+ Ok(Self {
+ ordering_converter,
+ ordering_type,
+ value: ScalarValue::try_from(&value_type)?,
+ best_ordering: None,
+ })
+ }
+
+ /// Apply a batch of `(value, ordering)` columns, keeping the value paired
with the
+ /// extremum ordering. Rows with a null ordering are ignored.
+ fn update_from(&mut self, value_arr: &ArrayRef, ordering_arr: &ArrayRef)
-> Result<()> {
+ if ordering_arr.is_empty() {
+ return Ok(());
+ }
+
+ let ordering_arr = canonicalize_float_ordering(ordering_arr);
+ let rows = self
+ .ordering_converter
+ .convert_columns(&[Arc::clone(&ordering_arr)])?;
+
+ // Find the index of the extremum ordering in this batch, ignoring
null orderings. `>=`
+ // makes the later row win a tie, which is what Spark's update does:
it evaluates
+ // `If(predicate(extremumOrdering, orderingExpr),
valueWithExtremumOrdering, valueExpr)`
+ // where `predicate` is the strict `oldExpr > newExpr` for `max_by`
(`<` for `min_by`), so
+ // an equal ordering makes the predicate false and the *new* row's
value is kept
+ // (`MaxByAndMinBy.scala`). The same strictness is why the signed-zero
canonicalization
+ // above matters: without it Arrow sees a strict inequality where
Spark sees a tie.
+ let mut best: Option<usize> = None;
+ for i in 0..ordering_arr.len() {
+ if ordering_arr.is_null(i) {
+ continue;
+ }
+ best = match best {
+ None => Some(i),
+ Some(b) if rows.row(i) >= rows.row(b) => Some(i),
+ Some(b) => Some(b),
+ };
+ }
+
+ let Some(b) = best else {
+ return Ok(());
+ };
+
+ let candidate = rows.row(b);
+ let take = match &self.best_ordering {
+ None => true,
+ Some(running) => candidate >= running.row(),
+ };
+
+ if take {
+ self.value = ScalarValue::try_from_array(value_arr, b)?;
+ self.best_ordering = Some(candidate.owned());
+ }
+
+ Ok(())
+ }
+
+ /// The running extremum ordering as a scalar, for the aggregation state.
+ fn ordering_scalar(&self) -> Result<ScalarValue> {
+ match &self.best_ordering {
+ Some(row) => {
+ let arrays = self
+ .ordering_converter
+ .convert_rows(std::iter::once(row.row()))?;
+ ScalarValue::try_from_array(&arrays[0], 0)
+ }
+ None => ScalarValue::try_from(&self.ordering_type),
+ }
+ }
+}
+
+impl Accumulator for MaxMinByAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ self.update_from(&values[0], &values[1])
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+ // State columns mirror the input columns: [value, ordering].
+ self.update_from(&states[0], &states[1])
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ Ok(vec![self.value.clone(), self.ordering_scalar()?])
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ Ok(self.value.clone())
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ + self.value.size()
+ + self
+ .best_ordering
+ .as_ref()
+ .map_or(0, |r| r.row().as_ref().len())
+ }
+}
+
+/// Vectorized grouped accumulator for `max_by` / `min_by`.
+///
+/// Each group keeps the best `(value, ordering)` pair as Arrow row-format
bytes. The ordering
+/// rows are byte-comparable, so selecting the extremum for a batch is a
single row conversion
+/// plus per-row byte comparisons, avoiding the per-group `ScalarValue` work
of the generic
+/// `GroupsAccumulatorAdapter`.
+struct MaxMinByGroupsAccumulator {
+ /// Converts and compares the ordering column. Its sort options encode the
wanted extremum
+ /// as the largest row bytes (see `extremum_sort_options`).
+ ordering_converter: RowConverter,
+ /// Converts the value column to and from row bytes. Sort options are
irrelevant here since
+ /// values are only stored, never compared.
+ value_converter: RowConverter,
+ /// Row bytes for a null value, used for groups that have not been updated
(or whose winning
+ /// value is null).
+ null_value_row: OwnedRow,
+ /// Row bytes for a null ordering, used for groups that have seen no
non-null ordering.
+ null_ordering_row: OwnedRow,
+ /// Per-group winning value (row bytes).
+ best_value: Vec<OwnedRow>,
+ /// Per-group winning ordering (row bytes).
+ best_ordering: Vec<OwnedRow>,
+ /// Per-group flag: has a non-null ordering been seen yet?
+ has_ordering: Vec<bool>,
+}
+
+impl MaxMinByGroupsAccumulator {
+ fn try_new(value_type: DataType, ordering_type: DataType, is_max: bool) ->
Result<Self> {
+ let ordering_converter =
RowConverter::new(vec![SortField::new_with_options(
+ ordering_type.clone(),
+ extremum_sort_options(is_max),
+ )])?;
+ let value_converter =
RowConverter::new(vec![SortField::new(value_type.clone())])?;
+ let null_ordering_row = ordering_converter
+ .convert_columns(&[new_null_array(&ordering_type, 1)])?
+ .row(0)
+ .owned();
+ let null_value_row = value_converter
+ .convert_columns(&[new_null_array(&value_type, 1)])?
+ .row(0)
+ .owned();
+ Ok(Self {
+ ordering_converter,
+ value_converter,
+ null_value_row,
+ null_ordering_row,
+ best_value: Vec::new(),
+ best_ordering: Vec::new(),
+ has_ordering: Vec::new(),
+ })
+ }
+
+ fn resize(&mut self, total_num_groups: usize) {
+ self.best_value
+ .resize(total_num_groups, self.null_value_row.clone());
+ self.best_ordering
+ .resize(total_num_groups, self.null_ordering_row.clone());
+ self.has_ordering.resize(total_num_groups, false);
+ }
+
+ /// Shared update/merge logic: `values[0]` is the value column,
`values[1]` the ordering
+ /// column. Rows with a null ordering are ignored; on a tie the later row
wins, matching the
+ /// strict predicate in Spark's update (see
`MaxMinByAccumulator::update_from`).
+ fn update_groups(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ self.resize(total_num_groups);
+ let value_rows = self
+ .value_converter
+ .convert_columns(&[Arc::clone(&values[0])])?;
+ let ordering_arr = canonicalize_float_ordering(&values[1]);
+ let ordering_rows = self
+ .ordering_converter
+ .convert_columns(&[Arc::clone(&ordering_arr)])?;
+
+ for (idx, &group_index) in group_indices.iter().enumerate() {
+ if let Some(filter) = opt_filter {
+ if !filter.is_valid(idx) || !filter.value(idx) {
+ continue;
+ }
+ }
+ if ordering_arr.is_null(idx) {
+ continue;
+ }
+ let candidate = ordering_rows.row(idx);
+ let take = !self.has_ordering[group_index]
+ || candidate >= self.best_ordering[group_index].row();
+ if take {
+ self.best_ordering[group_index] = candidate.owned();
+ self.best_value[group_index] = value_rows.row(idx).owned();
+ self.has_ordering[group_index] = true;
+ }
+ }
+ Ok(())
+ }
+}
+
+impl GroupsAccumulator for MaxMinByGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ self.update_groups(values, group_indices, opt_filter, total_num_groups)
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> Result<()> {
+ // State columns mirror the input columns: [value, ordering].
+ self.update_groups(values, group_indices, None, total_num_groups)
+ }
+
+ fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let value_rows = emit_to.take_needed(&mut self.best_value);
+ let _ = emit_to.take_needed(&mut self.best_ordering);
+ let _ = emit_to.take_needed(&mut self.has_ordering);
+ let arrays = self
+ .value_converter
+ .convert_rows(value_rows.iter().map(|r| r.row()))?;
+ Ok(Arc::clone(&arrays[0]))
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ let value_rows = emit_to.take_needed(&mut self.best_value);
+ let ordering_rows = emit_to.take_needed(&mut self.best_ordering);
+ let _ = emit_to.take_needed(&mut self.has_ordering);
+ let value_arrays = self
+ .value_converter
+ .convert_rows(value_rows.iter().map(|r| r.row()))?;
+ let ordering_arrays = self
+ .ordering_converter
+ .convert_rows(ordering_rows.iter().map(|r| r.row()))?;
+ Ok(vec![
+ Arc::clone(&value_arrays[0]),
+ Arc::clone(&ordering_arrays[0]),
+ ])
+ }
+
+ fn convert_to_state(
+ &self,
+ _values: &[ArrayRef],
+ _opt_filter: Option<&BooleanArray>,
+ ) -> Result<Vec<ArrayRef>> {
+ not_impl_err!("Input batch conversion to state not implemented")
+ }
+
+ fn size(&self) -> usize {
+ // `size_of::<OwnedRow>()` only covers the inline `Box<[u8]>`
pointer/len, not the row
+ // bytes behind it, so the per-group payloads have to be summed
separately. Leaving them
+ // out under-reports to the memory pool that drives spill decisions.
+ let row_bytes = |rows: &Vec<OwnedRow>| -> usize {
+ rows.iter().map(|r| r.row().as_ref().len()).sum::<usize>()
+ };
+ size_of_val(self)
+ + (self.best_value.capacity() + self.best_ordering.capacity())
+ * std::mem::size_of::<OwnedRow>()
+ + row_bytes(&self.best_value)
+ + row_bytes(&self.best_ordering)
+ + self.has_ordering.capacity()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{AsArray, Float32Array, Float64Array, Int32Array,
Int64Array, StringArray};
+
+ fn max_by_acc(value_type: DataType, ordering_type: DataType) ->
MaxMinByAccumulator {
+ MaxMinByAccumulator::try_new(value_type, ordering_type, true).unwrap()
+ }
+
+ fn min_by_acc(value_type: DataType, ordering_type: DataType) ->
MaxMinByAccumulator {
+ MaxMinByAccumulator::try_new(value_type, ordering_type, false).unwrap()
+ }
+
+ #[test]
+ fn max_by_basic() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c"]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![10, 50, 20]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn min_by_basic() {
+ let mut acc = min_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c"]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![10, 50, 20]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("a"));
+ }
+
+ #[test]
+ fn null_ordering_is_ignored() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"),
Some("b"), Some("c")]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![Some(10),
None, Some(5)]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ // The row with ordering=None (value "b") is ignored; max ordering is
10 -> "a".
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("a"));
+ }
+
+ #[test]
+ fn all_null_ordering_yields_null() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"),
Some("b")]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::Utf8(None));
+ }
+
+ #[test]
+ fn null_value_at_extremum_is_returned() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"),
None]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![Some(10),
Some(50)]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ // Max ordering 50 pairs with a null value.
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::Utf8(None));
+ }
+
+ #[test]
+ fn empty_group_yields_null() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::Utf8(None));
+ }
+
+ #[test]
+ fn max_by_nan_is_largest() {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float64);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c"]));
+ let ordering: ArrayRef = Arc::new(Float64Array::from(vec![1.0,
f64::NAN, 2.0]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ // Spark treats NaN as the largest value, matching arrow's row
ordering.
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn max_by_sign_bit_nan_is_still_largest() {
+ // `Double.compare` goes through `doubleToLongBits`, so a sign-bit-set
NaN is the same
+ // value as a positive NaN and still sorts above every finite
ordering. Arrow's raw-bit
+ // encoding would otherwise place it below -Infinity.
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float64);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c"]));
+ let ordering: ArrayRef = Arc::new(Float64Array::from(vec![1.0,
-f64::NAN, 2.0]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn min_by_sign_bit_nan_is_not_smallest() {
+ // The mirror of the case above: because every NaN is the *largest*
value, a sign-bit-set
+ // NaN loses to any finite ordering in `min_by` and must not be
selected. Without the
+ // canonicalization arrow's raw-bit encoding places it below
-Infinity, which would make it
+ // the minimum and return "a".
+ let mut acc = min_by_acc(DataType::Utf8, DataType::Float64);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let ordering: ArrayRef = Arc::new(Float64Array::from(vec![-f64::NAN,
2.0]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn max_by_ties_signed_zeros_and_keeps_last_row() {
+ // Spark compares the ordering with `SQLOrderingUtil.compareDoubles`,
which ties
+ // `-0.0 == 0.0`, and its update keeps the *new* row on a tie. So both
orders below must
+ // return the later row. Without the canonicalization, arrow's row
format ranks
+ // `0.0 > -0.0` strictly and the second case would return "a".
+ for (ordering, expected) in [
+ (vec![0.0_f64, -0.0_f64], "b"),
+ (vec![-0.0_f64, 0.0_f64], "b"),
+ ] {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float64);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let ordering_arr: ArrayRef =
Arc::new(Float64Array::from(ordering.clone()));
+ acc.update_batch(&[values, ordering_arr]).unwrap();
+ assert_eq!(
+ acc.evaluate().unwrap(),
+ ScalarValue::from(expected),
+ "max_by over ordering {ordering:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn min_by_ties_signed_zeros_and_keeps_last_row() {
+ // `min_by`'s predicate is the strict `oldExpr < newExpr`, so a tie
likewise keeps the new
+ // row. The signed zeros must not break the tie in either direction.
+ for (ordering, expected) in [
+ (vec![0.0_f64, -0.0_f64], "b"),
+ (vec![-0.0_f64, 0.0_f64], "b"),
+ ] {
+ let mut acc = min_by_acc(DataType::Utf8, DataType::Float64);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let ordering_arr: ArrayRef =
Arc::new(Float64Array::from(ordering.clone()));
+ acc.update_batch(&[values, ordering_arr]).unwrap();
+ assert_eq!(
+ acc.evaluate().unwrap(),
+ ScalarValue::from(expected),
+ "min_by over ordering {ordering:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn signed_zero_tie_holds_across_batches() {
+ // The running extremum is compared against the next batch's extremum
with the same
+ // encoding, so the tie must survive the batch boundary too: a second
batch whose ordering
+ // is the other zero still displaces the first batch's value.
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float64);
+ acc.update_batch(&[
+ Arc::new(StringArray::from(vec!["a"])) as ArrayRef,
+ Arc::new(Float64Array::from(vec![0.0_f64])) as ArrayRef,
+ ])
+ .unwrap();
+ acc.update_batch(&[
+ Arc::new(StringArray::from(vec!["b"])) as ArrayRef,
+ Arc::new(Float64Array::from(vec![-0.0_f64])) as ArrayRef,
+ ])
+ .unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn float32_signed_zero_tie() {
+ // `compareFloats` has the same shape as `compareDoubles`, so Float32
orderings need the
+ // same treatment. Both orders are checked: the `(-0.0, 0.0)` order
alone would pass even
+ // without the canonicalization, because arrow's ranking happens to
pick the later row too.
+ for (ordering, expected) in [
+ (vec![0.0_f32, -0.0_f32], "b"),
+ (vec![-0.0_f32, 0.0_f32], "b"),
+ ] {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let ordering_arr: ArrayRef =
Arc::new(Float32Array::from(ordering.clone()));
+ acc.update_batch(&[values, ordering_arr]).unwrap();
+ assert_eq!(
+ acc.evaluate().unwrap(),
+ ScalarValue::from(expected),
+ "max_by over f32 ordering {ordering:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn float32_max_by_sign_bit_nan_is_largest() {
+ // `compareFloats` routes through `floatToIntBits` just as
`compareDoubles` routes through
+ // `doubleToLongBits`, so a sign-bit-set NaN is the same value as a
positive one and is the
+ // largest ordering for Float32 too. Without the canonicalization
arrow ranks `-NaN` below
+ // -Infinity, so this would return "c".
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Float32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c"]));
+ let ordering: ArrayRef = Arc::new(Float32Array::from(vec![1.0,
-f32::NAN, 2.0]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn float32_min_by_sign_bit_nan_is_not_smallest() {
+ // Kept separate from the `max_by` case above rather than asserted
alongside it: a failure
+ // of the first assertion would stop the test before the second ever
ran, which is how a
+ // half-vacuous test hides.
+ let mut acc = min_by_acc(DataType::Utf8, DataType::Float32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
+ let ordering: ArrayRef = Arc::new(Float32Array::from(vec![-f32::NAN,
2.0]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::from("b"));
+ }
+
+ #[test]
+ fn groups_signed_zero_tie() {
+ // Same invariant in the grouped accumulator: group 0 sees (0.0, -0.0)
and group 1 sees
+ // (-0.0, 0.0); both must keep the later row.
+ let mut acc =
+ MaxMinByGroupsAccumulator::try_new(DataType::Utf8,
DataType::Float64, true).unwrap();
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c",
"d"]));
+ let ordering: ArrayRef = Arc::new(Float64Array::from(vec![
+ 0.0_f64, -0.0_f64, -0.0_f64, 0.0_f64,
+ ]));
+ acc.update_batch(&[values, ordering], &[0, 0, 1, 1], None, 2)
+ .unwrap();
+ let result = acc.evaluate(EmitTo::All).unwrap();
+ let result = result.as_string::<i32>();
+ assert_eq!(result.value(0), "b");
+ assert_eq!(result.value(1), "d");
+ }
+
+ #[test]
+ fn merge_matches_single_shot() {
+ let single = {
+ let mut acc = max_by_acc(DataType::Utf8, DataType::Int32);
+ let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b",
"c", "d", "e", "f"]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![1, 6, 3,
2, 5, 4]));
+ acc.update_batch(&[values, ordering]).unwrap();
+ acc.evaluate().unwrap()
+ };
+
+ let mut left = max_by_acc(DataType::Utf8, DataType::Int32);
+ left.update_batch(&[
+ Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
+ Arc::new(Int32Array::from(vec![1, 6, 3])) as ArrayRef,
+ ])
+ .unwrap();
+ let mut right = max_by_acc(DataType::Utf8, DataType::Int32);
+ right
+ .update_batch(&[
+ Arc::new(StringArray::from(vec!["d", "e", "f"])) as ArrayRef,
+ Arc::new(Int32Array::from(vec![2, 5, 4])) as ArrayRef,
+ ])
+ .unwrap();
+
+ let mut merged = max_by_acc(DataType::Utf8, DataType::Int32);
+ for acc in [&mut left, &mut right] {
+ let state = acc.state().unwrap();
+ let value_arr =
ScalarValue::iter_to_array(vec![state[0].clone()]).unwrap();
+ let ordering_arr =
ScalarValue::iter_to_array(vec![state[1].clone()]).unwrap();
+ merged.merge_batch(&[value_arr, ordering_arr]).unwrap();
+ }
+ assert_eq!(merged.evaluate().unwrap(), single);
+ }
+
+ // ---- GroupsAccumulator tests ----
+
+ fn max_by_groups(value_type: DataType, ordering_type: DataType) ->
MaxMinByGroupsAccumulator {
+ MaxMinByGroupsAccumulator::try_new(value_type, ordering_type,
true).unwrap()
+ }
+
+ fn min_by_groups(value_type: DataType, ordering_type: DataType) ->
MaxMinByGroupsAccumulator {
+ MaxMinByGroupsAccumulator::try_new(value_type, ordering_type,
false).unwrap()
+ }
+
+ fn eval_int(acc: &mut MaxMinByGroupsAccumulator) -> Vec<Option<i32>> {
+ acc.evaluate(EmitTo::All)
+ .unwrap()
+ .as_primitive::<arrow::datatypes::Int32Type>()
+ .iter()
+ .collect()
+ }
+
+ #[test]
+ fn groups_max_by_multi_group() {
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int32);
+ // group 0: values 10,30 orderings 1,2 -> 30 ; group 1: values 20,40
orderings 5,4 -> 20
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30,
40]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![1, 5, 2, 4]));
+ acc.update_batch(&[values, ordering], &[0, 1, 0, 1], None, 2)
+ .unwrap();
+ assert_eq!(eval_int(&mut acc), vec![Some(30), Some(20)]);
+ }
+
+ #[test]
+ fn groups_min_by_multi_group() {
+ let mut acc = min_by_groups(DataType::Int32, DataType::Int32);
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30,
40]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![1, 5, 2, 4]));
+ acc.update_batch(&[values, ordering], &[0, 1, 0, 1], None, 2)
+ .unwrap();
+ // group 0: min ordering 1 -> 10 ; group 1: min ordering 4 -> 40
+ assert_eq!(eval_int(&mut acc), vec![Some(10), Some(40)]);
+ }
+
+ #[test]
+ fn groups_null_ordering_and_empty_group() {
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int32);
+ // group 0: (10,null),(30,5) -> 30 ; group 1: no rows -> null ; group
2: (40,null) -> null
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(10),
Some(30), Some(40)]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(5),
None]));
+ acc.update_batch(&[values, ordering], &[0, 0, 2], None, 3)
+ .unwrap();
+ assert_eq!(eval_int(&mut acc), vec![Some(30), None, None]);
+ }
+
+ #[test]
+ fn groups_null_value_at_extremum() {
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int32);
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(10),
None]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![Some(1),
Some(9)]));
+ acc.update_batch(&[values, ordering], &[0, 0], None, 1)
+ .unwrap();
+ assert_eq!(eval_int(&mut acc), vec![None]);
+ }
+
+ #[test]
+ fn groups_merge_matches_single_shot() {
+ let single = {
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int32);
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4,
5, 6]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![1, 6, 3,
2, 5, 4]));
+ acc.update_batch(&[values, ordering], &[0, 0, 0, 0, 0, 0], None, 1)
+ .unwrap();
+ eval_int(&mut acc)
+ };
+
+ let mut left = max_by_groups(DataType::Int32, DataType::Int32);
+ left.update_batch(
+ &[
+ Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
+ Arc::new(Int32Array::from(vec![1, 6, 3])) as ArrayRef,
+ ],
+ &[0, 0, 0],
+ None,
+ 1,
+ )
+ .unwrap();
+ let mut right = max_by_groups(DataType::Int32, DataType::Int32);
+ right
+ .update_batch(
+ &[
+ Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef,
+ Arc::new(Int32Array::from(vec![2, 5, 4])) as ArrayRef,
+ ],
+ &[0, 0, 0],
+ None,
+ 1,
+ )
+ .unwrap();
+
+ let mut merged = max_by_groups(DataType::Int32, DataType::Int32);
+ for acc in [&mut left, &mut right] {
+ let state = acc.state(EmitTo::All).unwrap();
+ merged.merge_batch(&state, &[0], 1).unwrap();
+ }
+ assert_eq!(eval_int(&mut merged), single);
+ }
+
+ #[test]
+ fn groups_filter_is_respected() {
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int32);
+ let values: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30]));
+ let ordering: ArrayRef = Arc::new(Int32Array::from(vec![1, 9, 2]));
+ // Filter out the row with ordering 9, so the max becomes ordering 2
-> value 30.
+ let filter = BooleanArray::from(vec![true, false, true]);
+ acc.update_batch(&[values, ordering], &[0, 0, 0], Some(&filter), 1)
+ .unwrap();
+ assert_eq!(eval_int(&mut acc), vec![Some(30)]);
+ }
+
+ #[test]
+ fn groups_accumulator_size_counts_row_bytes() {
+ // Each group owns two `OwnedRow`s, and each one holds its bytes in a
separate `Box<[u8]>`.
+ // Counting only `size_of::<OwnedRow>()` per slot leaves that payload
invisible to the
+ // memory pool that drives spill decisions.
+ const GROUPS: usize = 1_000;
+ let mut acc = max_by_groups(DataType::Int32, DataType::Int64);
+ let values: ArrayRef = Arc::new(Int32Array::from((0..GROUPS as
i32).collect::<Vec<_>>()));
+ let ordering: ArrayRef = Arc::new(Int64Array::from((0..GROUPS as
i64).collect::<Vec<_>>()));
+ acc.update_batch(
+ &[values, ordering],
+ &(0..GROUPS).collect::<Vec<_>>(),
+ None,
+ GROUPS,
+ )
+ .unwrap();
+
+ // A row is a one-byte null sentinel plus the fixed-width encoding: 5
bytes for the Int32
+ // value and 9 for the Int64 ordering.
+ let payload = GROUPS * (5 + 9);
+ let headers = (acc.best_value.capacity() +
acc.best_ordering.capacity())
+ * std::mem::size_of::<OwnedRow>();
+ assert!(
+ acc.size() >= size_of_val(&acc) + headers + payload,
+ "size() = {} does not cover {payload} bytes of row payload on top
of {headers} bytes \
+ of slot headers",
+ acc.size()
+ );
+ }
+}
diff --git a/native/spark-expr/src/agg_funcs/mod.rs
b/native/spark-expr/src/agg_funcs/mod.rs
index a63a8ddb0a..d0d8083195 100644
--- a/native/spark-expr/src/agg_funcs/mod.rs
+++ b/native/spark-expr/src/agg_funcs/mod.rs
@@ -22,6 +22,7 @@ mod correlation;
mod covariance;
mod hll_plus_plus;
mod hll_plus_plus_const;
+mod max_min_by;
mod mode;
mod percentile;
mod quantile_summaries;
@@ -38,6 +39,7 @@ pub use avg_decimal::AvgDecimal;
pub use correlation::Correlation;
pub use covariance::Covariance;
pub use hll_plus_plus::{hllpp_precision, HllPlusPlus};
+pub use max_min_by::MaxMinBy;
pub use mode::Mode;
pub use percentile::SparkPercentile;
pub use regr::{Regr, RegrType};
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index e748e82cd7..be4bc9c341 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -414,7 +414,9 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
classOf[First] -> CometFirst,
classOf[Last] -> CometLast,
classOf[Max] -> CometMax,
+ classOf[MaxBy] -> CometMaxBy,
classOf[Min] -> CometMin,
+ classOf[MinBy] -> CometMinBy,
classOf[Mode] -> CometMode,
classOf[Percentile] -> CometPercentile,
classOf[RegrIntercept] -> CometRegrIntercept,
diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
index 4f4b0f5f41..aef649debc 100644
--- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
@@ -22,7 +22,7 @@ package org.apache.comet.serde
import scala.jdk.CollectionConverters._
import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression,
Literal}
-import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg,
BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count,
Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max,
Min, Mode, Percentile, RegrIntercept, RegrR2, RegrReplacement, RegrSlope,
RegrSXY, StddevPop, StddevSamp, Sum, VariancePop, VarianceSamp}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
ApproximatePercentile, Average, BitAndAgg, BitOrAgg, BitXorAgg,
BloomFilterAggregate, CentralMomentAgg, CollectList, CollectSet, Corr, Count,
Covariance, CovPopulation, CovSample, First, HyperLogLogPlusPlus, Last, Max,
MaxBy, MaxMinBy, Min, MinBy, Mode, Percentile, RegrIntercept, RegrR2,
RegrReplacement, RegrSlope, RegrSXY, StddevPop, StddevSamp, Sum, VariancePop,
VarianceSamp}
import org.apache.spark.sql.catalyst.util.ArrayData
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType,
DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType,
NumericType, ShortType, StringType, TimestampNTZType, TimestampType}
@@ -105,6 +105,89 @@ object CometMax extends CometAggregateExpressionSerde[Max]
{
}
}
+/**
+ * Shared serde for `max_by` and `min_by`. Both are 2-argument `MaxMinBy`
`DeclarativeAggregate`s
+ * differing only in the comparison direction, so the value/ordering handling
is identical.
+ */
+abstract class CometMaxMinBy[T <: MaxMinBy] extends
CometAggregateExpressionSerde[T] {
+
+ /** `true` for `max_by`, `false` for `min_by`. */
+ protected def isMax: Boolean
+
+ /** `"maximum"` or `"minimum"`, used in the non-determinism note. */
+ private def extremum: String = if (isMax) "maximum" else "minimum"
+
+ override def getCompatibleNotes(): Seq[String] = Seq(
+ s"This function is non-deterministic when multiple rows share the
$extremum ordering value." +
+ " Results may differ from Spark in that case.")
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(
+ "The value and ordering must both be fixed-length types (boolean,
integral, floating-point," +
+ " decimal, date, or timestamp). A variable-length or nested type such as
string, binary, or" +
+ " struct falls back to Spark.")
+
+ override def getSupportLevel(expr: T): SupportLevel = {
+ // Both the value and ordering must be fixed-length types.
+ //
+ // On its own a variable-length type never reaches here: Spark only uses
HashAggregate (the
+ // aggregate operator Comet accelerates) when the aggregation buffer is
mutable, and the buffer
+ // holds both the running value and the running ordering, so a StringType
in either position
+ // forces SortAggregate, which Comet does not convert.
+ //
+ // The check is still load-bearing, because a TypedImperativeAggregate
elsewhere in the same
+ // aggregate switches Spark to ObjectHashAggregate, which Comet does
convert. In that shape a
+ // string ordering would otherwise be compared by Arrow's row format as
raw UTF-8 bytes, while
+ // Spark compares collation sort keys. See the fallback cases in
max_by.sql.
+ //
+ // The native side compares the ordering column via Arrow's row format,
which supports all of
+ // the fixed-length orderable types allowed below.
+ if (!AggSerde.minMaxDataTypeSupported(expr.valueExpr.dataType)) {
+ Unsupported(Some(s"Unsupported value data type:
${expr.valueExpr.dataType}"))
+ } else if (!AggSerde.minMaxDataTypeSupported(expr.orderingExpr.dataType)) {
+ Unsupported(Some(s"Unsupported ordering data type:
${expr.orderingExpr.dataType}"))
+ } else {
+ Compatible()
+ }
+ }
+
+ override def convert(
+ aggExpr: AggregateExpression,
+ expr: T,
+ inputs: Seq[Attribute],
+ binding: Boolean,
+ conf: SQLConf): Option[ExprOuterClass.AggExpr] = {
+ val valueExpr = exprToProto(expr.valueExpr, inputs, binding)
+ val orderingExpr = exprToProto(expr.orderingExpr, inputs, binding)
+
+ if (valueExpr.isDefined && orderingExpr.isDefined) {
+ val aggBuilder = ExprOuterClass.AggExpr.newBuilder()
+ if (isMax) {
+ val builder = ExprOuterClass.MaxBy.newBuilder()
+ builder.setValue(valueExpr.get)
+ builder.setOrdering(orderingExpr.get)
+ aggBuilder.setMaxBy(builder)
+ } else {
+ val builder = ExprOuterClass.MinBy.newBuilder()
+ builder.setValue(valueExpr.get)
+ builder.setOrdering(orderingExpr.get)
+ aggBuilder.setMinBy(builder)
+ }
+ Some(aggBuilder.build())
+ } else {
+ withFallbackReason(aggExpr, "Child expression or data type not
supported")
+ None
+ }
+ }
+}
+
+object CometMaxBy extends CometMaxMinBy[MaxBy] {
+ override protected def isMax: Boolean = true
+}
+
+object CometMinBy extends CometMaxMinBy[MinBy] {
+ override protected def isMax: Boolean = false
+}
+
object CometCount extends CometAggregateExpressionSerde[Count] {
override def convert(
aggExpr: AggregateExpression,
diff --git
a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
new file mode 100644
index 0000000000..9c8354f06a
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql
@@ -0,0 +1,284 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- max_by(x, y) returns the value of x associated with the maximum value of y.
+--
+-- The value (x) must be a fixed-length type: Spark only uses HashAggregate
(the aggregate
+-- operator Comet accelerates) when the aggregation buffer is mutable, so
variable-length
+-- value types such as string force SortAggregate and fall back to Spark.
+--
+-- Ordering values are kept unique within each group so results are
deterministic (max_by is
+-- non-deterministic when several rows tie on the maximum ordering).
+
+-- ============================================================
+-- Setup: tables
+-- ============================================================
+
+statement
+CREATE TABLE mb_src(v int, ord int, grp string) USING parquet
+
+statement
+INSERT INTO mb_src VALUES
+ (10, 10, 'g1'), (20, 50, 'g1'), (30, 20, 'g1'),
+ (40, 40, 'g2'), (50, 5, 'g2'), (60, 30, 'g2'),
+ (70, 99, 'g3')
+
+-- ordering NULLs are ignored; a group of all-NULL orderings yields NULL
+statement
+CREATE TABLE mb_nulls(v int, ord int, grp string) USING parquet
+
+statement
+INSERT INTO mb_nulls VALUES
+ (1, 10, 'g1'), (2, NULL, 'g1'), (3, 5, 'g1'),
+ (4, NULL, 'g2'), (5, NULL, 'g2'),
+ (6, 7, 'g3'), (NULL, 100, 'g3')
+
+statement
+CREATE TABLE mb_empty(v int, ord int) USING parquet
+
+-- ============================================================
+-- Global aggregate (no GROUP BY)
+-- ============================================================
+
+query
+SELECT max_by(v, ord) FROM mb_src
+
+-- ============================================================
+-- GROUP BY
+-- ============================================================
+
+query
+SELECT grp, max_by(v, ord) FROM mb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- NULL handling: NULL orderings ignored; the value paired with the
+-- maximum ordering may itself be NULL; all-NULL orderings yield NULL
+-- ============================================================
+
+query
+SELECT grp, max_by(v, ord) FROM mb_nulls GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Empty table yields NULL
+-- ============================================================
+
+query
+SELECT max_by(v, ord) FROM mb_empty
+
+-- ============================================================
+-- Literal arguments (evaluated natively; constant folding is disabled)
+-- ============================================================
+
+query
+SELECT max_by(5, 10), max_by(CAST(NULL AS INT), 20)
+
+-- ============================================================
+-- Mixed with other aggregates
+-- ============================================================
+
+query
+SELECT grp, max_by(v, ord), count(*), max(ord)
+FROM mb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- BigInt value
+-- ============================================================
+
+statement
+CREATE TABLE mb_long(val bigint, ord int, grp string) USING parquet
+
+statement
+INSERT INTO mb_long VALUES
+ (1000000000000, 1, 'a'), (2000000000000, 3, 'a'), (3000000000000, 2, 'a'),
+ (4000000000000, 5, 'b'), (5000000000000, 4, 'b')
+
+query
+SELECT grp, max_by(val, ord) FROM mb_long GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Double value and double ordering (NaN is the maximum in Spark)
+-- ============================================================
+
+statement
+CREATE TABLE mb_dbl(v double, ord double, grp string) USING parquet
+
+statement
+INSERT INTO mb_dbl VALUES
+ (1.1, 1.5, 'g1'), (2.2, 2.5, 'g1'), (3.3, 0.5, 'g1'),
+ (4.4, 1.0, 'g2'), (5.5, CAST('NaN' AS DOUBLE), 'g2'), (6.6, 100.0, 'g2'),
+ (7.7, CAST('Infinity' AS DOUBLE), 'g3'), (8.8, 3.0, 'g3')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_dbl GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Decimal value and decimal ordering
+-- ============================================================
+
+statement
+CREATE TABLE mb_dec(v decimal(10,2), ord decimal(10,2), grp string) USING
parquet
+
+statement
+INSERT INTO mb_dec VALUES
+ (10.01, 1.50, 'g1'), (20.02, 9.99, 'g1'), (30.03, 5.00, 'g1'),
+ (40.04, 2.00, 'g2'), (50.05, 8.25, 'g2')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_dec GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Date / timestamp value and ordering
+-- ============================================================
+
+statement
+CREATE TABLE mb_dt(d date, ts timestamp, ord int, grp string) USING parquet
+
+statement
+INSERT INTO mb_dt VALUES
+ (DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', 1, 'g1'),
+ (DATE '2024-06-15', TIMESTAMP '2024-06-15 12:30:00', 3, 'g1'),
+ (DATE '2023-12-31', TIMESTAMP '2023-12-31 23:59:59', 2, 'g1'),
+ (DATE '2024-03-01', TIMESTAMP '2024-03-01 08:00:00', 1, 'g2')
+
+query
+SELECT grp, max_by(d, ord) FROM mb_dt GROUP BY grp ORDER BY grp
+
+query
+SELECT grp, max_by(ts, ord) FROM mb_dt GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Date / timestamp as the ordering column
+-- ============================================================
+
+query
+SELECT grp, max_by(ord, d) FROM mb_dt GROUP BY grp ORDER BY grp
+
+query
+SELECT grp, max_by(ord, ts) FROM mb_dt GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Negative and boundary ordering values
+-- ============================================================
+
+statement
+CREATE TABLE mb_bound(v int, ord bigint, grp string) USING parquet
+
+statement
+INSERT INTO mb_bound VALUES
+ (1, -100, 'g1'), (2, -5, 'g1'), (3, -50,
'g1'),
+ (4, -9223372036854775808, 'g2'), (5, 9223372036854775807, 'g2'), (6, 0,
'g2'),
+ (7, -2147483648, 'g3'), (8, 2147483647, 'g3')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_bound GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Double ordering with -Infinity and Infinity
+-- ============================================================
+
+statement
+CREATE TABLE mb_inf(v int, ord double, grp string) USING parquet
+
+statement
+INSERT INTO mb_inf VALUES
+ (1, CAST('-Infinity' AS DOUBLE), 'g1'), (2, -1.0, 'g1'), (3, CAST('Infinity'
AS DOUBLE), 'g1'),
+ (4, CAST('-Infinity' AS DOUBLE), 'g2'), (5, -2.0, 'g2')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_inf GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Multiple max_by in one query
+-- ============================================================
+
+query
+SELECT grp, max_by(v, ord), max_by(ord, v) FROM mb_src GROUP BY grp ORDER BY
grp
+
+-- ============================================================
+-- Value and ordering are the same column
+-- ============================================================
+
+query
+SELECT grp, max_by(ord, ord) FROM mb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Boolean value
+-- ============================================================
+
+statement
+CREATE TABLE mb_bool(v boolean, ord int, grp string) USING parquet
+
+statement
+INSERT INTO mb_bool VALUES
+ (true, 1, 'a'), (false, 3, 'a'), (true, 2, 'a'),
+ (false, 5, 'b'), (true, 4, 'b')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_bool GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Signed zeros in the ordering
+--
+-- Spark compares the ordering with SQLOrderingUtil.compareDoubles, whose `if
(x == y) 0`
+-- short-circuit ties -0.0 with 0.0. Arrow's row format instead ranks -0.0
strictly below 0.0, so
+-- the native side canonicalizes the ordering column before encoding it.
+--
+-- The tied rows deliberately share the same value (v), which keeps this query
deterministic. A
+-- fixture that gave the tied rows different values could not be asserted
here: which tied row wins
+-- depends on the inter-partition merge order, and this table spans several
partitions, so max_by is
+-- genuinely non-deterministic in that case (as Spark documents). The
order-dependent behaviour is
+-- pinned in the Rust unit tests in max_min_by.rs instead, where row order is
explicit. This case
+-- still exercises the canonicalization path end to end.
+--
+-- Negative zero must be written as -0.0D: an unsuffixed -0.0 is a DecimalType
literal, and Decimal
+-- has no signed zero, so CAST(-0.0 AS DOUBLE) would silently store +0.0.
+-- ============================================================
+
+statement
+CREATE TABLE mb_signed_zero(v int, ord double, grp string) USING parquet
+
+statement
+INSERT INTO mb_signed_zero VALUES
+ (1, -0.0D, 'g1'), (1, 0.0D, 'g1'), (2, -5.0D, 'g1'),
+ (3, 0.0D, 'g2'), (3, -0.0D, 'g2'), (4, -5.0D, 'g2')
+
+query
+SELECT grp, max_by(v, ord) FROM mb_signed_zero GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Variable-length value or ordering falls back to Spark
+--
+-- A plain max_by over a string is planned as SortAggregate, which Comet never
converts, so the
+-- serde's type check is not what stops it. Pairing it with a
TypedImperativeAggregate switches
+-- Spark to ObjectHashAggregate, which Comet does convert, and then
getSupportLevel is the only
+-- thing that keeps the aggregate off the native path. That matters most for a
string *ordering*:
+-- Arrow's row format compares raw UTF-8 bytes, while Spark compares collation
sort keys, so this
+-- check is load-bearing for correctness and not just an optimisation.
+-- ============================================================
+
+statement
+CREATE TABLE mb_varlen(v int, s string, grp string) USING parquet
+
+statement
+INSERT INTO mb_varlen VALUES
+ (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2')
+
+query expect_fallback(Unsupported value data type)
+SELECT grp, max_by(s, v), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER
BY grp
+
+query expect_fallback(Unsupported ordering data type)
+SELECT grp, max_by(v, s), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER
BY grp
diff --git
a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
new file mode 100644
index 0000000000..d015acc35d
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql
@@ -0,0 +1,266 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- min_by(x, y) returns the value of x associated with the minimum value of y.
+--
+-- The value (x) must be a fixed-length type: Spark only uses HashAggregate
(the aggregate
+-- operator Comet accelerates) when the aggregation buffer is mutable, so
variable-length
+-- value types such as string force SortAggregate and fall back to Spark.
+--
+-- Ordering values are kept unique within each group so results are
deterministic (min_by is
+-- non-deterministic when several rows tie on the minimum ordering).
+
+-- ============================================================
+-- Setup: tables
+-- ============================================================
+
+statement
+CREATE TABLE nb_src(v int, ord int, grp string) USING parquet
+
+statement
+INSERT INTO nb_src VALUES
+ (10, 10, 'g1'), (20, 50, 'g1'), (30, 20, 'g1'),
+ (40, 40, 'g2'), (50, 5, 'g2'), (60, 30, 'g2'),
+ (70, 99, 'g3')
+
+-- ordering NULLs are ignored; a group of all-NULL orderings yields NULL
+statement
+CREATE TABLE nb_nulls(v int, ord int, grp string) USING parquet
+
+statement
+INSERT INTO nb_nulls VALUES
+ (1, 10, 'g1'), (2, NULL, 'g1'), (3, 5, 'g1'),
+ (4, NULL, 'g2'), (5, NULL, 'g2'),
+ (6, 7, 'g3'), (NULL, 1, 'g3')
+
+statement
+CREATE TABLE nb_empty(v int, ord int) USING parquet
+
+-- ============================================================
+-- Global aggregate (no GROUP BY)
+-- ============================================================
+
+query
+SELECT min_by(v, ord) FROM nb_src
+
+-- ============================================================
+-- GROUP BY
+-- ============================================================
+
+query
+SELECT grp, min_by(v, ord) FROM nb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- NULL handling: NULL orderings ignored; the value paired with the
+-- minimum ordering may itself be NULL; all-NULL orderings yield NULL
+-- ============================================================
+
+query
+SELECT grp, min_by(v, ord) FROM nb_nulls GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Empty table yields NULL
+-- ============================================================
+
+query
+SELECT min_by(v, ord) FROM nb_empty
+
+-- ============================================================
+-- Literal arguments (evaluated natively; constant folding is disabled)
+-- ============================================================
+
+query
+SELECT min_by(5, 10), min_by(CAST(NULL AS INT), 20)
+
+-- ============================================================
+-- Mixed with other aggregates
+-- ============================================================
+
+query
+SELECT grp, min_by(v, ord), count(*), min(ord)
+FROM nb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- BigInt value
+-- ============================================================
+
+statement
+CREATE TABLE nb_long(val bigint, ord int, grp string) USING parquet
+
+statement
+INSERT INTO nb_long VALUES
+ (1000000000000, 1, 'a'), (2000000000000, 3, 'a'), (3000000000000, 2, 'a'),
+ (4000000000000, 5, 'b'), (5000000000000, 4, 'b')
+
+query
+SELECT grp, min_by(val, ord) FROM nb_long GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Double value and double ordering (NaN is the maximum in Spark, so it is
+-- never selected by min_by unless it is the only non-null ordering)
+-- ============================================================
+
+statement
+CREATE TABLE nb_dbl(v double, ord double, grp string) USING parquet
+
+statement
+INSERT INTO nb_dbl VALUES
+ (1.1, 1.5, 'g1'), (2.2, 2.5, 'g1'), (3.3, 0.5, 'g1'),
+ (4.4, 1.0, 'g2'), (5.5, CAST('NaN' AS DOUBLE), 'g2'), (6.6, 100.0, 'g2'),
+ (7.7, CAST('-Infinity' AS DOUBLE), 'g3'), (8.8, 3.0, 'g3')
+
+query
+SELECT grp, min_by(v, ord) FROM nb_dbl GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Decimal value and decimal ordering
+-- ============================================================
+
+statement
+CREATE TABLE nb_dec(v decimal(10,2), ord decimal(10,2), grp string) USING
parquet
+
+statement
+INSERT INTO nb_dec VALUES
+ (10.01, 1.50, 'g1'), (20.02, 9.99, 'g1'), (30.03, 5.00, 'g1'),
+ (40.04, 2.00, 'g2'), (50.05, 8.25, 'g2')
+
+query
+SELECT grp, min_by(v, ord) FROM nb_dec GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Date / timestamp value and ordering
+-- ============================================================
+
+statement
+CREATE TABLE nb_dt(d date, ts timestamp, ord int, grp string) USING parquet
+
+statement
+INSERT INTO nb_dt VALUES
+ (DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', 1, 'g1'),
+ (DATE '2024-06-15', TIMESTAMP '2024-06-15 12:30:00', 3, 'g1'),
+ (DATE '2023-12-31', TIMESTAMP '2023-12-31 23:59:59', 2, 'g1'),
+ (DATE '2024-03-01', TIMESTAMP '2024-03-01 08:00:00', 1, 'g2')
+
+query
+SELECT grp, min_by(d, ord) FROM nb_dt GROUP BY grp ORDER BY grp
+
+query
+SELECT grp, min_by(ts, ord) FROM nb_dt GROUP BY grp ORDER BY grp
+
+query
+SELECT grp, min_by(ord, d) FROM nb_dt GROUP BY grp ORDER BY grp
+
+query
+SELECT grp, min_by(ord, ts) FROM nb_dt GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Negative and boundary ordering values
+-- ============================================================
+
+statement
+CREATE TABLE nb_bound(v int, ord bigint, grp string) USING parquet
+
+statement
+INSERT INTO nb_bound VALUES
+ (1, -100, 'g1'), (2, -5, 'g1'), (3, -50,
'g1'),
+ (4, -9223372036854775808, 'g2'), (5, 9223372036854775807, 'g2'), (6, 0,
'g2'),
+ (7, -2147483648, 'g3'), (8, 2147483647, 'g3')
+
+query
+SELECT grp, min_by(v, ord) FROM nb_bound GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Multiple min_by and value = ordering column
+-- ============================================================
+
+query
+SELECT grp, min_by(v, ord), min_by(ord, v) FROM nb_src GROUP BY grp ORDER BY
grp
+
+query
+SELECT grp, min_by(ord, ord) FROM nb_src GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Boolean value
+-- ============================================================
+
+statement
+CREATE TABLE nb_bool(v boolean, ord int, grp string) USING parquet
+
+statement
+INSERT INTO nb_bool VALUES
+ (true, 1, 'a'), (false, 3, 'a'), (true, 2, 'a'),
+ (false, 5, 'b'), (true, 4, 'b')
+
+query
+SELECT grp, min_by(v, ord) FROM nb_bool GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- max_by and min_by in the same query
+-- ============================================================
+
+query
+SELECT grp, max_by(v, ord), min_by(v, ord) FROM nb_src GROUP BY grp ORDER BY
grp
+
+-- ============================================================
+-- Signed zeros in the ordering
+--
+-- Spark compares the ordering with SQLOrderingUtil.compareDoubles, whose `if
(x == y) 0`
+-- short-circuit ties -0.0 with 0.0. Arrow's row format instead ranks -0.0
strictly below 0.0, so
+-- the native side canonicalizes the ordering column before encoding it.
+--
+-- The tied rows deliberately share the same value (v), which keeps this query
deterministic. A
+-- fixture that gave the tied rows different values could not be asserted
here: which tied row wins
+-- depends on the inter-partition merge order, and this table spans several
partitions, so min_by is
+-- genuinely non-deterministic in that case (as Spark documents). The
order-dependent behaviour is
+-- pinned in the Rust unit tests in max_min_by.rs instead, where row order is
explicit. This case
+-- still exercises the canonicalization path end to end.
+--
+-- Negative zero must be written as -0.0D: an unsuffixed -0.0 is a DecimalType
literal, and Decimal
+-- has no signed zero, so CAST(-0.0 AS DOUBLE) would silently store +0.0.
+-- ============================================================
+
+statement
+CREATE TABLE mnb_signed_zero(v int, ord double, grp string) USING parquet
+
+statement
+INSERT INTO mnb_signed_zero VALUES
+ (1, -0.0D, 'g1'), (1, 0.0D, 'g1'), (2, 5.0D, 'g1'),
+ (3, 0.0D, 'g2'), (3, -0.0D, 'g2'), (4, 5.0D, 'g2')
+
+query
+SELECT grp, min_by(v, ord) FROM mnb_signed_zero GROUP BY grp ORDER BY grp
+
+-- ============================================================
+-- Variable-length value or ordering falls back to Spark
+--
+-- See the equivalent section in max_by.sql for why this needs a
TypedImperativeAggregate
+-- alongside it: without one, Spark plans SortAggregate and Comet never sees
the aggregate at all,
+-- so the serde's type check would go untested.
+-- ============================================================
+
+statement
+CREATE TABLE mnb_varlen(v int, s string, grp string) USING parquet
+
+statement
+INSERT INTO mnb_varlen VALUES
+ (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2')
+
+query expect_fallback(Unsupported value data type)
+SELECT grp, min_by(s, v), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp
ORDER BY grp
+
+query expect_fallback(Unsupported ordering data type)
+SELECT grp, min_by(v, s), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp
ORDER BY grp
diff --git
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala
index 7aa66fb0a9..2772d37ffb 100644
---
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala
+++
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala
@@ -168,6 +168,28 @@ object CometAggregateExpressionBenchmark extends
CometBenchmarkBase {
"approx_percentile_double_high_card",
"SELECT approx_percentile(c_double, 0.5) FROM parquetV1Table GROUP BY
high_card_grp"))
+ // max_by / min_by run natively only when both the value and ordering are
fixed-length types (a
+ // variable-length value or ordering forces Spark's SortAggregate, which
Comet does not run).
+ private val maxMinByAggregates = List(
+ AggExprConfig("max_by_int", "SELECT max_by(c_int, c_long) FROM
parquetV1Table GROUP BY grp"),
+ AggExprConfig(
+ "max_by_double",
+ "SELECT max_by(c_double, c_int) FROM parquetV1Table GROUP BY grp"),
+ AggExprConfig(
+ "max_by_decimal",
+ "SELECT max_by(c_decimal, c_int) FROM parquetV1Table GROUP BY grp"),
+ AggExprConfig(
+ "max_by_high_card",
+ "SELECT max_by(c_int, c_long) FROM parquetV1Table GROUP BY
high_card_grp"),
+ AggExprConfig("max_by_global", "SELECT max_by(c_int, c_long) FROM
parquetV1Table"),
+ AggExprConfig("min_by_int", "SELECT min_by(c_int, c_long) FROM
parquetV1Table GROUP BY grp"),
+ AggExprConfig(
+ "min_by_double",
+ "SELECT min_by(c_double, c_int) FROM parquetV1Table GROUP BY grp"),
+ AggExprConfig(
+ "min_by_high_card",
+ "SELECT min_by(c_int, c_long) FROM parquetV1Table GROUP BY
high_card_grp"))
+
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
val values = 1024 * 1024
@@ -192,7 +214,7 @@ object CometAggregateExpressionBenchmark extends
CometBenchmarkBase {
val allAggregates = basicAggregates ++ statisticalAggregates ++
bitwiseAggregates ++
multiKeyAggregates ++ multiAggregates ++ decimalAggregates ++
highCardinalityAggregates ++ percentileAggregates ++
approxPercentileAggregates ++
- approxCountDistinctAggregates
+ approxCountDistinctAggregates ++ maxMinByAggregates
allAggregates.foreach { config =>
runExpressionBenchmark(config.name, v, config.query,
config.extraCometConfigs)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
