andygrove commented on code in PR #4782: URL: https://github.com/apache/datafusion-comet/pull/4782#discussion_r3713602303
########## 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<()> { Review Comment: 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. ########## 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 { Review Comment: Fixed in both places. Extracted a `map_size` helper used by `ModeAccumulator::size` and `ModeGroupsAccumulator::size`: ```rust 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. ########## 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> { Review Comment: 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. -- 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]
