feat: support native aggregate function mode

[andygrove]

Which issue does this PR close?

Closes #3970.

Rationale for this change

The mode aggregate (the most frequent value in a group) is a mainstream statistical aggregate that previously fell back to Spark, preventing native execution of any query using it. Adding native support keeps these queries in Comet's pipeline.

What changes are included in this PR?

This PR was scaffolded with the implement-comet-expression project skill.

• Native Rust implementation in native/spark-expr/src/agg_funcs/mode.rs: a Mode aggregate UDF with both a global Accumulator and a vectorized GroupsAccumulator. State is a frequency map keyed by ScalarValue, serialized as a single struct<values: array<T>, counts: array<bigint>> buffer column so partial/final buffer schemas stay aligned with Spark's single-attribute TypedImperativeAggregate buffer.
• Protobuf Mode message and AggExpr oneof entry, plus the planner arm in planner.rs.
• CometMode serde and registration in QueryPlanSerde.aggrSerdeMap.
• A Mode branch in adjustOutputForNativeState (operators.scala) mapping the Spark binary buffer type to the native struct state type.
• A modeHasUnsupportedOrdering shim in CometTypeShim (spark-3.x / spark-4.x) because Mode.reverseOpt only exists on Spark 4.0+.
• Docs: mode marked supported in the expressions guide.

Scope and compatibility:

• Only the plain mode(col) form is supported. The mode(col, deterministic) and mode() WITHIN GROUP (ORDER BY col) forms (Spark 4.0+, which set reverseOpt) fall back to Spark.
• Registered as Incompatible (opt-in via spark.comet.expression.Mode.allowIncompatible=true): Spark breaks ties non-deterministically based on JVM hash-map iteration order, which a native hash map cannot reproduce bit-for-bit. Comet instead returns the smallest tied value deterministically.
• NULLs are ignored, empty input returns NULL, and float keys are normalized (-0.0 to 0.0, canonical NaN) to match Spark's counting. Supported input types are numeric, boolean, decimal, date, timestamp, timestamp_ntz, and default-collation string; other types fall back.

How are these changes tested?

• Rust unit tests in mode.rs covering most-frequent value, tie-break to smallest, NULL handling, empty input, float normalization, and partial/final merge equivalence for both the accumulator and the groups accumulator.
• A mode.sql file test exercising global and grouped aggregation, NULLs, all-NULL groups, mixed aggregates, HAVING, and boolean/integer/double/decimal/string/date/timestamp inputs, plus an unsupported-type fallback assertion. Verified on Spark 3.5 and Spark 4.1.

[mbutrovich]

First pass, thanks @andygrove! I verified null handling, empty-group-returns-NULL, the tie-break rationale, ordering via ScalarValue::partial_cmp, collation and type gating, and the state-schema rewrite against apple-spark .../aggregate/Mode.scala. Those all match. One item is a genuine semantic divergence from Spark that is not covered by the non-determinism label: the -0.0 normalization. Details and a test below. The rest are reuse/efficiency/cleanup items.

[mbutrovich]

normalize_key collapses both NaN and -0.0 before keying the frequency map. The NaN half is correct, but collapsing -0.0 into 0.0 diverges from Spark. The Spark-side semantics are verified:

• Spark's Mode.update keys an OpenHashMap[AnyRef, Long] via buffer.changeValue(InternalRow.copyValue(key), ...) with no normalization (Mode.scala:63).
• OpenHashMap/OpenHashSet key comparison is _data(pos) equals k (core/.../util/collection/OpenHashSet.scala:122), and that file carries an explicit comment that equals distinguishes 0.0/-0.0 and collapses NaN/NaN (OpenHashSet.scala:116). So Spark treats -0.0 and 0.0 as distinct keys and all NaN as one key.
• NormalizeFloatingNumbers does not touch aggregate arguments. Its apply only rewrites WINDOW and JOIN patterns (sql/catalyst/.../optimizer/NormalizeFloatingNumbers.scala:62). There is no Aggregate case, so a mode(v) argument reaches the aggregate un-normalized.

So the doc comment at mode.rs:40-41 claiming this matches NormalizeFloatingNumbers does not hold for aggregate inputs. The Spark-equivalent key is doubleToLongBits: canonicalize NaN, keep -0.0 distinct from 0.0. Suggested fix: drop the -0.0 branch from normalize_key, keep only NaN canonicalization, then key on bit-equality (arrow's Hashable/total_cmp keeps -0.0 and 0.0 distinct once NaN is canonical).

One thing I could not confirm without running: whether this divergence is observable end to end, which depends on two harness details. First, whether the Parquet write/read round-trip preserves -0.0 on the Spark baseline side. Second, how the row comparison treats -0.0 vs 0.0. The existing group b fixture at mode.sql:115-120 uses -0.0 and passed CI, which suggests one of those is masking the difference today. A test that flips the winner to a different magnitude sidesteps both, so it is unambiguous. Does this hold up if you run it?

statement
CREATE TABLE mode_neg_zero(v double, grp string) USING parquet

statement
INSERT INTO mode_neg_zero VALUES
  (CAST(-0.0 AS DOUBLE), 'a'), (CAST(0.0 AS DOUBLE), 'a'),
  (CAST(5.0 AS DOUBLE), 'a'), (CAST(5.0 AS DOUBLE), 'a')

-- Hypothesis: Spark keeps -0.0 and 0.0 separate (counts -0.0:1, 0.0:1, 5.0:2) so mode = 5.0,
-- while Comet collapses to 0.0:2, 5.0:2, a tie it breaks to the smaller value 0.0.
-- 5.0 vs 0.0 is caught regardless of how -0.0/0.0 compare.
query
SELECT grp, mode(v) FROM mode_neg_zero GROUP BY grp ORDER BY grp

Keep a NaN case to lock in that NaN must still collapse:

statement
CREATE TABLE mode_nan(v double, grp string) USING parquet

statement
INSERT INTO mode_nan VALUES
  (CAST('NaN' AS DOUBLE), 'a'), (CAST('NaN' AS DOUBLE), 'a'), (CAST(1.0 AS DOUBLE), 'a')

query
SELECT grp, mode(v) FROM mode_nan GROUP BY grp ORDER BY grp

If the fix lands, the existing group b fixture (mode.sql:115) and the Rust test float_zero_and_nan_normalized (mode.rs:454) encode the collapsed result, so they would need updating.

[mbutrovich]

Cross-linking a related finding, since the two are easy to conflate and the fixes point in opposite directions.

#4817 (max_by/min_by) has the mirror-image bug on the same input. There the ordering column is compared with SQLOrderingUtil.compareDoubles, which is if (x == y) 0 else Double.compare(x, y) and explicitly documents -0.0 == 0.0 (SQLOrderingUtil.scala:27-31), so Spark ties the two and Comet must collapse them. Here mode keys on OpenHashSet, whose equals distinguishes -0.0 from 0.0, so Comet must stop collapsing them.

Same two input values, opposite correct behavior, because the two aggregates reach different Spark comparison paths. Whichever way each is fixed, please add a one-line comment at each site naming the Spark path that governs it (OpenHashSet.equals here, SQLOrderingUtil.compareDoubles there), so the next person does not "fix" one to match the other.

[andygrove]

Thanks — this was the right catch, and it turned out to have a wrinkle worth recording.

Your analysis holds for every Spark version this repo currently supports. I confirmed all three
claims: branch-3.4, branch-3.5, branch-4.0 and branch-4.1 all key on
InternalRow.copyValue(key) with no normalization, OpenHashSet's equals carries that explicit
0.0/-0.0 vs NaN/NaN comment, and NormalizeFloatingNumbers.apply really is
transformWithPruning(_.containsAnyPattern(WINDOW, JOIN)), so an aggregate argument is never
normalized.

But Spark 4.2.0 reversed it. SPARK-57329
("mode() returns incorrect result when input contains both -0.0 and 0.0") treats the split counts as
Spark's own correctness bug and normalizes the key at update time via a new ModeKeyNormalizer
(DOUBLE_NORMALIZER/FLOAT_NORMALIZER). Its reasoning is the mirror of yours: Spark treats
-0.0 = 0.0 under SQL semantics everywhere else, so mode should too. The fix landed in
branch-4.2 after v4.2.0-rc1, so released 4.2.0 has it — I checked the tags — and this repo
already builds a spark-4.2 profile.

So correct behaviour is version-dependent, and neither always-collapsing nor never-collapsing is
right across the profiles we build. Rather than fix 3.4-4.1 and newly break 4.2, I version-gated it:
a normalize_neg_zero field on the Mode proto message, set from isSpark42Plus in the serde
(same shape as BloomFilterVersion above it and setIsSpark4Plus in CometCast), with the fold
gated on it natively.

NaN canonicalization is now unconditional, since doubleToLongBits collapses NaN on every
supported version. On your Hashable point: it turns out no extra work is needed, because
ScalarValue's PartialEq and Hash for Float32/Float64 are both defined on to_bits()
(datafusion-common/src/scalar/mod.rs, the Fl wrapper). So once NaN is canonical the zeros stay
distinct on their own.

On your open question — whether the divergence is observable end to end — it was not, and the reason
is worth knowing:

The existing group b fixture at mode.sql:115-120 uses -0.0 and passed CI, which suggests one
of those is masking the difference today.

The masking was in the fixture data, not the comparison. CAST(-0.0 AS DOUBLE) does not produce a
negative zero: an unsuffixed -0.0 is a DecimalType literal, and Decimal has no signed zero, so
the cast yields +0.0. I verified by reading back doubleToRawLongBits from the Parquet table — all
the "negative" zeros were bits=0. The column never contained a negative zero, so the fixture was
vacuous regardless of the bug. -0.0D, -CAST(0.0 AS DOUBLE) and CAST('-0.0' AS DOUBLE) all work;
SPARK-57329's own reproducer uses the D suffix for exactly this reason.

Switched the fixtures to -0.0D and added the SPARK-57329 shape (-0.0:2, 0.0:2, 5.0:3) as
mode_signed_zero, so the winner is 5.0 when the zeros stay apart and 0.0 when they fold, with no
tie either way. Then I checked it is not vacuous by forcing normalize_neg_zero to the wrong value
and confirming failure:

!== Spark Answer - 2 ==   == Comet Answer - 2 ==
![b,-0.0]                 [b,0.0]

Worth noting for future float fixtures: the harness does distinguish -0.0 from 0.0 in results.
I still designed the new case so the two candidate answers differ in magnitude, so it cannot depend
on that.

Also added a Rust test per direction plus a merge-path one, and updated
float_zero_and_nan_normalized, which encoded the collapsed result.

[andygrove]

Comments added at both sites, and the warning was well placed — the two do point in opposite
directions, and there is now a third direction to keep straight.

mode.rs records that the governing path is OpenHashSet.equals / doubleToLongBits, that
NormalizeFloatingNumbers does not reach aggregate arguments, and that SPARK-57329 changed the
answer in 4.2.0 (details in the thread above), ending with an explicit "do not simplify this to
always normalize, because max_by/min_by need the opposite treatment".

I am working through #4817 next and will put the mirror-image comment there naming
SQLOrderingUtil.compareDoubles. Given what turned up here, I will check that one across
branch-3.4 through branch-4.2 before implementing rather than assuming it is stable — if
compareDoubles has been touched the same way Mode was, that PR needs the same version gate.

[mbutrovich]

The frequency map keys on ScalarValue with per-row ScalarValue::try_from_array (see also line 195). DataFusion's idiomatic pattern for a primitive frequency/dedup map is a monomorphized HashMap<Hashable<T::Native>, _> (datafusion/functions-aggregate/src/median.rs:313, .../approx_distinct and count_distinct/native.rs). Hashable alone does not reproduce Spark's float key semantics (it keeps -0.0/0.0 distinct, which is correct per the finding above, but does not collapse NaN, so a canonicalization step is still needed). Given mode supports many types, the type-generic ScalarValue map is a defensible simplicity choice. Either keep it and add a one-line comment that it is intentionally type-generic, or monomorphize the hot primitive paths for speed. Not a defect.

[andygrove]

Kept it type-generic and added the comment. mode covers every primitive plus decimal, string and
the temporal types, so one generic map beats a kernel per type until a profile says otherwise — the
comment now says that explicitly so it reads as a decision rather than an oversight.

Your parenthetical turned out to be the useful part: Hashable alone would not have been enough, but
neither is it needed. ScalarValue's PartialEq and Hash for the float variants are both defined
on to_bits(), so bit-equality is already what the map gives us; only the NaN canonicalization step
has to be explicit. Noted in the normalize_key doc comment.

[mbutrovich]

(see also line 348) size() counts only capacity * size_of::<(ScalarValue, i64)>(), which omits the heap bytes of Utf8(String) / Binary(Vec<u8>) / boxed decimal keys. For string modes this under-reports memory to the pool that drives spill decisions. Compare count_distinct/native.rs, which includes element bytes via estimate_memory_size. Fix: sum key.size() per entry (ScalarValue has a size() helper) instead of a flat per-slot size.

[andygrove]

Fixed in both places. Extracted a map_size helper used by ModeAccumulator::size and
ModeGroupsAccumulator::size:

fn map_size(map: &HashMap<ScalarValue, i64>) -> usize {
    map.capacity() * size_of::<(ScalarValue, i64)>()
        + map
            .keys()
            .map(|k| k.size().saturating_sub(size_of::<ScalarValue>()))
            .sum::<usize>()
}

ScalarValue::size() includes the inline size_of::<ScalarValue>(), which capacity * size_of::<(ScalarValue, i64)>()
already accounts for, so the per-key term subtracts it to avoid double counting and adds only the
heap payload.

[mbutrovich]

default_value duplicates the AggregateUDFImpl trait default (udaf.rs, both ScalarValue::try_from(&data_type)). Redundant, can be removed.

[andygrove]

Removed. return_type returns self.data_type, which is what the trait default feeds to
ScalarValue::try_from, so it was exactly equivalent. The all-NULL-input coverage
(empty_input_is_null and the mode_all_null fixture) still passes.

[mbutrovich]

(and line 345) evaluate/state call ScalarValue::iter_to_array, which errors on an empty iterator. This is currently unreachable because the grouped-aggregate stream short-circuits when there are zero groups, but the dependency is implicit. To make the invariant explicit and self-documenting, add debug_assert!(!emitted.is_empty()) (or an early return of an empty array) at the emit site.

[andygrove]

Added debug_assert! at both grouped emit sites with a comment naming the dependency, so the
invariant is stated rather than implied:

let emitted = emit_to.take_needed(&mut self.groups);
// `ScalarValue::iter_to_array` errors on an empty iterator. The grouped-aggregate stream
// never emits zero groups, so this is unreachable; assert it rather than leaving the
// dependency implicit.
debug_assert!(!emitted.is_empty(), "mode: evaluate called with no groups");

and the same in state, where build_state funnels into the same call.

[mbutrovich]

Beyond the correctness tests above: add a timestamp_ntz compared query (declared supported in isSupportedType, only timestamp is exercised), and a Spark-4.x fallback assertion for mode(col, true) and mode() WITHIN GROUP (ORDER BY col) to lock in modeHasUnsupportedOrdering (only the type-fallback path is asserted today).

[andygrove]

Both added.

timestamp_ntz is now a compared query in mode.sql over a dedicated mode_ntz table (including a
NULL), so the type is exercised rather than just declared in isSupportedType.

The Spark 4.x ordering assertions went into a new mode_within_group.sql rather than this file,
because MinSparkVersion is a file-level directive and these forms do not parse on 3.x. It covers
mode(col, true), and WITHIN GROUP ascending and descending, plus a grouped variant.

One correction to what I first wrote there: I had also asserted mode(col, false) falls back, and it
does not — it runs natively, which is correct. ModeBuilder.build only constructs
new Mode(child, true) for the true case and rewrites mode(col, false) to the plain
Mode(child) with reverseOpt = None. So the false form genuinely is the plain form. That is now
a native query with a comment recording why, which incidentally pins that
modeHasUnsupportedOrdering must key off reverseOpt rather than the argument count.

The file also carries a sentinel native query, so a whole-expression regression cannot make the
fallback assertions pass vacuously.

[mbutrovich]

Spark ASC WITHIN GROUP (reverseOpt = Some(false)) returns the smallest tied value, identical to Comet's deterministic tie-break, so that specific form could be Compatible rather than a fallback. Optional, phrased as a change if you want it: add a TODO referencing it.

[andygrove]

Added the TODO rather than the behaviour change, to keep this PR scoped:

// TODO the ASC form (`reverseOpt = Some(false)`) returns the smallest tied value, which is
// exactly Comet's tie-break, so it could be served natively as `Compatible`.

Worth noting that reverseOpt = Some(false) is only reachable through
WITHIN GROUP (ORDER BY col) ascending — mode(col, false) collapses to reverseOpt = None in
ModeBuilder, so it is already native. That narrows the follow-up to the WITHIN GROUP ASC form
alone.

[andygrove]

Thanks for the thorough pass @mbutrovich — all items addressed in 47b0dfd. Detail is in the
individual threads; summarising the one item that changed the conclusion, since it affects how you
read the diff.

The -0.0 fix is version-dependent, so it is gated rather than dropped.

Your reasoning is correct for Spark 3.4 through 4.1, and I verified each claim against the release
branches. But Spark 4.2.0 reversed the behaviour:
SPARK-57329 treats the split -0.0/0.0 counts
as Spark's own correctness bug and normalizes the key at update time. It landed in branch-4.2 after
v4.2.0-rc1, so released 4.2.0 has it, and this repo already builds a spark-4.2 profile.

Spark	-0.0 and 0.0	mode of (-0.0 x2, 0.0 x2, 5.0 x3)
3.4 - 4.1	distinct keys	5.0
4.2.0+	one key	0.0

So simply dropping the -0.0 branch would have fixed 3.4-4.1 and newly broken 4.2. Instead there is
now a normalize_neg_zero field on the Mode proto message, set from isSpark42Plus in the serde
(the same shape as BloomFilterVersion directly above it, and setIsSpark4Plus in CometCast), with
the fold gated on it in the native accumulators. NaN canonicalization stayed unconditional.

The existing signed-zero coverage was vacuous, for a reason worth knowing. You flagged that the
group b fixture used -0.0 and passed CI anyway. The masking was in the data: CAST(-0.0 AS DOUBLE)
does not produce a negative zero, because an unsuffixed -0.0 is a DecimalType literal and Decimal
has no signed zero. I read doubleToRawLongBits back out of the Parquet table and every "negative"
zero was bits=0. Fixtures now use -0.0D, and I confirmed the new mode_signed_zero case is not
vacuous by forcing the flag to the wrong value and watching it fail.

Everything else from the review is in: size() now counts key heap bytes, the redundant
default_value is gone, debug_assert! at both grouped emit sites, a comment recording that the
ScalarValue map is intentionally type-generic, a timestamp_ntz compared query, and a new
mode_within_group.sql for the Spark 4.x ordered forms. One of my own additions there was wrong and
is corrected: mode(col, false) should not fall back, because ModeBuilder rewrites it to the plain
Mode(child).

Verified on Spark 4.1: 11 Rust unit tests, and all 23 files in expressions/aggregate/.

[rich7420]

Please include self.groups.capacity() * size_of::<HashMap<ScalarValue, i64>>() in size(). With 1,000,000 all-NULL groups, I measured 48 MB of vector storage while size() still reports 56 bytes. Please add a regression test so this allocation is included in spill accounting.

[andygrove]

Fixed in 3e6f125, and your 48 MB figure is exactly right — the Vec's slot array was invisible because size_of_val(self) only covers its inline pointer/len/capacity, and map_size legitimately reports 0 for an empty map.

size_of_val(self)
    + self.groups.capacity() * size_of::<HashMap<ScalarValue, i64>>()
    + self.groups.iter().map(map_size).sum::<usize>()

I used capacity() rather than len() so the accounting follows what the Vec actually holds after resize_with over-allocates.

The regression test is groups_accumulator_size_accounts_for_group_slots. It uses 10,000 all-NULL groups rather than 1,000,000 to keep the test fast, and it reproduces your measurement precisely — against the old size() it fails with:

size() = 56 does not cover 480000 bytes of empty group slots

Same 56 bytes you saw, and 48 bytes per group slot, which is the 48 MB at your group count. I checked it fails before the fix rather than only passing after.

[rich7420]

@andygrove thanks for the patch!! just left a comment

[mbutrovich]

There's a small gap on the testing: I think there are only double columns and not float, but I'm not terribly concerned about that. Thanks @andygrove!

[rich7420]

thanks for the update!

[sunchao]

Review summary

This PR adds opt-in native execution of the mode aggregate. No new blocking code findings at head 601397fed84c97e68da73e577df8e7f7c30315c3, reviewed against base 392da2ca7a99fd223644eff650406f61580eb46b. The implementation looks reasonable. Benchmark evidence and completed CI remain outstanding validation items.

Assessment

• Correctness: Partial/final state merging, null handling, signed-zero version differences, and fallback guards checked out. The earlier memory-accounting findings are fixed.
• Design and complexity: One frequency map per group and one structured state column fit the existing aggregation framework. The generic ScalarValue implementation is a reasonable simplicity tradeoff.
• Performance: Still unproven. The PR provides no Spark-versus-Comet benchmarks. Each input row requires scalar extraction and hashing, and memory reporting scans all stored keys. Please add benchmarks covering both repeated and mostly distinct values, including strings. These are measurement targets, not demonstrated regressions.

Validation

• 12/12 native unit tests passed using the unchanged PR source with cached DataFusion 55 / Arrow 59.3 dependencies. This was standalone compilation of the mode implementation and its tests, not a full native build.
• 8/8 selected Spark SQL fixture runs passed, including the two mode fixtures and 32 additional mode query checks covering:
  • mode mixed with COUNT(DISTINCT ...) and percentile;
  • aggregate filters and empty inputs;
  • 131,072 groups;
  • native/JVM shuffle with AQE enabled and disabled.
    Six fixture runs were mode-specific; the other two selected runs were make_time_shuffle fixtures.
• 18 Spark component cases passed across Spark 4.1.3 and 4.2.0. These exercised Spark's own Mode aggregation and serialized state merging as an oracle, rather than end-to-end Comet queries.

Build limitation: Local native compilation hit unavailable aws-smithy-runtime-api 1.16.0 in the configured package mirror. Spark testing used a checksum-verified CI library built from merge commit 4fbbd64e1c37462c623c4b6e077c5c86b2258eb7. Its aggregation code and dependencies match the reviewed head. The additional changes concern Iceberg writes. The native library SHA-256 was 463040833a8941b7a719de178fa6605ad8cded99ba89594387f71538cbe3adae.

One confirmed limitation

Forcing DataFusion's partial-aggregation skipping through the development/testing configuration causes mode to fail because convert_to_state is unimplemented. Comet explicitly disables that feature by default, so this does not establish a defect in the supported execution path.

CI

At publication: 31 successful, 37 running or queued, 7 skipped, no failures reported. CI run.

[andygrove]

Merged. Thanks @rich7420 @mbutrovich @sunchao !