andygrove commented on code in PR #4782: URL: https://github.com/apache/datafusion-comet/pull/4782#discussion_r3713607462
########## 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 { + /// Collapse `-0.0`/`0.0` and every `NaN` to a canonical form for one float variant. + macro_rules! normalize_float { + ($variant:path, $f:expr, $nan:expr) => { + if $f == 0.0 { + $variant(Some(0.0)) + } else if $f.is_nan() { + $variant(Some($nan)) + } else { + $variant(Some($f)) + } + }; + } + match value { + ScalarValue::Float32(Some(f)) => normalize_float!(ScalarValue::Float32, f, f32::NAN), + ScalarValue::Float64(Some(f)) => normalize_float!(ScalarValue::Float64, f, f64::NAN), + other => other, + } +} + +/// 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<()> { + if array.is_null(idx) { + return Ok(()); + } + let key = normalize_key(ScalarValue::try_from_array(array, idx)?); + *map.entry(key).or_insert(0) += 1; + Ok(()) +} + +/// Fold row `row` of the struct-state columns (`{values, counts}`) into `map`. +fn merge_state_row( + map: &mut HashMap<ScalarValue, i64>, + values_list: &arrow::array::ListArray, + counts_list: &arrow::array::ListArray, + row: usize, +) -> Result<()> { + if values_list.is_null(row) { + return Ok(()); + } + let values = values_list.value(row); + let counts = counts_list.value(row); + let counts = counts + .as_primitive_opt::<Int64Type>() + .ok_or_else(|| internal_datafusion_err!("mode state counts must be Int64"))?; + for i in 0..values.len() { + if values.is_null(i) { + continue; + } + let key = normalize_key(ScalarValue::try_from_array(&values, i)?); + *map.entry(key).or_insert(0) += counts.value(i); + } + Ok(()) +} + +/// Pick the mode from a frequency map: the value with the highest count, breaking ties by the +/// smallest value. Returns a null scalar of `data_type` when the map is empty. +fn eval_mode(counts: &HashMap<ScalarValue, i64>, data_type: &DataType) -> Result<ScalarValue> { + let mut best: Option<(&ScalarValue, i64)> = None; + for (value, &count) in counts.iter() { + let wins = match best { + None => true, + Some((best_value, best_count)) => { + count > best_count + || (count == best_count + && value.partial_cmp(best_value) == Some(Ordering::Less)) + } + }; + if wins { + best = Some((value, count)); + } + } + match best { + Some((value, _)) => Ok(value.clone()), + None => ScalarValue::try_from(data_type), + } +} + +/// Non-grouped accumulator backing global `mode` aggregation. +#[derive(Debug)] +pub struct ModeAccumulator { + counts: HashMap<ScalarValue, i64>, + data_type: DataType, +} + +impl ModeAccumulator { + fn new(data_type: DataType) -> Self { + Self { + counts: HashMap::new(), + data_type, + } + } +} + +impl Accumulator for ModeAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let array = &values[0]; + for i in 0..array.len() { + count_values(&mut self.counts, array, i)?; + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let structs = states[0].as_struct(); + let values_list = structs.column(0).as_list::<i32>(); + let counts_list = structs.column(1).as_list::<i32>(); + for row in 0..structs.len() { + merge_state_row(&mut self.counts, values_list, counts_list, row)?; + } + Ok(()) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + let array = build_state(&self.data_type, &[&self.counts])?; + Ok(vec![ScalarValue::Struct(Arc::new(array))]) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + eval_mode(&self.counts, &self.data_type) + } + + fn size(&self) -> usize { + size_of_val(self) + self.counts.capacity() * size_of::<(ScalarValue, i64)>() + } +} + +/// Vectorized grouped accumulator: one frequency map per group. +#[derive(Debug)] +pub struct ModeGroupsAccumulator { + groups: Vec<HashMap<ScalarValue, i64>>, + data_type: DataType, +} + +impl ModeGroupsAccumulator { + fn new(data_type: DataType) -> Self { + Self { + groups: Vec::new(), + data_type, + } + } + + fn resize(&mut self, total_num_groups: usize) { + if self.groups.len() < total_num_groups { + self.groups.resize_with(total_num_groups, HashMap::new); + } + } +} + +impl GroupsAccumulator for ModeGroupsAccumulator { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.resize(total_num_groups); + let array = &values[0]; + for (idx, &group_index) in group_indices.iter().enumerate() { + if let Some(f) = opt_filter { + if !f.is_valid(idx) || !f.value(idx) { + continue; + } + } + count_values(&mut self.groups[group_index], array, idx)?; + } + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + _opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.resize(total_num_groups); + let structs = values[0].as_struct(); + let values_list = structs.column(0).as_list::<i32>(); + let counts_list = structs.column(1).as_list::<i32>(); + for (row, &group_index) in group_indices.iter().enumerate() { + merge_state_row(&mut self.groups[group_index], values_list, counts_list, row)?; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + let emitted = emit_to.take_needed(&mut self.groups); + let mut results = Vec::with_capacity(emitted.len()); + for map in &emitted { + results.push(eval_mode(map, &self.data_type)?); + } + ScalarValue::iter_to_array(results) Review Comment: Added `debug_assert!` at both grouped emit sites with a comment naming the dependency, so the invariant is stated rather than implied: ```rust 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. ########## spark/src/test/resources/sql-tests/expressions/aggregate/mode.sql: ########## @@ -0,0 +1,186 @@ +-- 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. + +-- Comet's `mode` is opt-in via allowIncompatible because Spark breaks ties non-deterministically. +-- Every compared query below has a single value with the strictly-highest frequency per group so +-- that Comet's smallest-value tie-break agrees with Spark's arbitrary choice. +-- Config: spark.comet.expression.Mode.allowIncompatible=true + +-- ============================================================ +-- Setup: tables Review Comment: 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. ########## spark/src/main/scala/org/apache/comet/serde/aggregates.scala: ########## @@ -823,6 +823,71 @@ object CometCollectSet extends CometAggregateExpressionSerde[CollectSet] { } } +object CometMode extends CometAggregateExpressionSerde[Mode] with CometTypeShim { + + private val tieBreakReason = + "mode breaks ties non-deterministically in Spark (the result depends on JVM hash-map" + + " iteration order); Comet returns the smallest of the tied values instead" + + " (https://github.com/apache/datafusion-comet/issues/3970)" + + override def getIncompatibleReasons(): Seq[String] = Seq(tieBreakReason) + + private def isSupportedType(dt: DataType): Boolean = dt match { + case BooleanType => true + case ByteType | ShortType | IntegerType | LongType => true + case FloatType | DoubleType => true + case _: DecimalType => true + case DateType | TimestampType | TimestampNTZType => true + case StringType => true + case _ => false + } + + override def getSupportLevel(expr: Mode): SupportLevel = { Review Comment: Added the TODO rather than the behaviour change, to keep this PR scoped: ```scala // 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. ########## 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: 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. -- 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]
