feat: support native aggregate function mode#4782
Conversation
Add native support for the Spark `mode` aggregate, the most frequent value within a group. Spark breaks ties on the default `mode(col)` form non-deterministically (the chosen value depends on JVM hash-map iteration order), so the function is registered as Incompatible and opt-in via allowIncompatible; Comet resolves ties deterministically by returning the smallest tied value. NULLs are ignored, empty input returns NULL, and float keys are normalized to match Spark. The deterministic-flag and WITHIN GROUP forms fall back to Spark. Closes apache#3970
b6cf253 to
fb2f52b
Compare
andygrove
changed the title
mode
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
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.
|
|
||
| /// Normalize a scalar key so that Spark's floating-point normalization is honoured: `-0.0` and | ||
| /// `0.0` collapse to the same key and all `NaN` bit patterns collapse to a canonical `NaN`. | ||
| fn normalize_key(value: ScalarValue) -> ScalarValue { |
There was a problem hiding this comment.
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.updatekeys anOpenHashMap[AnyRef, Long]viabuffer.changeValue(InternalRow.copyValue(key), ...)with no normalization (Mode.scala:63). OpenHashMap/OpenHashSetkey comparison is_data(pos) equals k(core/.../util/collection/OpenHashSet.scala:122), and that file carries an explicit comment thatequalsdistinguishes0.0/-0.0and collapsesNaN/NaN(OpenHashSet.scala:116). So Spark treats-0.0and0.0as distinct keys and allNaNas one key.NormalizeFloatingNumbersdoes not touch aggregate arguments. Itsapplyonly rewritesWINDOWandJOINpatterns (sql/catalyst/.../optimizer/NormalizeFloatingNumbers.scala:62). There is noAggregatecase, so amode(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 grpKeep 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 grpIf 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.
| } | ||
|
|
||
| /// Add each non-null value in `array` to `map`, normalizing float keys. | ||
| fn count_values(map: &mut HashMap<ScalarValue, i64>, array: &ArrayRef, idx: usize) -> Result<()> { |
There was a problem hiding this comment.
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.
| eval_mode(&self.counts, &self.data_type) | ||
| } | ||
|
|
||
| fn size(&self) -> usize { |
There was a problem hiding this comment.
(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.
| Ok(self.data_type.clone()) | ||
| } | ||
|
|
||
| fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { |
There was a problem hiding this comment.
default_value duplicates the AggregateUDFImpl trait default (udaf.rs, both ScalarValue::try_from(&data_type)). Redundant, can be removed.
| for map in &emitted { | ||
| results.push(eval_mode(map, &self.data_type)?); | ||
| } | ||
| ScalarValue::iter_to_array(results) |
There was a problem hiding this comment.
(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.
| -- Config: spark.comet.expression.Mode.allowIncompatible=true | ||
|
|
||
| -- ============================================================ | ||
| -- Setup: tables |
There was a problem hiding this comment.
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).
| case _ => false | ||
| } | ||
|
|
||
| override def getSupportLevel(expr: Mode): SupportLevel = { |
There was a problem hiding this comment.
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.
2 participants
Which issue does this PR close?
Closes #3970.
Rationale for this change
The
modeaggregate (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-expressionproject skill.native/spark-expr/src/agg_funcs/mode.rs: aModeaggregate UDF with both a globalAccumulatorand a vectorizedGroupsAccumulator. State is a frequency map keyed byScalarValue, serialized as a singlestruct<values: array<T>, counts: array<bigint>>buffer column so partial/final buffer schemas stay aligned with Spark's single-attributeTypedImperativeAggregatebuffer.Modemessage andAggExproneof entry, plus the planner arm inplanner.rs.CometModeserde and registration inQueryPlanSerde.aggrSerdeMap.Modebranch inadjustOutputForNativeState(operators.scala) mapping the Spark binary buffer type to the native struct state type.modeHasUnsupportedOrderingshim inCometTypeShim(spark-3.x / spark-4.x) becauseMode.reverseOptonly exists on Spark 4.0+.modemarked supported in the expressions guide.Scope and compatibility:
mode(col)form is supported. Themode(col, deterministic)andmode() WITHIN GROUP (ORDER BY col)forms (Spark 4.0+, which setreverseOpt) fall back to Spark.Incompatible(opt-in viaspark.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.-0.0to0.0, canonicalNaN) 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?
mode.rscovering 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.mode.sqlfile 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.