andygrove commented on code in PR #6457: URL: https://github.com/apache/datafusion-comet/pull/6457#discussion_r4147896373
########## spark/src/test/resources/sql-tests/expressions/aggregate/min_max_floating_point.sql: ########## @@ -0,0 +1,73 @@ +-- 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. + +-- min and max order floats the way Spark's SQL ordering does: NaN is larger than every other +-- value, Infinity included, and -0.0 equals 0.0. Spark's Max keeps greatest(buffer, input), so +-- of equal values the first one seen is returned: max over -0.0 then 0.0 returns -0.0. Min is +-- the same with least. +-- +-- `-d` of the NaN rows is a NaN with the sign bit set on every platform. Arithmetic produces +-- that NaN on x86-64, and IEEE 754 total order sorts it below -Infinity. + +-- Strict floating-point mode no longer needs to fall back for min and max. +-- ConfigMatrix: spark.comet.exec.strictFloatingPoint=false,true + +statement +CREATE TABLE mm_float(id INT, g INT, d DOUBLE, f FLOAT) USING parquet + +statement +INSERT INTO mm_float VALUES + (1, 1, 1.0D, float('1.0')), + (2, 1, double('NaN'), float('NaN')), + (3, 1, -1.0D, float('-1.0')), + (4, 2, double('-0.0'), float('-0.0')), + (5, 2, 0.0D, float('0.0')), + (6, 3, 0.0D, float('0.0')), + (7, 3, double('-0.0'), float('-0.0')), + (8, 4, double('-Infinity'), float('-Infinity')), + (9, 5, double('Infinity'), float('Infinity')), + (10, 6, double('NaN'), float('NaN')), + (11, 7, NULL, NULL), + (12, 8, double('Infinity'), float('Infinity')), + (13, 8, double('NaN'), float('NaN')) + +-- Without grouping +query +SELECT max(d), min(d), max(-d), min(-d), max(f), min(f), max(-f), min(-f) FROM mm_float + +-- Of equal values the first one seen wins +query +SELECT max(d), min(d), max(-d), min(-d), max(f), min(f) FROM mm_float WHERE g IN (2, 3) + +-- Grouped. Groups 4 and 5 hold only an infinity, group 6 only NaN and group 7 only NULL. +query +SELECT g, max(d), min(d), max(-d), min(-d), max(f), min(f), max(-f), min(-f) +FROM mm_float GROUP BY g ORDER BY g + +-- Window frames: growing, sliding and the whole partition. Each frame is ordered by id, so the +-- first of two equal values is well defined. The sliding frames use FOLLOWING offsets because the +-- suite turns off constant folding, which leaves `1 PRECEDING` as an expression Spark's own window +-- operator cannot evaluate. +query +SELECT id, + max(d) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), + min(-d) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), + max(-d) OVER (ORDER BY id ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING), + min(d) OVER (ORDER BY id ROWS BETWEEN CURRENT ROW AND 2 FOLLOWING), Review Comment: The `f` windows all started at `UNBOUNDED PRECEDING`, so only the plain accumulator ran over `FLOAT`. The fixture now has `max(-f)` and `min(f)` over `CURRENT ROW AND 1 FOLLOWING` and `CURRENT ROW AND 2 FOLLOWING`, and the Rust sliding test runs every sequence as `FLOAT` as well. ########## native/spark-expr/src/agg_funcs/min_max.rs: ########## @@ -0,0 +1,738 @@ +// 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 crate::float_semantics::{float_gt, float_lt}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::NullBuffer; +use arrow::compute::{not, nullif, prep_null_mask_filter}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Field, FieldRef, Float32Type, Float64Type}; +use datafusion::common::{internal_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 num::Float; +use std::collections::VecDeque; +use std::fmt::Debug; +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; + +/// Spark's `max` or `min` over Float32 or Float64 values. +/// +/// Spark orders floats with `SQLOrderingUtil.compareDoubles`, in which NaN is larger than every +/// other value and `-0.0` equals `0.0`. Its `Max` updates the buffer with `greatest(max, input)`, +/// which replaces the buffer only with a strictly larger value, so of equal values the first one +/// seen is kept: `max` over `-0.0` and then `0.0` returns `-0.0`. `Min` does the same with +/// `least`. Every accumulator here follows that rule, including over a sliding window frame. +/// +/// DataFusion's `max` and `min` order floats by IEEE 754 total order in their batch kernels, in +/// which a NaN with the sign bit set is the smallest value, and the grouped versions start each +/// group at the most negative (or positive) finite value, so a group holding only `-Infinity` +/// returns that finite value. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMinMax { + signature: Signature, + is_max: bool, +} + +impl SparkMinMax { + pub fn new(is_max: bool) -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + is_max, + } + } +} + +impl AggregateUDFImpl for SparkMinMax { + fn name(&self) -> &str { + if self.is_max { + "max" + } else { + "min" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + match args.return_field.data_type() { + DataType::Float32 => Ok(Box::new(MinMaxAccumulator::<Float32Type>::new(self.is_max))), + DataType::Float64 => Ok(Box::new(MinMaxAccumulator::<Float64Type>::new(self.is_max))), + other => internal_err!("Spark {} expects a float, got {other}", self.name()), + } + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format_state_name(args.name, self.name()), + args.return_field.data_type().clone(), + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + match args.return_field.data_type() { + DataType::Float32 => Ok(Box::new(MinMaxGroupsAccumulator::<Float32Type>::new( + self.is_max, + ))), + DataType::Float64 => Ok(Box::new(MinMaxGroupsAccumulator::<Float64Type>::new( + self.is_max, + ))), + other => internal_err!("Spark {} expects a float, got {other}", self.name()), + } + } + + fn create_sliding_accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + match args.return_field.data_type() { + DataType::Float32 => Ok(Box::new(SlidingMinMaxAccumulator::<Float32Type>::new( + self.is_max, + ))), + DataType::Float64 => Ok(Box::new(SlidingMinMaxAccumulator::<Float64Type>::new( + self.is_max, + ))), + other => internal_err!("Spark {} expects a float, got {other}", self.name()), + } + } +} + +/// Whether `candidate` replaces a running maximum `current`: the same as [`float_gt`], so equal +/// values do not and the first one seen is kept. Written so that one comparison decides the usual +/// answer, which is no. That makes the grouped and sliding loops faster, which cannot vectorize, +/// while `float_gt` is faster in a loop over a slice that can. +// The negated comparison is the point: `!(candidate <= current)` also holds when either value +// is NaN, in one comparison. The `partial_cmp` form clippy suggests makes grouped `max` about 50% +// slower. +#[allow(clippy::neg_cmp_op_on_partial_ord)] +#[inline] +fn replaces_max<T: Float>(candidate: T, current: T) -> bool { + // Only a NaN `current` must be kept, since NaN is the largest value. + !(candidate <= current) && !current.is_nan() +} + +/// Whether `candidate` replaces a running minimum `current`: the same as [`float_lt`]. See +/// [`replaces_max`]. +#[allow(clippy::neg_cmp_op_on_partial_ord)] +#[inline] +fn replaces_min<T: Float>(candidate: T, current: T) -> bool { + // Only a NaN `candidate` must be passed over. + !(candidate >= current) && !candidate.is_nan() +} + +/// Whether `max` (with `is_max`) or `min` replaces `current` with `candidate`. +#[inline] +fn replaces<T: Float>(is_max: bool, candidate: T, current: T) -> bool { + if is_max { + replaces_max(candidate, current) + } else { + replaces_min(candidate, current) + } +} + +/// `max` or `min` without grouping. +#[derive(Debug)] +struct MinMaxAccumulator<T: ArrowPrimitiveType> { + value: Option<T::Native>, + is_max: bool, +} + +impl<T: ArrowPrimitiveType> MinMaxAccumulator<T> { + fn new(is_max: bool) -> Self { + Self { + value: None, + is_max, + } + } +} + +impl<T> Accumulator for MinMaxAccumulator<T> +where + T: ArrowPrimitiveType + Debug, + T::Native: Float, +{ + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let values = values[0].as_primitive::<T>(); + self.value = match values.nulls() { + None => fold(self.value, values.values().iter().copied(), self.is_max), + Some(_) => fold(self.value, values.iter().flatten(), self.is_max), + }; + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.update_batch(states) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![self.evaluate()?]) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + ScalarValue::new_primitive::<T>(self.value, &T::DATA_TYPE) + } + + fn size(&self) -> usize { + size_of_val(self) + } +} + +/// Folds `values` into `current`, keeping the first of equal values. +#[inline] +fn fold<T: Float>(current: Option<T>, values: impl Iterator<Item = T>, is_max: bool) -> Option<T> { + // Monomorphized per direction. + fn run<T: Float>( + current: Option<T>, + mut values: impl Iterator<Item = T>, + replaces: impl Fn(T, T) -> bool, + ) -> Option<T> { + let mut best = current.or_else(|| values.next())?; + for value in values { + if replaces(value, best) { + best = value; + } + } + Some(best) + } + // `float_gt` and `float_lt` rather than `replaces_max` and `replaces_min`: this loop over a + // slice vectorizes with them and is several times faster. + if is_max { + run(current, values, float_gt) + } else { + run(current, values, float_lt) + } +} + +/// `max` or `min` per group. A group keeps the first value it sees as is, and then only strictly +/// preferred ones. +#[derive(Debug)] +struct MinMaxGroupsAccumulator<T: ArrowPrimitiveType> { + values: Vec<T::Native>, + /// Whether each group has seen a value. A group that has not evaluates to null. + seen: Vec<bool>, + is_max: bool, +} + +impl<T: ArrowPrimitiveType> MinMaxGroupsAccumulator<T> +where + T::Native: Float, +{ + fn new(is_max: bool) -> Self { + Self { + values: vec![], + seen: vec![], + is_max, + } + } + + /// Monomorphized per direction, with a tight loop for the common case of no nulls and no + /// filter. + fn update<F>( + &mut self, + values: &PrimitiveArray<T>, + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + replaces: F, + ) where + F: Fn(T::Native, T::Native) -> bool, + { + let (values_by_group, seen_by_group) = (&mut self.values, &mut self.seen); + debug_assert!(group_indices + .iter() + .all(|&group| group < values_by_group.len())); + let mut accumulate = |group: usize, value: T::Native| { + // SAFETY: `update_batch` sized both vectors to `total_num_groups`, and every group + // index is below it, as `GroupsAccumulator` requires. DataFusion's own grouped `max` + // relies on the same guarantee; checking each index costs a third of this loop. + let (current, seen) = unsafe { + ( + values_by_group.get_unchecked_mut(group), + seen_by_group.get_unchecked_mut(group), + ) + }; + // `|` rather than `||`, so that the usual case takes no extra branch. + if !*seen | replaces(value, *current) { + *current = value; + } + *seen = true; + }; + if values.null_count() == 0 && opt_filter.is_none() { + for (&group, &value) in group_indices.iter().zip(values.values().iter()) { + accumulate(group, value); + } + return; + } + for (row, &group) in group_indices.iter().enumerate() { + let filtered_out = + opt_filter.is_some_and(|filter| !filter.is_valid(row) || !filter.value(row)); + if !filtered_out && values.is_valid(row) { + accumulate(group, values.value(row)); + } + } + } +} + +impl<T> GroupsAccumulator for MinMaxGroupsAccumulator<T> +where + T: ArrowPrimitiveType + Debug + Send + Sync, + T::Native: Float, +{ + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + // A new group takes its first value whatever it holds, since it has not been seen. + self.values.resize(total_num_groups, T::Native::default()); + self.seen.resize(total_num_groups, false); + let values = values[0].as_primitive::<T>(); + if self.is_max { + self.update(values, group_indices, opt_filter, replaces_max); + } else { + self.update(values, group_indices, opt_filter, replaces_min); + } + Ok(()) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + self.update_batch(values, group_indices, None, total_num_groups) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + let values = emit_to.take_needed(&mut self.values); + let seen = emit_to.take_needed(&mut self.seen); + Ok(Arc::new(PrimitiveArray::<T>::new( + values.into(), + Some(NullBuffer::from(seen)), + ))) + } + + fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + Ok(vec![self.evaluate(emit_to)?]) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result<Vec<ArrayRef>> { + let values = Arc::clone(&values[0]); + let Some(filter) = opt_filter else { + return Ok(vec![values]); + }; + // A row the filter drops, or whose filter value is null, contributes nothing. + let dropped = not(&prep_null_mask_filter(filter))?; + Ok(vec![nullif(&values, &dropped)?]) + } + + fn size(&self) -> usize { + size_of_val(self) + + self.values.capacity() * size_of::<T::Native>() + + self.seen.capacity() * size_of::<bool>() + } +} + +/// `max` or `min` over a sliding window frame, which retracts values in the order it added them. +/// +/// A deque holds the values that can still become the result, oldest first. Adding a value drops +/// the newer values that it is strictly preferred to, but not equal ones, so the front is always +/// the first of the best values in the frame. +#[derive(Debug)] +struct SlidingMinMaxAccumulator<T: ArrowPrimitiveType> { + /// Candidates, each with the position at which it was added. + candidates: VecDeque<(u64, T::Native)>, + added: u64, + retracted: u64, + is_max: bool, +} + +impl<T: ArrowPrimitiveType> SlidingMinMaxAccumulator<T> +where + T::Native: Float, +{ + fn new(is_max: bool) -> Self { + Self { + candidates: VecDeque::new(), + added: 0, + retracted: 0, + is_max, + } + } + + fn add(&mut self, value: T::Native) { + while self + .candidates + .back() + .is_some_and(|&(_, newest)| replaces(self.is_max, value, newest)) + { + self.candidates.pop_back(); + } + self.candidates.push_back((self.added, value)); + self.added += 1; + } + + fn retract_oldest(&mut self) { + if self.retracted == self.added { + return; + } + if self + .candidates + .front() + .is_some_and(|&(position, _)| position == self.retracted) + { + self.candidates.pop_front(); + } + self.retracted += 1; + } +} + +impl<T> Accumulator for SlidingMinMaxAccumulator<T> +where + T: ArrowPrimitiveType + Debug, + T::Native: Float, +{ + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + for value in values[0].as_primitive::<T>().iter().flatten() { + self.add(value); + } + Ok(()) + } + + fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + // Nulls were never added, so only the valid values leave the frame. + for _ in 0..values[0].len() - values[0].null_count() { + self.retract_oldest(); + } + Ok(()) + } + + fn supports_retract_batch(&self) -> bool { + true + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.update_batch(states) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![self.evaluate()?]) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + let value = self.candidates.front().map(|&(_, value)| value); + ScalarValue::new_primitive::<T>(value, &T::DATA_TYPE) + } + + fn size(&self) -> usize { + size_of_val(self) + self.candidates.capacity() * size_of::<(u64, T::Native)>() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::float_semantics::compare_floats; + use arrow::array::{Float32Array, Float64Array}; + + const EDGE_VALUES: [Option<f64>; 10] = [ + Some(f64::NEG_INFINITY), + Some(-1.0), + Some(-0.0), + Some(0.0), + Some(1.0), + Some(f64::INFINITY), + Some(f64::NAN), + // A NaN with the sign bit set, as arithmetic produces on x86-64. + Some(f64::from_bits(0xfff8_0000_0000_0000)), + // A NaN with a payload. + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + None, + ]; + + /// Pseudo-random sequences of edge values, the same on every run. + fn sequences() -> Vec<Vec<Option<f64>>> { + let mut state = 0x2545_f491_4f6c_dd1du64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + (0..300) + .map(|i| { + (0..i % 23) + .map(|_| EDGE_VALUES[(next() % EDGE_VALUES.len() as u64) as usize]) + .collect() + }) + .collect() + } + + /// Spark's `max` (or `min`): fold in order with `compareDoubles`, replacing only with a + /// strictly greater (or smaller) value. + fn spark_fold(values: &[Option<f64>], is_max: bool) -> Option<f64> { Review Comment: Done. `EDGE_VALUES` and `spark_extreme` are `#[cfg(test)] pub(crate)` in `float_semantics`, next to `NEGATIVE_NAN` and `PAYLOAD_NAN`, and both test modules use them, so `greatest`/`least` now covers `1.0` too. ########## native/core/src/execution/planner.rs: ########## @@ -3522,11 +3524,19 @@ impl PhysicalPlanner { } Some(AggExprStruct::Min(expr)) => { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; - Ok((by_name("min")?, vec![child])) + if is_float(&child.data_type(&schema)?) { Review Comment: Done. Both window arms call `min_max_udaf` and wrap the result in `WindowFunctionDefinition::AggregateUDF`. `by_name("min")` resolved to the same UDAF that `min_udaf()` returns, so windows over other types are unchanged. `is_float` is gone in favor of `DataType::is_floating()`. -- 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]
