andygrove commented on code in PR #4782: URL: https://github.com/apache/datafusion-comet/pull/4782#discussion_r3713598755
########## native/spark-expr/src/agg_funcs/mode.rs: ########## @@ -0,0 +1,496 @@ +/* + * 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::{Array, ArrayRef, AsArray, BooleanArray, StructArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, Int64Type}; +use datafusion::common::{internal_datafusion_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::cmp::Ordering; +use std::collections::HashMap; +use std::mem::size_of; +use std::sync::Arc; + +/// Spark's `mode` aggregate: returns the most frequent value within a group, ignoring NULLs. +/// +/// Spark breaks ties on the default `mode(col)` form non-deterministically (the value is chosen +/// by JVM `OpenHashMap` iteration order), which a native hash map cannot reproduce bit-for-bit. +/// Comet resolves ties deterministically by returning the smallest value, so this function is +/// registered as `Incompatible` on the Scala side and is opt-in via `allowIncompatible`. +/// +/// Float keys are normalized before counting (`-0.0` becomes `0.0` and every `NaN` becomes a +/// canonical `NaN`) to match Spark's `NormalizeFloatingNumbers` behaviour so that counts agree. +/// +/// Spark's `Mode` is a `TypedImperativeAggregate` with a single aggregation-buffer attribute, so +/// the intermediate state is a single struct field `{ values: list<T>, counts: list<i64> }` (a +/// parallel-array encoding of the frequency map) to keep the partial/final buffer schemas aligned +/// with Spark. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Mode { + name: String, + signature: Signature, + data_type: DataType, +} + +impl Mode { + pub fn new(data_type: DataType) -> Self { + Self { + name: "mode".to_string(), + signature: Signature::any(1, Volatility::Immutable), + data_type, + } + } +} + +/// Fields of the single struct state column `{values: list<T>, counts: list<i64>}`. +fn state_struct_fields(data_type: &DataType) -> Fields { + let values_list = DataType::List(Arc::new(Field::new_list_field(data_type.clone(), true))); + let counts_list = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))); + Fields::from(vec![ + Field::new("values", values_list, false), + Field::new("counts", counts_list, false), + ]) +} + +/// Build the single-column struct state array holding one `{values, counts}` row per map. +fn build_state(data_type: &DataType, maps: &[&HashMap<ScalarValue, i64>]) -> Result<StructArray> { + let mut value_lists = Vec::with_capacity(maps.len()); + let mut count_lists = Vec::with_capacity(maps.len()); + for map in maps { + let mut values = Vec::with_capacity(map.len()); + let mut counts = Vec::with_capacity(map.len()); + for (value, &count) in map.iter() { + values.push(value.clone()); + counts.push(ScalarValue::Int64(Some(count))); + } + value_lists.push(ScalarValue::List(ScalarValue::new_list( + &values, data_type, true, + ))); + count_lists.push(ScalarValue::List(ScalarValue::new_list( + &counts, + &DataType::Int64, + true, + ))); + } + let values = ScalarValue::iter_to_array(value_lists)?; + let counts = ScalarValue::iter_to_array(count_lists)?; + Ok(StructArray::new( + state_struct_fields(data_type), + vec![values, counts], + None, + )) +} + +impl AggregateUDFImpl for Mode { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(self.data_type.clone()) + } + + fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> { + ScalarValue::try_from(&self.data_type) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(ModeAccumulator::new(self.data_type.clone()))) + } + + fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(&self.name, "freq"), + DataType::Struct(state_struct_fields(&self.data_type)), + false, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + Ok(Box::new(ModeGroupsAccumulator::new(self.data_type.clone()))) + } +} + +/// 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 { Review Comment: 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](https://issues.apache.org/jira/browse/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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
