mbutrovich commented on code in PR #4782: URL: https://github.com/apache/datafusion-comet/pull/4782#discussion_r3659285110
########## 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: Cross-linking a related finding, since the two are easy to conflate and the fixes point in opposite directions. #4817 (`max_by`/`min_by`) has the mirror-image bug on the same input. There the ordering column is compared with `SQLOrderingUtil.compareDoubles`, which is `if (x == y) 0 else Double.compare(x, y)` and explicitly documents `-0.0 == 0.0` (`SQLOrderingUtil.scala:27-31`), so Spark ties the two and Comet must collapse them. Here `mode` keys on `OpenHashSet`, whose `equals` distinguishes `-0.0` from `0.0`, so Comet must stop collapsing them. Same two input values, opposite correct behavior, because the two aggregates reach different Spark comparison paths. Whichever way each is fixed, please add a one-line comment at each site naming the Spark path that governs it (`OpenHashSet.equals` here, `SQLOrderingUtil.compareDoubles` there), so the next person does not "fix" one to match the other. -- 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]
