rich7420 commented on code in PR #4782: URL: https://github.com/apache/datafusion-comet/pull/4782#discussion_r3951344766
########## native/spark-expr/src/agg_funcs/mode.rs: ########## @@ -0,0 +1,646 @@ +/* + * 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, not_impl_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 +/// +/// Spark keys the frequency map on the boxed input value and compares keys with +/// `OpenHashSet`'s `_data(pos) equals k` (`core/.../util/collection/OpenHashSet.scala:122`), i.e. +/// `java.lang.Double.equals`, which is defined via `doubleToLongBits`. That collapses every `NaN` +/// bit pattern to one key but keeps `-0.0` and `0.0` apart. Note that +/// `NormalizeFloatingNumbers` does *not* apply here: its `apply` only rewrites `WINDOW` and +/// `JOIN` patterns, so an aggregate's argument reaches `Mode` un-normalized. +/// +/// Spark 4.2.0 changed this. SPARK-57329 ("mode() returns incorrect result when input contains +/// both -0.0 and 0.0") treats the split `-0.0`/`0.0` counts as a bug and normalizes the key at +/// update time, so from 4.2.0 on the two fold into a single key. `normalize_neg_zero` therefore +/// tracks the Spark version Comet is running against: it is `false` for Spark 3.4 through 4.1 and +/// `true` for 4.2.0+. `NaN` canonicalization is unconditional because every supported version +/// collapses `NaN` via `doubleToLongBits`. +/// +/// Do not "simplify" this to always normalize: `max_by`/`min_by` need the opposite treatment, +/// because they compare the ordering column with `SQLOrderingUtil.compareDoubles`, which ties +/// `-0.0 == 0.0` on every version. +/// +/// 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, + /// Whether `-0.0` folds into `0.0` before being used as a key (Spark 4.2.0+; SPARK-57329). + normalize_neg_zero: bool, +} + +impl Mode { + pub fn new(data_type: DataType, normalize_neg_zero: bool) -> Self { + Self { + name: "mode".to_string(), + signature: Signature::any(1, Volatility::Immutable), + data_type, + normalize_neg_zero, + } + } +} + +/// 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 accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(ModeAccumulator::new( + self.data_type.clone(), + self.normalize_neg_zero, + ))) + } + + 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(), + self.normalize_neg_zero, + ))) + } +} + +/// Canonicalize a float key so that map lookups reproduce Spark's key equality. +/// +/// `ScalarValue`'s `PartialEq`/`Hash` for `Float32`/`Float64` are both defined on `to_bits()`, so +/// distinct `NaN` bit patterns would otherwise be distinct keys and `-0.0` is naturally kept apart +/// from `0.0`. Collapsing `NaN` to one canonical value therefore reproduces `doubleToLongBits` +/// equality, which is what Spark's `OpenHashSet` uses. `-0.0` is folded into `0.0` only when +/// `normalize_neg_zero` is set, i.e. only on Spark 4.2.0+ (SPARK-57329); see [`Mode`]. +fn normalize_key(value: ScalarValue, normalize_neg_zero: bool) -> ScalarValue { + macro_rules! normalize_float { + ($variant:path, $f:expr, $nan:expr) => { + if $f.is_nan() { + $variant(Some($nan)) + } else if normalize_neg_zero && $f == 0.0 { + // `-0.0 == 0.0` in IEEE 754, so this catches negative zero only. + $variant(Some(0.0)) + } 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`, canonicalizing float keys. +/// +/// The map is intentionally keyed on the type-generic `ScalarValue` rather than a monomorphized +/// `HashMap<Hashable<T::Native>, _>`: `mode` supports every primitive type plus decimal, string and +/// the temporal types, so one generic map is simpler than a kernel per type. Revisit if the hot +/// primitive paths ever show up in a profile. +fn count_values( + map: &mut HashMap<ScalarValue, i64>, + array: &ArrayRef, + idx: usize, + normalize_neg_zero: bool, +) -> Result<()> { + if array.is_null(idx) { + return Ok(()); + } + let key = normalize_key(ScalarValue::try_from_array(array, idx)?, normalize_neg_zero); + *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, + normalize_neg_zero: bool, +) -> 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)?, normalize_neg_zero); + *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), + } +} + +/// Heap bytes held by the frequency map's keys, on top of the map's own slot allocation. +/// +/// `HashMap::capacity` only accounts for the inline `(ScalarValue, i64)` slots, which misses the +/// `String`/`Vec<u8>`/boxed-decimal payloads behind variable-length keys. Under-reporting those +/// would hide real memory from the pool that drives spill decisions. +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>() +} + +/// Non-grouped accumulator backing global `mode` aggregation. +#[derive(Debug)] +pub struct ModeAccumulator { + counts: HashMap<ScalarValue, i64>, + data_type: DataType, + normalize_neg_zero: bool, +} + +impl ModeAccumulator { + fn new(data_type: DataType, normalize_neg_zero: bool) -> Self { + Self { + counts: HashMap::new(), + data_type, + normalize_neg_zero, + } + } +} + +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, self.normalize_neg_zero)?; + } + 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, + self.normalize_neg_zero, + )?; + } + 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) + map_size(&self.counts) + } +} + +/// Vectorized grouped accumulator: one frequency map per group. +#[derive(Debug)] +pub struct ModeGroupsAccumulator { + groups: Vec<HashMap<ScalarValue, i64>>, + data_type: DataType, + normalize_neg_zero: bool, +} + +impl ModeGroupsAccumulator { + fn new(data_type: DataType, normalize_neg_zero: bool) -> Self { + Self { + groups: Vec::new(), + data_type, + normalize_neg_zero, + } + } + + 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, + self.normalize_neg_zero, + )?; + } + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + 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, + self.normalize_neg_zero, + )?; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + 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"); + 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) + } + + fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let emitted = emit_to.take_needed(&mut self.groups); + // As in `evaluate`: `build_state` funnels into `ScalarValue::iter_to_array`, which needs a + // non-empty iterator. + debug_assert!(!emitted.is_empty(), "mode: state called with no groups"); + let refs: Vec<&HashMap<ScalarValue, i64>> = emitted.iter().collect(); + Ok(vec![Arc::new(build_state(&self.data_type, &refs)?)]) + } + + fn convert_to_state( + &self, + _values: &[ArrayRef], + _opt_filter: Option<&BooleanArray>, + ) -> Result<Vec<ArrayRef>> { + not_impl_err!("Input batch conversion to state not implemented") + } + + fn size(&self) -> usize { + size_of_val(self) + self.groups.iter().map(map_size).sum::<usize>() Review Comment: 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. -- 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]
