From 7cc2160659f183e1630e2f4485a7c856dfe58776 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 29 Sep 2026 08:46:23 -0600 Subject: [PATCH 1/3] refactor: move native -0.0 and NaN handling into one float_semantics module Part of #6385. Each native expression that follows one of Spark's float rules carried its own copy of that rule. This moves them onto a new `float_semantics` module in `spark-expr`, with no change in behavior. - `normalize_float`, `canonicalize_nan`, `compare_floats` and `hash_input` are the per-value rules: `NormalizeNaNAndZero`, `Double.equals`, `compareDoubles`, and the `Murmur3Hash`/`XxHash64` input. - `normalize_floats`, `normalize_nested_floats` and `has_float_leaf` normalize arrays. `NormalizeNaNAndZero` and `NormalizeNestedFloats` move into the module, and `NormalizeNaNAndZero::wrap_if_needed` takes over the planner's check in `create_normalized_key_expr`. - `spark_comparator` replaces both `nested_equality` (nested `=` and `IN`) and the comparator local to `array_min`/`array_max`. hll_plus_plus, max_min_by, mode, percentile, array_extrema, array_position, arrays_overlap, nested_comparison, the float hash macros, and the planner's sort keys and range partition bounds now call the module. --- native/core/src/execution/planner.rs | 26 +- .../spark-expr/src/agg_funcs/hll_plus_plus.rs | 27 +- native/spark-expr/src/agg_funcs/max_min_by.rs | 83 ++--- native/spark-expr/src/agg_funcs/mode.rs | 23 +- native/spark-expr/src/agg_funcs/percentile.rs | 33 +- .../src/array_funcs/array_extrema.rs | 113 +------ .../src/array_funcs/array_position.rs | 6 +- .../src/array_funcs/arrays_overlap.rs | 14 +- native/spark-expr/src/array_funcs/mod.rs | 2 - .../src/array_funcs/nested_comparison.rs | 98 +----- .../spark-expr/src/float_semantics/compare.rs | 305 ++++++++++++++++++ native/spark-expr/src/float_semantics/mod.rs | 179 ++++++++++ .../normalize.rs} | 172 +++++++++- native/spark-expr/src/hash_funcs/utils.rs | 32 +- native/spark-expr/src/lib.rs | 4 +- .../spark-expr/src/math_funcs/internal/mod.rs | 3 - .../src/math_funcs/internal/normalize_nan.rs | 133 -------- 17 files changed, 744 insertions(+), 509 deletions(-) create mode 100644 native/spark-expr/src/float_semantics/compare.rs create mode 100644 native/spark-expr/src/float_semantics/mod.rs rename native/spark-expr/src/{array_funcs/nested_float_normalize.rs => float_semantics/normalize.rs} (78%) delete mode 100644 native/spark-expr/src/math_funcs/internal/normalize_nan.rs diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index cd7c9c6003e..5a8370ad8ac 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -149,11 +149,12 @@ use datafusion_comet_proto::{ spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; use datafusion_comet_spark_expr::{ - jvm_udf::JvmScalarUdfExpr, spark_in_list, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, - CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, - GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, HllUnionAgg, IfExpr, - ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, Regr, RegrType, SparkCastOptions, Stddev, - SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, + float_semantics::normalize_floats, jvm_udf::JvmScalarUdfExpr, spark_in_list, ApproxPercentile, + ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, + DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, + HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, Regr, RegrType, + SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, + WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -1015,15 +1016,10 @@ impl PhysicalPlanner { input_schema: SchemaRef, ) -> Result, ExecutionError> { let child = self.create_expr(spark_expr, Arc::clone(&input_schema))?; - let data_type = child.data_type(input_schema.as_ref())?; - // Spark may already have normalized a partition or join key. - if matches!(data_type, DataType::Float32 | DataType::Float64) - && child.downcast_ref::().is_none() - { - Ok(Arc::new(NormalizeNaNAndZero::new(data_type, child))) - } else { - Ok(child) - } + Ok(NormalizeNaNAndZero::wrap_if_needed( + child, + input_schema.as_ref(), + )?) } /// Only constant literals are supported as scan defaults. @@ -3705,7 +3701,7 @@ impl PhysicalPlanner { .iter() .map(|scalar_vec| { ScalarValue::iter_to_array(scalar_vec.iter().cloned()) - .map(|array| NormalizeNaNAndZero::normalize_array(&array)) + .map(|array| normalize_floats(&array)) }) .collect::, _>>()?; diff --git a/native/spark-expr/src/agg_funcs/hll_plus_plus.rs b/native/spark-expr/src/agg_funcs/hll_plus_plus.rs index 3cb39f0acd7..080127225b7 100644 --- a/native/spark-expr/src/agg_funcs/hll_plus_plus.rs +++ b/native/spark-expr/src/agg_funcs/hll_plus_plus.rs @@ -24,12 +24,10 @@ //! results are bit-identical to Spark. use crate::agg_funcs::hll_plus_plus_const::{BIAS_DATA, RAW_ESTIMATE_DATA, THRESHOLDS}; +use crate::float_semantics::normalize_floats; use crate::hash_funcs::create_xxhash64_hashes; -use crate::math_funcs::internal::normalize_float; -use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, Float32Array, Float64Array, Int64Array, -}; -use arrow::datatypes::{DataType, Field, FieldRef, Float32Type, Float64Type}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion::common::{not_impl_err, Result, ScalarValue}; use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion::logical_expr::{ @@ -120,25 +118,6 @@ impl AggregateUDFImpl for HllPlusPlus { } } -/// Normalize a float/double column the way Spark's `NormalizeNaNAndZero` does before hashing: -/// every NaN becomes the canonical NaN and `-0.0` becomes `0.0`. Returns the input unchanged for -/// non-floating-point types. -fn normalize_floats(array: &ArrayRef) -> ArrayRef { - match array.data_type() { - DataType::Float32 => { - let normalized: Float32Array = - array.as_primitive::().unary(normalize_float); - Arc::new(normalized) - } - DataType::Float64 => { - let normalized: Float64Array = - array.as_primitive::().unary(normalize_float); - Arc::new(normalized) - } - _ => Arc::clone(array), - } -} - /// Hash a value column with Spark's `XxHash64` (seed 42), normalizing floats first, reusing /// `buf` as scratch to avoid a per-batch allocation. The buffer is re-seeded (not just cleared) /// because `create_xxhash64_hashes` folds each value into the existing seed. diff --git a/native/spark-expr/src/agg_funcs/max_min_by.rs b/native/spark-expr/src/agg_funcs/max_min_by.rs index 147f1eb92bf..f0a82fee737 100644 --- a/native/spark-expr/src/agg_funcs/max_min_by.rs +++ b/native/spark-expr/src/agg_funcs/max_min_by.rs @@ -15,9 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{new_null_array, Array, ArrayRef, AsArray, BooleanArray}; +use crate::float_semantics::normalize_floats; +use arrow::array::{new_null_array, Array, ArrayRef, BooleanArray}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, FieldRef, Float32Type, Float64Type}; +use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::row::{OwnedRow, RowConverter, SortField}; use datafusion::common::{not_impl_err, Result, ScalarValue}; use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; @@ -39,6 +40,30 @@ use std::sync::Arc; /// buffer and, on a tie in the ordering, the later row wins. Because ties across /// partitions are processed in an unspecified order, Spark documents the function as /// non-deterministic when several rows share the extremum ordering. +/// +/// # Float orderings +/// +/// A floating-point ordering column is canonicalized with [`normalize_floats`] so that Arrow's +/// row-format byte order reproduces Spark's comparison for this aggregate. +/// +/// Spark compares the ordering with `SQLOrderingUtil.compareDoubles`/`compareFloats`, wired in via +/// `PhysicalDoubleType.ordering`/`PhysicalFloatType.ordering`. That is +/// `if (x == y) 0 else java.lang.Double.compare(x, y)`, which has two consequences Arrow's row +/// format does not share: +/// +/// * `-0.0` and `0.0` tie, because the `x == y` short-circuit is IEEE equality. Arrow encodes +/// floats by flipping the bits off the sign, a total order placing `-0.0` strictly below `0.0`. +/// * every `NaN` is one value and sorts above `+Infinity`, because `Double.compare` goes through +/// `doubleToLongBits`. Arrow uses the raw bits, so a sign-bit-set `NaN` would sort below +/// `-Infinity` instead. +/// +/// Folding `-0.0` into `0.0` and every `NaN` into the canonical `NaN` makes the row bytes agree +/// with `compareDoubles` on both counts. This is verified identical on Spark 3.4 through master. +/// +/// Note this is the opposite of what `mode` needs: `mode` keys a hash map via +/// `OpenHashSet`'s `equals` (`java.lang.Double.equals`), which distinguishes `-0.0` from `0.0`, so +/// it must *not* fold them. Same two input values, different Spark comparison path, opposite +/// correct behaviour. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaxMinBy { name: String, @@ -145,56 +170,6 @@ fn extremum_sort_options(is_max: bool) -> SortOptions { } } -/// Canonicalize a floating-point ordering column so that Arrow's row-format byte order reproduces -/// Spark's comparison for this aggregate. -/// -/// Spark compares the ordering with `SQLOrderingUtil.compareDoubles`/`compareFloats`, wired in via -/// `PhysicalDoubleType.ordering`/`PhysicalFloatType.ordering`. That is -/// `if (x == y) 0 else java.lang.Double.compare(x, y)`, which has two consequences Arrow's row -/// format does not share: -/// -/// * `-0.0` and `0.0` tie, because the `x == y` short-circuit is IEEE equality. Arrow encodes -/// floats by flipping the bits off the sign, a total order placing `-0.0` strictly below `0.0`. -/// * every `NaN` is one value and sorts above `+Infinity`, because `Double.compare` goes through -/// `doubleToLongBits`. Arrow uses the raw bits, so a sign-bit-set `NaN` would sort below -/// `-Infinity` instead. -/// -/// Folding `-0.0` into `0.0` and every `NaN` into the canonical `NaN` makes the row bytes agree -/// with `compareDoubles` on both counts. This is verified identical on Spark 3.4 through master. -/// -/// Note this is the opposite of what `mode` needs: `mode` keys a hash map via -/// `OpenHashSet`'s `equals` (`java.lang.Double.equals`), which distinguishes `-0.0` from `0.0`, so -/// it must *not* fold them. Same two input values, different Spark comparison path, opposite -/// correct behaviour. -fn canonicalize_float_ordering(array: &ArrayRef) -> ArrayRef { - match array.data_type() { - DataType::Float32 => Arc::new(array.as_primitive::().unary::<_, Float32Type>( - |v| { - if v.is_nan() { - f32::NAN - } else if v == 0.0 { - // `-0.0 == 0.0` in IEEE 754, so this catches negative zero only. - 0.0 - } else { - v - } - }, - )), - DataType::Float64 => Arc::new(array.as_primitive::().unary::<_, Float64Type>( - |v| { - if v.is_nan() { - f64::NAN - } else if v == 0.0 { - 0.0 - } else { - v - } - }, - )), - _ => Arc::clone(array), - } -} - /// Accumulator that tracks the running `(value, ordering)` pair for the extremum ordering. #[derive(Debug)] struct MaxMinByAccumulator { @@ -231,7 +206,7 @@ impl MaxMinByAccumulator { return Ok(()); } - let ordering_arr = canonicalize_float_ordering(ordering_arr); + let ordering_arr = normalize_floats(ordering_arr); let rows = self .ordering_converter .convert_columns(&[Arc::clone(&ordering_arr)])?; @@ -389,7 +364,7 @@ impl MaxMinByGroupsAccumulator { let value_rows = self .value_converter .convert_columns(&[Arc::clone(&values[0])])?; - let ordering_arr = canonicalize_float_ordering(&values[1]); + let ordering_arr = normalize_floats(&values[1]); let ordering_rows = self .ordering_converter .convert_columns(&[Arc::clone(&ordering_arr)])?; diff --git a/native/spark-expr/src/agg_funcs/mode.rs b/native/spark-expr/src/agg_funcs/mode.rs index 7cba965329c..cc0f9529f0d 100644 --- a/native/spark-expr/src/agg_funcs/mode.rs +++ b/native/spark-expr/src/agg_funcs/mode.rs @@ -17,6 +17,7 @@ * under the License. */ +use crate::float_semantics::{canonicalize_nan, normalize_float}; 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}; @@ -171,21 +172,15 @@ impl AggregateUDFImpl for Mode { /// 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), + ScalarValue::Float32(Some(f)) if normalize_neg_zero => { + ScalarValue::Float32(Some(normalize_float(f))) + } + ScalarValue::Float64(Some(f)) if normalize_neg_zero => { + ScalarValue::Float64(Some(normalize_float(f))) + } + ScalarValue::Float32(Some(f)) => ScalarValue::Float32(Some(canonicalize_nan(f))), + ScalarValue::Float64(Some(f)) => ScalarValue::Float64(Some(canonicalize_nan(f))), other => other, } } diff --git a/native/spark-expr/src/agg_funcs/percentile.rs b/native/spark-expr/src/agg_funcs/percentile.rs index 78e585a3965..a5484312a33 100644 --- a/native/spark-expr/src/agg_funcs/percentile.rs +++ b/native/spark-expr/src/agg_funcs/percentile.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::float_semantics::compare_floats; use arrow::array::{ Array, ArrayRef, AsArray, BooleanArray, Float64Array, Float64Builder, ListArray, }; @@ -297,38 +298,24 @@ fn spark_percentile(values: &mut [f64], percentile: f64) -> Option { let lower = position.floor() as usize; let higher = position.ceil() as usize; - let (_, lower_value, _) = values.select_nth_unstable_by(lower, spark_double_cmp); + let (_, lower_value, _) = values.select_nth_unstable_by(lower, |x, y| compare_floats(*x, *y)); let lower_value = *lower_value; if lower == higher { return Some(lower_value); } - let (_, higher_value, _) = values.select_nth_unstable_by(higher, spark_double_cmp); + let (_, higher_value, _) = values.select_nth_unstable_by(higher, |x, y| compare_floats(*x, *y)); let higher_value = *higher_value; - if spark_double_cmp(&lower_value, &higher_value) == Ordering::Equal { + if compare_floats(lower_value, higher_value) == Ordering::Equal { return Some(lower_value); } Some((higher as f64 - position) * lower_value + (position - lower as f64) * higher_value) } -fn spark_double_cmp(x: &f64, y: &f64) -> Ordering { - if x == y || (x.is_nan() && y.is_nan()) { - Ordering::Equal - } else if x.is_nan() { - Ordering::Greater - } else if y.is_nan() { - Ordering::Less - } else { - x.partial_cmp(y) - .expect("non-NaN values should be comparable") - } -} - #[cfg(test)] mod tests { - use super::{spark_double_cmp, spark_percentile}; - use std::cmp::Ordering; + use super::spark_percentile; #[test] fn interpolates_with_full_spark_precision() { @@ -340,10 +327,10 @@ mod tests { } #[test] - fn matches_spark_double_ordering_for_nan_and_zero() { - assert_eq!(spark_double_cmp(&f64::NAN, &1.0), Ordering::Greater); - assert_eq!(spark_double_cmp(&1.0, &f64::NAN), Ordering::Less); - assert_eq!(spark_double_cmp(&f64::NAN, &f64::NAN), Ordering::Equal); - assert_eq!(spark_double_cmp(&-0.0, &0.0), Ordering::Equal); + fn orders_nan_above_every_other_value() { + // A NaN with the sign bit set, which IEEE 754 total order would put first. + let mut values = vec![f64::from_bits(0xfff8_0000_0000_0000), 2.0, 1.0]; + assert_eq!(spark_percentile(&mut values, 0.0), Some(1.0)); + assert!(spark_percentile(&mut values, 1.0).unwrap().is_nan()); } } diff --git a/native/spark-expr/src/array_funcs/array_extrema.rs b/native/spark-expr/src/array_funcs/array_extrema.rs index 5fcdcc9dcfe..df048eae972 100644 --- a/native/spark-expr/src/array_funcs/array_extrema.rs +++ b/native/spark-expr/src/array_funcs/array_extrema.rs @@ -18,13 +18,13 @@ use std::cmp::Ordering; use std::sync::Arc; +use crate::float_semantics::{compare_floats, spark_comparator}; use arrow::array::{ - make_array, make_comparator, new_empty_array, Array, ArrayRef, AsArray, BooleanArray, - DynComparator, ListArray, MutableArrayData, PrimitiveArray, PrimitiveBuilder, StructArray, - UInt32Array, + make_array, new_empty_array, Array, ArrayRef, AsArray, BooleanArray, ListArray, + MutableArrayData, PrimitiveArray, PrimitiveBuilder, StructArray, UInt32Array, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::compute::{take, SortOptions}; +use arrow::compute::take; use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; use datafusion::common::{exec_err, Result, ScalarValue}; use datafusion::functions_nested::min_max::{array_max_udf, array_min_udf}; @@ -130,6 +130,11 @@ where let values = array.values().as_primitive::(); let buffer = values.values(); let nulls = values.nulls(); + let ordering = if is_min { + Ordering::Less + } else { + Ordering::Greater + }; let mut result = PrimitiveBuilder::::with_capacity(array.len()); for (row, offsets) in array.offsets().windows(2).enumerate() { let mut best: Option = None; @@ -140,15 +145,8 @@ where if nulls.is_some_and(|nulls| nulls.is_null(start + index)) { continue; } - let replace = match best { - None => true, - Some(current) if is_min => { - candidate < current || (!candidate.is_nan() && current.is_nan()) - } - Some(current) => { - candidate > current || (candidate.is_nan() && !current.is_nan()) - } - }; + let replace = + best.is_none_or(|current| compare_floats(candidate, current) == ordering); if replace { // Copy the winning value, never normalize its zero sign or NaN bits. best = Some(candidate); @@ -162,7 +160,7 @@ where fn nested_extrema(array: &ListArray, is_min: bool) -> Result { let values = array.values(); - let compare = spark_comparator(values)?; + let compare = spark_comparator(values.as_ref(), values.as_ref())?; let nulls = values.nulls(); let ordering = if is_min { Ordering::Less @@ -297,92 +295,5 @@ fn take_extrema_values(values: &ArrayRef, indices: &UInt32Array) -> Result Result { - match array.data_type() { - DataType::Float32 => Ok(float_comparator::(array)), - DataType::Float64 => Ok(float_comparator::(array)), - DataType::List(_) => { - let array = array.as_list::(); - let compare = spark_comparator(array.values())?; - let offsets = array.offsets().clone(); - Ok(nulls_first(array.nulls().cloned(), move |left, right| { - let left_start = offsets[left] as usize; - let right_start = offsets[right] as usize; - let left_len = offsets[left + 1] as usize - left_start; - let right_len = offsets[right + 1] as usize - right_start; - for offset in 0..left_len.min(right_len) { - let ordering = compare(left_start + offset, right_start + offset); - if ordering != Ordering::Equal { - return ordering; - } - } - left_len.cmp(&right_len) - })) - } - DataType::Struct(_) => { - let array = array.as_struct(); - let fields = array - .columns() - .iter() - .map(spark_comparator) - .collect::>>()?; - Ok(nulls_first(array.nulls().cloned(), move |left, right| { - fields - .iter() - .map(|compare| compare(left, right)) - .find(|&ordering| ordering != Ordering::Equal) - .unwrap_or(Ordering::Equal) - })) - } - _ => Ok(make_comparator( - array.as_ref(), - array.as_ref(), - SortOptions { - descending: false, - nulls_first: true, - }, - )?), - } -} - -fn float_comparator(array: &ArrayRef) -> DynComparator -where - T::Native: Float, -{ - let values = array.as_primitive::().values().clone(); - nulls_first(array.nulls().cloned(), move |left, right| { - let left = values[left]; - let right = values[right]; - if left == right || (left.is_nan() && right.is_nan()) { - Ordering::Equal - } else if left > right || left.is_nan() { - Ordering::Greater - } else { - Ordering::Less - } - }) -} - -fn nulls_first( - nulls: Option, - compare: impl Fn(usize, usize) -> Ordering + Send + Sync + 'static, -) -> DynComparator { - match nulls { - None => Box::new(compare), - Some(nulls) => { - Box::new( - move |left, right| match (nulls.is_null(left), nulls.is_null(right)) { - (true, true) => Ordering::Equal, - (true, false) => Ordering::Less, - (false, true) => Ordering::Greater, - (false, false) => compare(left, right), - }, - ) - } - } -} - #[cfg(test)] mod tests; diff --git a/native/spark-expr/src/array_funcs/array_position.rs b/native/spark-expr/src/array_funcs/array_position.rs index 060501c5f5b..b5f7d74c2ac 100644 --- a/native/spark-expr/src/array_funcs/array_position.rs +++ b/native/spark-expr/src/array_funcs/array_position.rs @@ -29,11 +29,10 @@ use datafusion::common::{exec_err, DataFusionError, Result as DataFusionResult, use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; -use num::Float; use std::cmp::Ordering; use std::sync::Arc; -use super::nested_float_normalize::{has_float_leaf, normalize_nested_floats}; +use crate::float_semantics::{compare_floats, has_float_leaf, normalize_nested_floats}; /// Spark array_position() function that returns the 1-based position of an element in an array. /// Returns 0 if the element is not found (Spark behavior differs from DataFusion which returns null). @@ -182,11 +181,10 @@ where let start = w[0].as_usize(); let end = w[1].as_usize(); let search_val = element_typed.value(row_index); - let search_is_nan = search_val.is_nan(); for i in start..end { if !values_typed.is_null(i) { let v = values_typed.value(i); - if (search_is_nan && v.is_nan()) || v == search_val { + if compare_floats(v, search_val).is_eq() { result[row_index] = (i - start + 1) as i64; break; } diff --git a/native/spark-expr/src/array_funcs/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index 542687c7d16..13af6636dab 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -51,7 +51,7 @@ use std::hash::Hash; use std::ops::Range; use std::sync::Arc; -use super::nested_float_normalize::{has_float_leaf, normalize_nested_floats}; +use crate::float_semantics::{canonicalize_nan, has_float_leaf, normalize_nested_floats}; #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkArraysOverlap { @@ -300,11 +300,7 @@ impl OverlapKey for f32 { type Key = u32; fn overlap_key(self) -> u32 { - if self.is_nan() { - f32::NAN.to_bits() - } else { - self.to_bits() - } + canonicalize_nan(self).to_bits() } } @@ -312,11 +308,7 @@ impl OverlapKey for f64 { type Key = u64; fn overlap_key(self) -> u64 { - if self.is_nan() { - f64::NAN.to_bits() - } else { - self.to_bits() - } + canonicalize_nan(self).to_bits() } } diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 77208a20035..4c234163ba8 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -25,7 +25,6 @@ mod flatten; mod get_array_struct_fields; mod list_extract; mod nested_comparison; -mod nested_float_normalize; mod sequence; mod size; @@ -39,6 +38,5 @@ pub use flatten::SparkFlatten; pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; pub use nested_comparison::{spark_comparison, spark_in_list}; -pub use nested_float_normalize::NormalizeNestedFloats; pub use sequence::spark_sequence; pub use size::{spark_size, SparkSizeFunc}; diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index bc2226aa2bb..8f50d7cf72d 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -17,13 +17,13 @@ //! Spark equality for nested floating-point values without materializing normalized columns. -use super::nested_float_normalize::{ - has_float_leaf, normalize_nested_floats, NormalizeNestedFloats, +use crate::float_semantics::{ + has_float_leaf, normalize_nested_floats, spark_comparator, NormalizeNestedFloats, }; -use arrow::array::{make_comparator, Array, ArrayRef, AsArray, BooleanArray, OffsetSizeTrait}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::{not, or_kleene, SortOptions}; -use arrow::datatypes::{DataType, Float32Type, Float64Type, Schema}; +use arrow::compute::{not, or_kleene}; +use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; use datafusion::common::{internal_err, DFSchema, Result, ScalarValue}; use datafusion::logical_expr::{ColumnarValue, Operator}; @@ -34,90 +34,10 @@ use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; -type Equality = Box bool + Send + Sync>; - fn needs_spark_equality(dt: &DataType) -> bool { dt.is_nested() && has_float_leaf(dt) } -// Inner nulls participate in structural equality. Outer SQL nulls are handled separately. -fn with_nulls(left: &ArrayRef, right: &ArrayRef, equal: Equality) -> Equality { - if left.null_count() == 0 && right.null_count() == 0 { - return equal; - } - let left = left.nulls().cloned(); - let right = right.nulls().cloned(); - Box::new(move |i, j| { - let l = left.as_ref().is_some_and(|n| n.is_null(i)); - let r = right.as_ref().is_some_and(|n| n.is_null(j)); - if l || r { - l == r - } else { - equal(i, j) - } - }) -} - -fn list_equality(left: &ArrayRef, right: &ArrayRef) -> Result { - let l = left.as_list::(); - let r = right.as_list::(); - let equal = nested_equality(l.values(), r.values())?; - let l = l.offsets().clone(); - let r = r.offsets().clone(); - Ok(Box::new(move |i, j| { - let (ls, le) = (l[i].as_usize(), l[i + 1].as_usize()); - let (rs, re) = (r[j].as_usize(), r[j + 1].as_usize()); - le - ls == re - rs && (ls..le).zip(rs..re).all(|(i, j)| equal(i, j)) - })) -} - -fn nested_equality(left: &ArrayRef, right: &ArrayRef) -> Result { - if !DFSchema::datatype_is_logically_equal(left.data_type(), right.data_type()) { - return internal_err!( - "Nested equality requires matching types, got {} and {}", - left.data_type(), - right.data_type() - ); - } - if !has_float_leaf(left.data_type()) { - let cmp = make_comparator(left.as_ref(), right.as_ref(), SortOptions::default())?; - return Ok(Box::new(move |i, j| cmp(i, j).is_eq())); - } - let equal: Equality = match left.data_type() { - DataType::Float32 => { - let l = left.as_primitive::().values().clone(); - let r = right.as_primitive::().values().clone(); - Box::new(move |i, j| l[i] == r[j] || (l[i].is_nan() && r[j].is_nan())) - } - DataType::Float64 => { - let l = left.as_primitive::().values().clone(); - let r = right.as_primitive::().values().clone(); - Box::new(move |i, j| l[i] == r[j] || (l[i].is_nan() && r[j].is_nan())) - } - DataType::List(_) => list_equality::(left, right)?, - DataType::LargeList(_) => list_equality::(left, right)?, - DataType::FixedSizeList(_, width) => { - let l = left.as_fixed_size_list(); - let r = right.as_fixed_size_list(); - let equal = nested_equality(l.values(), r.values())?; - let width = *width as usize; - Box::new(move |i, j| (0..width).all(|k| equal(i * width + k, j * width + k))) - } - DataType::Struct(_) => { - let equal = left - .as_struct() - .columns() - .iter() - .zip(right.as_struct().columns()) - .map(|(l, r)| nested_equality(l, r)) - .collect::>>()?; - Box::new(move |i, j| equal.iter().all(|eq| eq(i, j))) - } - _ => return internal_err!("Unsupported nested equality type {}", left.data_type()), - }; - Ok(with_nulls(left, right, equal)) -} - struct Operand { array: ArrayRef, scalar: bool, @@ -149,7 +69,8 @@ impl Operand { { BooleanArray::new_null(len) } else { - let equal = nested_equality(&self.array, &other.array)?; + // Inner nulls take part in structural equality. Outer SQL nulls are handled here. + let compare = spark_comparator(self.array.as_ref(), other.array.as_ref())?; let nulls = match (self.scalar, other.scalar) { (true, true) => None, (true, false) => other.array.nulls().cloned(), @@ -157,10 +78,11 @@ impl Operand { (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), }; let values = BooleanBuffer::collect_bool(len, |row| { - equal( + compare( if self.scalar { 0 } else { row }, if other.scalar { 0 } else { row }, ) + .is_eq() }); BooleanArray::new(values, nulls) }; @@ -400,7 +322,7 @@ mod tests { FixedSizeListArray, Float64Array, Int32Array, LargeListArray, ListArray, StructArray, }; use arrow::buffer::OffsetBuffer; - use arrow::datatypes::Field; + use arrow::datatypes::{Field, Float32Type, Float64Type}; use datafusion::physical_expr::expressions::{Column, Literal}; use std::collections::hash_map::DefaultHasher; diff --git a/native/spark-expr/src/float_semantics/compare.rs b/native/spark-expr/src/float_semantics/compare.rs new file mode 100644 index 00000000000..ebfb4d533c6 --- /dev/null +++ b/native/spark-expr/src/float_semantics/compare.rs @@ -0,0 +1,305 @@ +// 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 super::{compare_floats, has_float_leaf}; +use arrow::array::{make_comparator, Array, AsArray, DynComparator, OffsetSizeTrait}; +use arrow::compute::SortOptions; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use datafusion::common::{internal_err, DFSchema, Result}; +use num::Float; +use std::cmp::Ordering; +use std::ops::Range; + +/// Builds a comparator of `left[i]` against `right[j]` in Spark's SQL ordering, in which floats +/// compare as [`compare_floats`] does at any depth. Inner nulls sort first, lists compare element +/// by element and then by length, and structs compare field by field. +/// +/// Subtrees without a float leaf use Arrow's comparator, which orders them the same way. The two +/// arrays must have the same type, ignoring field names and nullability. +pub fn spark_comparator(left: &dyn Array, right: &dyn Array) -> Result { + if !DFSchema::datatype_is_logically_equal(left.data_type(), right.data_type()) { + return internal_err!( + "Spark comparison requires matching types, got {} and {}", + left.data_type(), + right.data_type() + ); + } + comparator(left, right) +} + +fn comparator(left: &dyn Array, right: &dyn Array) -> Result { + if !has_float_leaf(left.data_type()) { + let options = SortOptions { + descending: false, + nulls_first: true, + }; + return Ok(make_comparator(left, right, options)?); + } + match left.data_type() { + DataType::Float32 => Ok(float_comparator::(left, right)), + DataType::Float64 => Ok(float_comparator::(left, right)), + DataType::List(_) => list_comparator::(left, right), + DataType::LargeList(_) => list_comparator::(left, right), + DataType::FixedSizeList(_, _) => fixed_size_list_comparator(left, right), + DataType::Struct(_) => struct_comparator(left, right), + dt => internal_err!("Unsupported type for Spark comparison: {dt}"), + } +} + +fn float_comparator(left: &dyn Array, right: &dyn Array) -> DynComparator +where + T::Native: Float, +{ + let l = left.as_primitive::().values().clone(); + let r = right.as_primitive::().values().clone(); + nulls_first(left, right, move |i, j| compare_floats(l[i], r[j])) +} + +fn list_comparator( + left: &dyn Array, + right: &dyn Array, +) -> Result { + let (l, r) = (left.as_list::(), right.as_list::()); + let compare = comparator(l.values().as_ref(), r.values().as_ref())?; + let (l, r) = (l.offsets().clone(), r.offsets().clone()); + Ok(nulls_first(left, right, move |i, j| { + let left = l[i].as_usize()..l[i + 1].as_usize(); + let right = r[j].as_usize()..r[j + 1].as_usize(); + lexicographic(&compare, left, right) + })) +} + +fn fixed_size_list_comparator(left: &dyn Array, right: &dyn Array) -> Result { + let (l, r) = (left.as_fixed_size_list(), right.as_fixed_size_list()); + let compare = comparator(l.values().as_ref(), r.values().as_ref())?; + let (l, r) = (l.value_length() as usize, r.value_length() as usize); + Ok(nulls_first(left, right, move |i, j| { + lexicographic(&compare, i * l..(i + 1) * l, j * r..(j + 1) * r) + })) +} + +/// Compares two runs of child values element by element, then by length. +fn lexicographic(compare: &DynComparator, left: Range, right: Range) -> Ordering { + let lengths = left.len().cmp(&right.len()); + left.zip(right) + .map(|(i, j)| compare(i, j)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(lengths) +} + +fn struct_comparator(left: &dyn Array, right: &dyn Array) -> Result { + let fields = left + .as_struct() + .columns() + .iter() + .zip(right.as_struct().columns()) + .map(|(l, r)| comparator(l.as_ref(), r.as_ref())) + .collect::>>()?; + Ok(nulls_first(left, right, move |i, j| { + fields + .iter() + .map(|compare| compare(i, j)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(Ordering::Equal) + })) +} + +/// Orders nulls before values. A null slot never reaches `compare`, because the child values +/// under a null list or struct can be anything. +fn nulls_first( + left: &dyn Array, + right: &dyn Array, + compare: impl Fn(usize, usize) -> Ordering + Send + Sync + 'static, +) -> DynComparator { + let left = left.nulls().filter(|nulls| nulls.null_count() > 0).cloned(); + let right = right + .nulls() + .filter(|nulls| nulls.null_count() > 0) + .cloned(); + if left.is_none() && right.is_none() { + return Box::new(compare); + } + Box::new(move |i, j| { + let l = left.as_ref().is_some_and(|nulls| nulls.is_null(i)); + let r = right.as_ref().is_some_and(|nulls| nulls.is_null(j)); + match (l, r) { + (false, false) => compare(i, j), + (true, true) => Ordering::Equal, + (true, false) => Ordering::Less, + (false, true) => Ordering::Greater, + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, FixedSizeListArray, Float64Array, Int32Array, LargeListArray, ListArray, + StructArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::Field; + use std::sync::Arc; + + /// Wraps every value in a one-element list, large list and fixed-size list, and in a struct, + /// so that each nesting compares the same pairs of leaves. + fn nestings(values: ArrayRef) -> Vec { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let lengths = || std::iter::repeat_n(1, values.len()); + vec![ + Arc::clone(&values), + Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(lengths()), + Arc::clone(&values), + None, + )), + Arc::new(LargeListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(lengths()), + Arc::clone(&values), + None, + )), + Arc::new(FixedSizeListArray::new( + Arc::clone(&field), + 1, + Arc::clone(&values), + None, + )), + Arc::new(StructArray::new( + vec![field].into(), + vec![Arc::clone(&values)], + None, + )), + ] + } + + #[test] + fn floats_compare_as_spark_orders_them_at_every_depth() -> Result<()> { + let left = vec![ + Some(-0.0), + Some(f64::from_bits(0xfff8_0000_0000_0000)), + None, + Some(1.0), + Some(f64::INFINITY), + ]; + let right = vec![ + Some(0.0), + Some(f64::INFINITY), + None, + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + Some(f64::NEG_INFINITY), + Some(-1.0), + ]; + let expected = |l: Option, r: Option| match (l, r) { + (Some(l), Some(r)) => compare_floats(l, r), + (None, None) => Ordering::Equal, + (None, Some(_)) => Ordering::Less, + (Some(_), None) => Ordering::Greater, + }; + let nested_left = nestings(Arc::new(Float64Array::from(left.clone()))); + let nested_right = nestings(Arc::new(Float64Array::from(right.clone()))); + for (l_array, r_array) in nested_left.iter().zip(&nested_right) { + let compare = spark_comparator(l_array.as_ref(), r_array.as_ref())?; + for (i, &l) in left.iter().enumerate() { + for (j, &r) in right.iter().enumerate() { + assert_eq!( + compare(i, j), + expected(l, r), + "{l:?} vs {r:?} in {}", + l_array.data_type() + ); + } + } + } + Ok(()) + } + + #[test] + fn lists_compare_elements_then_length_and_skip_null_slots() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Float64, true)); + // The last row of each side is a null list over values that differ between the sides. + let left: ArrayRef = Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths([2, 1, 0, 1]), + Arc::new(Float64Array::from(vec![ + Some(0.0), + Some(1.0), + Some(f64::NAN), + Some(1.0), + ])), + Some(NullBuffer::from(vec![true, true, true, false])), + )); + let right: ArrayRef = Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths([1, 2, 1, 1]), + Arc::new(Float64Array::from(vec![ + Some(-0.0), + Some(f64::INFINITY), + Some(5.0), + None, + Some(9.0), + ])), + Some(NullBuffer::from(vec![true, true, true, false])), + )); + let compare = spark_comparator(left.as_ref(), right.as_ref())?; + // `[0.0, 1.0]` extends `[-0.0]`, so it is greater. + assert_eq!(compare(0, 0), Ordering::Greater); + // `[NaN]` against `[Infinity, 5.0]`: the first element decides before the length. + assert_eq!(compare(1, 1), Ordering::Greater); + // `[]` is a prefix of `[null]`. + assert_eq!(compare(2, 2), Ordering::Less); + assert_eq!(compare(3, 3), Ordering::Equal); + assert_eq!(compare(3, 2), Ordering::Less); + assert_eq!(compare(2, 3), Ordering::Greater); + Ok(()) + } + + #[test] + fn struct_fields_without_floats_break_ties() -> Result<()> { + let structs = |floats: Vec, ints: Vec| -> ArrayRef { + Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("f", DataType::Float64, true)), + Arc::new(Float64Array::from(floats)) as ArrayRef, + ), + ( + Arc::new(Field::new("i", DataType::Int32, true)), + Arc::new(Int32Array::from(ints)) as ArrayRef, + ), + ])) + }; + let left = structs(vec![-0.0, f64::NAN, 1.0], vec![2, 1, 1]); + let right = structs( + vec![0.0, f64::from_bits(0xfff8_0000_0000_0000), 2.0], + vec![1, 1, 0], + ); + let compare = spark_comparator(left.as_ref(), right.as_ref())?; + assert_eq!(compare(0, 0), Ordering::Greater); + assert_eq!(compare(1, 1), Ordering::Equal); + assert_eq!(compare(2, 2), Ordering::Less); + Ok(()) + } + + #[test] + fn mismatched_types_are_rejected() { + let floats: ArrayRef = Arc::new(Float64Array::from(vec![1.0])); + let ints: ArrayRef = Arc::new(Int32Array::from(vec![1])); + assert!(spark_comparator(floats.as_ref(), ints.as_ref()).is_err()); + } +} diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs new file mode 100644 index 00000000000..9ff1b0dc3b8 --- /dev/null +++ b/native/spark-expr/src/float_semantics/mod.rs @@ -0,0 +1,179 @@ +// 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. + +//! Spark's semantics for `-0.0` and NaN, in one place for all native expressions. +//! +//! Arrow orders floats by IEEE 754 total order: `-0.0` sorts below `0.0`, and NaNs compare by +//! their bits, so a NaN with the sign bit set sorts below `-Infinity`. Spark follows one of these +//! rules instead, each inherited from the Java API that a function's implementation calls: +//! +//! - SQL ordering, `SQLOrderingUtil.compareDoubles`: `-0.0` equals `0.0`, all NaNs are equal, and +//! NaN sorts above every other value. Comparisons, sorting, `min`/`max` and the array functions +//! that compare elements use it. See [`compare_floats`] and [`spark_comparator`]. +//! - `NormalizeNaNAndZero`, which Spark applies to grouping, join and window partition keys: +//! `-0.0` becomes `0.0` and every NaN becomes the canonical NaN. See [`normalize_float`], +//! [`normalize_floats`] and [`normalize_nested_floats`]. +//! - `java.lang.Double.equals`, which boxed keys such as those in `OpenHashSet` use: all NaNs are +//! equal, but `-0.0` and `0.0` are distinct. See [`canonicalize_nan`]. +//! - `Murmur3Hash` and `XxHash64`: `-0.0` hashes as `0.0`, and NaN as the canonical NaN. See +//! [`hash_input`]. +//! +//! A native expression must follow the rule of the Spark function it replaces, so build it from +//! these helpers rather than a local copy. Non-canonical NaNs are not a corner case: on x86-64 +//! every NaN that arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. + +mod compare; +mod normalize; + +pub use compare::spark_comparator; +pub use normalize::{ + has_float_leaf, normalize_floats, normalize_nested_floats, NormalizeNaNAndZero, + NormalizeNestedFloats, +}; + +use num::Float; +use std::cmp::Ordering; + +/// Spark's `NormalizeNaNAndZero`: every NaN becomes the canonical NaN and `-0.0` becomes `0.0`, +/// so two values that Spark's SQL ordering treats as equal also have the same bits. +#[inline] +pub fn normalize_float(v: T) -> T { + if v.is_nan() { + T::nan() + } else if v == T::neg_zero() { + T::zero() + } else { + v + } +} + +/// Canonicalizes NaN but keeps the sign of zero. This is the equality of a boxed +/// `java.lang.Double`, whose `equals` compares `doubleToLongBits`. +#[inline] +pub fn canonicalize_nan(v: T) -> T { + if v.is_nan() { + T::nan() + } else { + v + } +} + +/// Spark's SQL ordering, `SQLOrderingUtil.compareDoubles`: `-0.0` equals `0.0`, all NaNs are +/// equal, and NaN is greater than every other value, including positive infinity. +#[inline] +pub fn compare_floats(left: T, right: T) -> Ordering { + // IEEE 754 already treats the two zeros as equal. Only a NaN leaves the values unordered. + left.partial_cmp(&right) + .unwrap_or_else(|| left.is_nan().cmp(&right.is_nan())) +} + +/// The value that Spark's `Murmur3Hash` and `XxHash64` hash in place of a float. `-0.0` hashes as +/// `0.0`, so the two zeros hash alike. +/// +/// Spark reads the other bits through `doubleToLongBits`, which also canonicalizes NaN. This does +/// not yet, so a NaN whose bits are not canonical hashes differently from Spark (#6385). +#[inline] +pub fn hash_input(v: T) -> T { + if v == T::zero() { + // `-0.0 == 0.0` in IEEE 754, so this catches negative zero as well. + T::zero() + } else { + v + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A NaN with the sign bit set, which arithmetic produces on x86-64. + const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); + /// A signaling NaN with a payload. + const PAYLOAD_NAN: f64 = f64::from_bits(0x7ff0_0000_0000_0001); + const NEGATIVE_NAN_F32: f32 = f32::from_bits(0xffc0_0000); + + #[test] + fn normalize_float_folds_negative_zero_and_nan() { + for (value, expected) in [ + (-0.0, 0.0), + (0.0, 0.0), + (NEGATIVE_NAN, f64::NAN), + (PAYLOAD_NAN, f64::NAN), + (f64::NEG_INFINITY, f64::NEG_INFINITY), + (-1.5, -1.5), + ] { + assert_eq!(normalize_float(value).to_bits(), expected.to_bits()); + } + assert_eq!(normalize_float(-0.0f32).to_bits(), 0.0f32.to_bits()); + assert_eq!( + normalize_float(NEGATIVE_NAN_F32).to_bits(), + f32::NAN.to_bits() + ); + } + + #[test] + fn canonicalize_nan_keeps_the_sign_of_zero() { + for (value, expected) in [ + (-0.0, -0.0), + (0.0, 0.0), + (NEGATIVE_NAN, f64::NAN), + (PAYLOAD_NAN, f64::NAN), + (-1.5, -1.5), + ] { + assert_eq!(canonicalize_nan(value).to_bits(), expected.to_bits()); + } + assert_eq!(canonicalize_nan(-0.0f32).to_bits(), (-0.0f32).to_bits()); + assert_eq!( + canonicalize_nan(NEGATIVE_NAN_F32).to_bits(), + f32::NAN.to_bits() + ); + } + + #[test] + fn compare_floats_matches_compare_doubles() { + // Ascending under `compareDoubles`, with values that compare equal grouped together. + let groups = [ + vec![f64::NEG_INFINITY], + vec![-1.0], + vec![-0.0, 0.0], + vec![f64::MIN_POSITIVE], + vec![f64::INFINITY], + vec![f64::NAN, NEGATIVE_NAN, PAYLOAD_NAN], + ]; + for (i, left) in groups.iter().enumerate() { + for (j, right) in groups.iter().enumerate() { + for (&l, &r) in left.iter().flat_map(|l| right.iter().map(move |r| (l, r))) { + assert_eq!(compare_floats(l, r), i.cmp(&j), "{l:?} vs {r:?}"); + } + } + } + assert_eq!( + compare_floats(NEGATIVE_NAN_F32, f32::INFINITY), + Ordering::Greater + ); + assert_eq!(compare_floats(-0.0f32, 0.0f32), Ordering::Equal); + } + + #[test] + fn hash_input_folds_negative_zero() { + assert_eq!(hash_input(-0.0f64).to_bits(), 0); + assert_eq!(hash_input(-0.0f32).to_bits(), 0); + assert_eq!(hash_input(0.0f64).to_bits(), 0); + assert_eq!(hash_input(-1.5f64), -1.5); + assert_eq!(hash_input(f64::NEG_INFINITY), f64::NEG_INFINITY); + } +} diff --git a/native/spark-expr/src/array_funcs/nested_float_normalize.rs b/native/spark-expr/src/float_semantics/normalize.rs similarity index 78% rename from native/spark-expr/src/array_funcs/nested_float_normalize.rs rename to native/spark-expr/src/float_semantics/normalize.rs index fdb93f5a153..2fb2ec0dc8c 100644 --- a/native/spark-expr/src/array_funcs/nested_float_normalize.rs +++ b/native/spark-expr/src/float_semantics/normalize.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::math_funcs::internal::normalize_float; +use super::normalize_float; use arrow::array::{ Array, ArrayRef, AsArray, FixedSizeListArray, Float32Array, Float64Array, LargeListArray, ListArray, StructArray, @@ -29,6 +29,96 @@ use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; +/// Spark's `NormalizeNaNAndZero` expression, which applies [`normalize_float`] to a Float32 or +/// Float64 child. +#[derive(Debug, Eq)] +pub struct NormalizeNaNAndZero { + pub data_type: DataType, + pub child: Arc, +} + +impl PartialEq for NormalizeNaNAndZero { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) && self.data_type.eq(&other.data_type) + } +} + +impl Hash for NormalizeNaNAndZero { + fn hash(&self, state: &mut H) { + self.child.hash(state); + self.data_type.hash(state); + } +} + +impl NormalizeNaNAndZero { + pub fn new(data_type: DataType, child: Arc) -> Self { + Self { data_type, child } + } + + /// Wraps a Float32 or Float64 key so that keys equal under Spark's SQL ordering also compare + /// equal in Arrow. Other types are returned as is, and so is a key that is already wrapped, + /// because Spark may have normalized a partition or join key itself. + pub fn wrap_if_needed( + child: Arc, + schema: &Schema, + ) -> Result> { + let data_type = child.data_type(schema)?; + if matches!(data_type, DataType::Float32 | DataType::Float64) + && child.downcast_ref::().is_none() + { + Ok(Arc::new(Self::new(data_type, child))) + } else { + Ok(child) + } + } +} + +impl PhysicalExpr for NormalizeNaNAndZero { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn data_type(&self, input_schema: &Schema) -> Result { + self.child.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.child.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let cv = self.child.evaluate(batch)?; + let array = cv.into_array(batch.num_rows())?; + + match &self.data_type { + DataType::Float32 | DataType::Float64 => { + Ok(ColumnarValue::Array(normalize_floats(&array))) + } + dt => panic!("Unexpected data type {dt:?}"), + } + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(NormalizeNaNAndZero::new( + self.data_type.clone(), + Arc::clone(&children[0]), + ))) + } +} + +impl Display for NormalizeNaNAndZero { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FloatNormalize [child: {}]", self.child) + } +} + /// Normalizes nested IN operands, preserving constants for static membership lookup. #[derive(Debug, Eq)] pub struct NormalizeNestedFloats { @@ -119,7 +209,27 @@ impl PhysicalExpr for NormalizeNestedFloats { } } -pub(super) fn has_float_leaf(dt: &DataType) -> bool { +/// Applies [`normalize_float`] to a Float32 or Float64 array. Any other array is returned as is, +/// including a nested one: [`normalize_nested_floats`] reaches floats inside lists and structs. +pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { + match array.data_type() { + DataType::Float32 => { + let normalized: Float32Array = + array.as_primitive::().unary(normalize_float); + Arc::new(normalized) + } + DataType::Float64 => { + let normalized: Float64Array = + array.as_primitive::().unary(normalize_float); + Arc::new(normalized) + } + _ => Arc::clone(array), + } +} + +/// Whether a Float32 or Float64 field occurs at any depth of a list or struct type. The keys and +/// values of a map are not searched. +pub fn has_float_leaf(dt: &DataType) -> bool { match dt { DataType::Float32 | DataType::Float64 => true, DataType::List(field) | DataType::LargeList(field) | DataType::FixedSizeList(field, _) => { @@ -138,22 +248,13 @@ pub(super) fn has_float_leaf(dt: &DataType) -> bool { /// keeps the rebuild proportional to the float data, and it also leaves empty structs alone: /// `StructArray::new` cannot infer a length from zero columns and would panic. An empty struct /// reaches this code through Iceberg's `_partition` metadata column on an unpartitioned table. -pub(super) fn normalize_nested_floats(array: &ArrayRef) -> ArrayRef { +pub fn normalize_nested_floats(array: &ArrayRef) -> ArrayRef { if !has_float_leaf(array.data_type()) { return Arc::clone(array); } match array.data_type() { - DataType::Float32 => { - let normalized: Float32Array = - array.as_primitive::().unary(normalize_float); - Arc::new(normalized) - } - DataType::Float64 => { - let normalized: Float64Array = - array.as_primitive::().unary(normalize_float); - Arc::new(normalized) - } + DataType::Float32 | DataType::Float64 => normalize_floats(array), DataType::List(field) => { let list = array.as_list::(); let normalized_values = normalize_nested_floats(list.values()); @@ -398,6 +499,51 @@ mod tests { assert_eq!(normalized.value(5), 1.5); } + /// Range partition boundaries use `normalize_floats` to match sort keys, which are only + /// normalized when they are floats themselves, so a nested array must come back unchanged. + #[test] + fn test_normalize_floats_leaves_nested_arrays_alone() { + let flat: ArrayRef = Arc::new(Float32Array::from(vec![ + Some(-0.0), + Some(f32::from_bits(0xffc0_0001)), + None, + ])); + let normalized = normalize_floats(&flat); + let normalized = normalized.as_primitive::(); + assert_eq!(normalized.value(0).to_bits(), 0.0f32.to_bits()); + assert_eq!(normalized.value(1).to_bits(), f32::NAN.to_bits()); + assert!(normalized.is_null(2)); + + let mut builder = ListBuilder::new(Float64Builder::new()); + builder.values().append_value(-0.0); + builder.append(true); + let nested: ArrayRef = Arc::new(builder.finish()); + assert!(Arc::ptr_eq(&normalize_floats(&nested), &nested)); + } + + #[test] + fn test_wrap_float_keys_once() -> Result<()> { + use datafusion::physical_expr::expressions::Column; + let list = DataType::List(Arc::new(Field::new("item", DataType::Float64, true))); + let schema = Schema::new(vec![ + Field::new("f32", DataType::Float32, true), + Field::new("f64", DataType::Float64, true), + Field::new("i32", DataType::Int32, true), + Field::new("list", list, true), + ]); + for (index, wrapped) in [(0, true), (1, true), (2, false), (3, false)] { + let column: Arc = + Arc::new(Column::new(schema.field(index).name(), index)); + let key = NormalizeNaNAndZero::wrap_if_needed(Arc::clone(&column), &schema)?; + assert_eq!(key.downcast_ref::().is_some(), wrapped); + assert_eq!(Arc::ptr_eq(&key, &column), !wrapped); + // A key that is already normalized is not wrapped a second time. + let again = NormalizeNaNAndZero::wrap_if_needed(Arc::clone(&key), &schema)?; + assert!(Arc::ptr_eq(&again, &key)); + } + Ok(()) + } + #[test] fn test_normalize_nested_list_floats() { let mut builder = ListBuilder::new(Float64Builder::new()); diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index 7411b8dbfe5..55beb6cd378 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -116,7 +116,7 @@ macro_rules! hash_array_primitive { #[macro_export] macro_rules! hash_array_primitive_float { - ($array_type: ident, $column: ident, $ty: ident, $ty2: ident, $hashes: ident, $hash_method: ident) => { + ($array_type: ident, $column: ident, $hashes: ident, $hash_method: ident) => { let array = $column .as_any() .downcast_ref::<$array_type>() @@ -132,25 +132,15 @@ macro_rules! hash_array_primitive_float { if array.null_count() == 0 { // Fast path: no nulls, use direct indexing for i in 0..values.len() { - let value = values[i]; - // Spark uses 0 as hash for -0.0, see `Murmur3Hash` expression. - if value == 0.0 && value.is_sign_negative() { - $hashes[i] = $hash_method((0 as $ty2).to_le_bytes(), $hashes[i]); - } else { - $hashes[i] = $hash_method((value as $ty).to_le_bytes(), $hashes[i]); - } + let value = $crate::float_semantics::hash_input(values[i]); + $hashes[i] = $hash_method(value.to_le_bytes(), $hashes[i]); } } else { // Slow path: check nulls for i in 0..values.len() { if !array.is_null(i) { - let value = values[i]; - // Spark uses 0 as hash for -0.0, see `Murmur3Hash` expression. - if value == 0.0 && value.is_sign_negative() { - $hashes[i] = $hash_method((0 as $ty2).to_le_bytes(), $hashes[i]); - } else { - $hashes[i] = $hash_method((value as $ty).to_le_bytes(), $hashes[i]); - } + let value = $crate::float_semantics::hash_input(values[i]); + $hashes[i] = $hash_method(value.to_le_bytes(), $hashes[i]); } } } @@ -362,12 +352,12 @@ macro_rules! hash_list_with_primitive_elements { DataType::Float32 => { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, $hash_method, - |v: f32| if v == 0.0 && v.is_sign_negative() { (0_i32).to_le_bytes() } else { v.to_le_bytes() }); + |v: f32| $crate::float_semantics::hash_input(v).to_le_bytes()); } DataType::Float64 => { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, $hash_method, - |v: f64| if v == 0.0 && v.is_sign_negative() { (0_i64).to_le_bytes() } else { v.to_le_bytes() }); + |v: f64| $crate::float_semantics::hash_input(v).to_le_bytes()); } DataType::Boolean => { let elem_array = $values.as_any().downcast_ref::().unwrap(); @@ -461,12 +451,12 @@ macro_rules! hash_list_with_primitive_elements { DataType::Float32 => { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, $hash_method, - |v: f32| if v == 0.0 && v.is_sign_negative() { (0_i32).to_le_bytes() } else { v.to_le_bytes() }); + |v: f32| $crate::float_semantics::hash_input(v).to_le_bytes()); } DataType::Float64 => { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, $hash_method, - |v: f64| if v == 0.0 && v.is_sign_negative() { (0_i64).to_le_bytes() } else { v.to_le_bytes() }); + |v: f64| $crate::float_semantics::hash_input(v).to_le_bytes()); } DataType::Boolean => { let elem_array = $values.as_any().downcast_ref::().unwrap(); @@ -901,8 +891,6 @@ macro_rules! create_hashes_internal { $crate::hash_array_primitive_float!( Float32Array, col, - f32, - i32, $hashes_buffer, $hash_method ); @@ -911,8 +899,6 @@ macro_rules! create_hashes_internal { $crate::hash_array_primitive_float!( Float64Array, col, - f64, - i64, $hashes_buffer, $hash_method ); diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index b3379b6f5c9..2ec0388b06a 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -60,6 +60,8 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; +pub mod float_semantics; +pub use float_semantics::{NormalizeNaNAndZero, NormalizeNestedFloats}; mod hll_scalar; pub use hll_scalar::spark_hll_sketch_estimate; pub use hll_scalar::spark_hll_union; @@ -93,7 +95,7 @@ pub use math_funcs::{ create_negate_expr, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, spark_make_decimal, spark_modulo, spark_pow, spark_round, spark_sqrt, spark_unhex, spark_unscaled_value, CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr, - NormalizeNaNAndZero, WideDecimalBinaryExpr, WideDecimalOp, + WideDecimalBinaryExpr, WideDecimalOp, }; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; pub use string_funcs::*; diff --git a/native/spark-expr/src/math_funcs/internal/mod.rs b/native/spark-expr/src/math_funcs/internal/mod.rs index 98191f78196..f103e416eb9 100644 --- a/native/spark-expr/src/math_funcs/internal/mod.rs +++ b/native/spark-expr/src/math_funcs/internal/mod.rs @@ -18,12 +18,9 @@ mod checkoverflow; mod decimal_rescale_check; mod make_decimal; -mod normalize_nan; mod unscaled_value; pub use checkoverflow::CheckOverflow; pub use decimal_rescale_check::DecimalRescaleCheckOverflow; pub use make_decimal::spark_make_decimal; -pub(crate) use normalize_nan::normalize_float; -pub use normalize_nan::NormalizeNaNAndZero; pub use unscaled_value::spark_unscaled_value; diff --git a/native/spark-expr/src/math_funcs/internal/normalize_nan.rs b/native/spark-expr/src/math_funcs/internal/normalize_nan.rs deleted file mode 100644 index 85afef5666e..00000000000 --- a/native/spark-expr/src/math_funcs/internal/normalize_nan.rs +++ /dev/null @@ -1,133 +0,0 @@ -// 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::compute::unary; -use arrow::datatypes::{DataType, Schema}; -use arrow::{ - array::{as_primitive_array, ArrayRef, Float32Array, Float64Array}, - datatypes::{Float32Type, Float64Type}, - record_batch::RecordBatch, -}; -use datafusion::logical_expr::ColumnarValue; -use datafusion::physical_expr::PhysicalExpr; -use std::hash::Hash; -use std::{ - fmt::{Display, Formatter}, - sync::Arc, -}; - -#[derive(Debug, Eq)] -pub struct NormalizeNaNAndZero { - pub data_type: DataType, - pub child: Arc, -} - -impl PartialEq for NormalizeNaNAndZero { - fn eq(&self, other: &Self) -> bool { - self.child.eq(&other.child) && self.data_type.eq(&other.data_type) - } -} - -impl Hash for NormalizeNaNAndZero { - fn hash(&self, state: &mut H) { - self.child.hash(state); - self.data_type.hash(state); - } -} - -impl NormalizeNaNAndZero { - pub fn new(data_type: DataType, child: Arc) -> Self { - Self { data_type, child } - } - - /// Normalize scalar floating-point comparison keys, leaving other types unchanged. - /// Sorting and range-partition boundaries must use the same representation. - pub fn normalize_array(array: &ArrayRef) -> ArrayRef { - match array.data_type() { - DataType::Float32 => { - let input = as_primitive_array::(array); - let result: Float32Array = unary(input, normalize_float); - Arc::new(result) - } - DataType::Float64 => { - let input = as_primitive_array::(array); - let result: Float64Array = unary(input, normalize_float); - Arc::new(result) - } - _ => Arc::clone(array), - } - } -} - -impl PhysicalExpr for NormalizeNaNAndZero { - fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - Display::fmt(self, f) - } - - fn data_type(&self, input_schema: &Schema) -> datafusion::common::Result { - self.child.data_type(input_schema) - } - - fn nullable(&self, input_schema: &Schema) -> datafusion::common::Result { - self.child.nullable(input_schema) - } - - fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result { - let cv = self.child.evaluate(batch)?; - let array = cv.into_array(batch.num_rows())?; - - match &self.data_type { - DataType::Float32 | DataType::Float64 => { - Ok(ColumnarValue::Array(Self::normalize_array(&array))) - } - dt => panic!("Unexpected data type {dt:?}"), - } - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.child] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> datafusion::common::Result> { - Ok(Arc::new(NormalizeNaNAndZero::new( - self.data_type.clone(), - Arc::clone(&children[0]), - ))) - } -} - -/// Normalize a floating point value by converting all NaN representations to a canonical NaN -/// and negative zero to positive zero. This is used for Spark's comparison semantics. -#[inline] -pub(crate) fn normalize_float(v: T) -> T { - if v.is_nan() { - T::nan() - } else if v == T::neg_zero() { - T::zero() - } else { - v - } -} - -impl Display for NormalizeNaNAndZero { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "FloatNormalize [child: {}]", self.child) - } -} From c0d9c945022424b91b520450f3025b3686d9da5a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 29 Sep 2026 09:32:40 -0600 Subject: [PATCH 2/3] Restore hot-loop speed and narrow the float_semantics API - compare_floats uses the if-chain form again, so `.is_eq()` in array_position and percentile compiles to a single comparison. - float_extrema calls new float_lt/float_gt predicates, which compile to the same instructions as the inline tests they replace. The three-way comparison was up to 9x slower in that loop. - spark_equality gives nested `=` and `IN` back the length check that skips comparing the elements of lists of different lengths. - float_semantics is a private module again, re-exporting normalize_floats, canonicalize_nan and the two expressions. - The Parquet variant cast canonicalizes NaN with canonicalize_nan, and gains a test. - Tests share the NaN constants and check the per-value rules in one table, and the max_by doc defers to the module docs. --- native/core/src/execution/planner.rs | 4 +- .../core/src/parquet/cast_column/variant.rs | 5 +- .../src/parquet/cast_column/variant/tests.rs | 24 ++++ native/spark-expr/src/agg_funcs/max_min_by.rs | 25 +--- native/spark-expr/src/agg_funcs/percentile.rs | 3 +- .../src/array_funcs/array_extrema.rs | 14 +- .../src/array_funcs/nested_comparison.rs | 7 +- .../spark-expr/src/float_semantics/compare.rs | 95 ++++++++----- native/spark-expr/src/float_semantics/mod.rs | 126 ++++++++++-------- native/spark-expr/src/lib.rs | 6 +- 10 files changed, 185 insertions(+), 124 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5a8370ad8ac..9e1f423f9b5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -149,8 +149,8 @@ use datafusion_comet_proto::{ spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; use datafusion_comet_spark_expr::{ - float_semantics::normalize_floats, jvm_udf::JvmScalarUdfExpr, spark_in_list, ApproxPercentile, - ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, + jvm_udf::JvmScalarUdfExpr, normalize_floats, spark_in_list, ApproxPercentile, ArrayInsert, Avg, + AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, Regr, RegrType, SparkCastOptions, Stddev, SumDecimal, ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, diff --git a/native/core/src/parquet/cast_column/variant.rs b/native/core/src/parquet/cast_column/variant.rs index 35287733c16..5e18f603b0b 100644 --- a/native/core/src/parquet/cast_column/variant.rs +++ b/native/core/src/parquet/cast_column/variant.rs @@ -26,6 +26,7 @@ use arrow::{ }; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion_comet_common::SparkError; +use datafusion_comet_spark_expr::canonicalize_nan; use parquet::variant::{ unshred_variant, ListBuilder, MetadataBuilder, ObjectBuilder, ObjectFieldBuilder, ParentState, ReadOnlyMetadataBuilder, ValueBuilder, Variant, VariantArray, VariantBuilderExt, @@ -852,8 +853,8 @@ fn spark_typed_scalar<'m, 'v>(value: Variant<'m, 'v>) -> Variant<'m, 'v> { .map(Variant::Decimal4) .unwrap_or(value), Variant::String(s) => Variant::from(s), - Variant::Float(v) if v.is_nan() => Variant::Float(f32::NAN), - Variant::Double(v) if v.is_nan() => Variant::Double(f64::NAN), + Variant::Float(v) => Variant::Float(canonicalize_nan(v)), + Variant::Double(v) => Variant::Double(canonicalize_nan(v)), _ => value, } } diff --git a/native/core/src/parquet/cast_column/variant/tests.rs b/native/core/src/parquet/cast_column/variant/tests.rs index 52f33ecc3eb..93d0f6a7bcf 100644 --- a/native/core/src/parquet/cast_column/variant/tests.rs +++ b/native/core/src/parquet/cast_column/variant/tests.rs @@ -980,3 +980,27 @@ fn normalize_null_parent_ignores_empty_children() { assert!(output.is_null(0)); assert!(output.as_struct().column(0).is_null(0)); } + +/// Spark writes typed floats through `floatToIntBits`/`doubleToLongBits`, so every NaN comes out +/// canonical while `-0.0` keeps its sign. +#[test] +fn typed_scalar_canonicalizes_nan_and_keeps_negative_zero() { + for bits in [0xfff8_0000_0000_0000, 0x7ff0_0000_0000_0001] { + let Variant::Double(v) = spark_typed_scalar(Variant::Double(f64::from_bits(bits))) else { + panic!("expected a double"); + }; + assert_eq!(v.to_bits(), f64::NAN.to_bits()); + } + let Variant::Float(v) = spark_typed_scalar(Variant::Float(f32::from_bits(0xffc0_0000))) else { + panic!("expected a float"); + }; + assert_eq!(v.to_bits(), f32::NAN.to_bits()); + let Variant::Double(v) = spark_typed_scalar(Variant::Double(-0.0)) else { + panic!("expected a double"); + }; + assert_eq!(v.to_bits(), (-0.0f64).to_bits()); + let Variant::Float(v) = spark_typed_scalar(Variant::Float(-0.0)) else { + panic!("expected a float"); + }; + assert_eq!(v.to_bits(), (-0.0f32).to_bits()); +} diff --git a/native/spark-expr/src/agg_funcs/max_min_by.rs b/native/spark-expr/src/agg_funcs/max_min_by.rs index f0a82fee737..cf85e26ed2c 100644 --- a/native/spark-expr/src/agg_funcs/max_min_by.rs +++ b/native/spark-expr/src/agg_funcs/max_min_by.rs @@ -43,27 +43,10 @@ use std::sync::Arc; /// /// # Float orderings /// -/// A floating-point ordering column is canonicalized with [`normalize_floats`] so that Arrow's -/// row-format byte order reproduces Spark's comparison for this aggregate. -/// -/// Spark compares the ordering with `SQLOrderingUtil.compareDoubles`/`compareFloats`, wired in via -/// `PhysicalDoubleType.ordering`/`PhysicalFloatType.ordering`. That is -/// `if (x == y) 0 else java.lang.Double.compare(x, y)`, which has two consequences Arrow's row -/// format does not share: -/// -/// * `-0.0` and `0.0` tie, because the `x == y` short-circuit is IEEE equality. Arrow encodes -/// floats by flipping the bits off the sign, a total order placing `-0.0` strictly below `0.0`. -/// * every `NaN` is one value and sorts above `+Infinity`, because `Double.compare` goes through -/// `doubleToLongBits`. Arrow uses the raw bits, so a sign-bit-set `NaN` would sort below -/// `-Infinity` instead. -/// -/// Folding `-0.0` into `0.0` and every `NaN` into the canonical `NaN` makes the row bytes agree -/// with `compareDoubles` on both counts. This is verified identical on Spark 3.4 through master. -/// -/// Note this is the opposite of what `mode` needs: `mode` keys a hash map via -/// `OpenHashSet`'s `equals` (`java.lang.Double.equals`), which distinguishes `-0.0` from `0.0`, so -/// it must *not* fold them. Same two input values, different Spark comparison path, opposite -/// correct behaviour. +/// Spark compares the ordering with `SQLOrderingUtil.compareDoubles`, while the accumulators rank +/// it in Arrow's row format, which orders floats by IEEE 754 total order. The ordering column goes +/// through [`normalize_floats`] first, which makes the two agree. This is verified identical on +/// Spark 3.4 through master. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaxMinBy { name: String, diff --git a/native/spark-expr/src/agg_funcs/percentile.rs b/native/spark-expr/src/agg_funcs/percentile.rs index a5484312a33..87219db2534 100644 --- a/native/spark-expr/src/agg_funcs/percentile.rs +++ b/native/spark-expr/src/agg_funcs/percentile.rs @@ -28,7 +28,6 @@ use datafusion::logical_expr::{ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature, }; use datafusion::physical_expr::expressions::format_state_name; -use std::cmp::Ordering; use std::mem::{size_of, size_of_val}; use std::sync::Arc; @@ -306,7 +305,7 @@ fn spark_percentile(values: &mut [f64], percentile: f64) -> Option { let (_, higher_value, _) = values.select_nth_unstable_by(higher, |x, y| compare_floats(*x, *y)); let higher_value = *higher_value; - if compare_floats(lower_value, higher_value) == Ordering::Equal { + if compare_floats(lower_value, higher_value).is_eq() { return Some(lower_value); } diff --git a/native/spark-expr/src/array_funcs/array_extrema.rs b/native/spark-expr/src/array_funcs/array_extrema.rs index df048eae972..76d9121b303 100644 --- a/native/spark-expr/src/array_funcs/array_extrema.rs +++ b/native/spark-expr/src/array_funcs/array_extrema.rs @@ -18,7 +18,7 @@ use std::cmp::Ordering; use std::sync::Arc; -use crate::float_semantics::{compare_floats, spark_comparator}; +use crate::float_semantics::{float_gt, float_lt, spark_comparator}; use arrow::array::{ make_array, new_empty_array, Array, ArrayRef, AsArray, BooleanArray, ListArray, MutableArrayData, PrimitiveArray, PrimitiveBuilder, StructArray, UInt32Array, @@ -130,11 +130,6 @@ where let values = array.values().as_primitive::(); let buffer = values.values(); let nulls = values.nulls(); - let ordering = if is_min { - Ordering::Less - } else { - Ordering::Greater - }; let mut result = PrimitiveBuilder::::with_capacity(array.len()); for (row, offsets) in array.offsets().windows(2).enumerate() { let mut best: Option = None; @@ -145,8 +140,11 @@ where if nulls.is_some_and(|nulls| nulls.is_null(start + index)) { continue; } - let replace = - best.is_none_or(|current| compare_floats(candidate, current) == ordering); + let replace = match best { + None => true, + Some(current) if is_min => float_lt(candidate, current), + Some(current) => float_gt(candidate, current), + }; if replace { // Copy the winning value, never normalize its zero sign or NaN bits. best = Some(candidate); diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index 8f50d7cf72d..395a5ddc3f5 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -18,7 +18,7 @@ //! Spark equality for nested floating-point values without materializing normalized columns. use crate::float_semantics::{ - has_float_leaf, normalize_nested_floats, spark_comparator, NormalizeNestedFloats, + has_float_leaf, normalize_nested_floats, spark_equality, NormalizeNestedFloats, }; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::buffer::{BooleanBuffer, NullBuffer}; @@ -70,7 +70,7 @@ impl Operand { BooleanArray::new_null(len) } else { // Inner nulls take part in structural equality. Outer SQL nulls are handled here. - let compare = spark_comparator(self.array.as_ref(), other.array.as_ref())?; + let equal = spark_equality(self.array.as_ref(), other.array.as_ref())?; let nulls = match (self.scalar, other.scalar) { (true, true) => None, (true, false) => other.array.nulls().cloned(), @@ -78,11 +78,10 @@ impl Operand { (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), }; let values = BooleanBuffer::collect_bool(len, |row| { - compare( + equal( if self.scalar { 0 } else { row }, if other.scalar { 0 } else { row }, ) - .is_eq() }); BooleanArray::new(values, nulls) }; diff --git a/native/spark-expr/src/float_semantics/compare.rs b/native/spark-expr/src/float_semantics/compare.rs index ebfb4d533c6..fd7be4bc444 100644 --- a/native/spark-expr/src/float_semantics/compare.rs +++ b/native/spark-expr/src/float_semantics/compare.rs @@ -31,17 +31,35 @@ use std::ops::Range; /// Subtrees without a float leaf use Arrow's comparator, which orders them the same way. The two /// arrays must have the same type, ignoring field names and nullability. pub fn spark_comparator(left: &dyn Array, right: &dyn Array) -> Result { - if !DFSchema::datatype_is_logically_equal(left.data_type(), right.data_type()) { - return internal_err!( + check_types(left, right)?; + comparator(left, right, false) +} + +/// Builds a test of whether `left[i]` equals `right[j]` in the ordering of [`spark_comparator`]. +/// Lists of different lengths are unequal without comparing their elements. +pub fn spark_equality( + left: &dyn Array, + right: &dyn Array, +) -> Result bool + Send + Sync>> { + check_types(left, right)?; + let compare = comparator(left, right, true)?; + Ok(Box::new(move |i, j| compare(i, j).is_eq())) +} + +fn check_types(left: &dyn Array, right: &dyn Array) -> Result<()> { + if DFSchema::datatype_is_logically_equal(left.data_type(), right.data_type()) { + Ok(()) + } else { + internal_err!( "Spark comparison requires matching types, got {} and {}", left.data_type(), right.data_type() - ); + ) } - comparator(left, right) } -fn comparator(left: &dyn Array, right: &dyn Array) -> Result { +/// With `equality` set, the comparator only has to tell equal from unequal values. +fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { if !has_float_leaf(left.data_type()) { let options = SortOptions { descending: false, @@ -52,10 +70,10 @@ fn comparator(left: &dyn Array, right: &dyn Array) -> Result { match left.data_type() { DataType::Float32 => Ok(float_comparator::(left, right)), DataType::Float64 => Ok(float_comparator::(left, right)), - DataType::List(_) => list_comparator::(left, right), - DataType::LargeList(_) => list_comparator::(left, right), - DataType::FixedSizeList(_, _) => fixed_size_list_comparator(left, right), - DataType::Struct(_) => struct_comparator(left, right), + DataType::List(_) => list_comparator::(left, right, equality), + DataType::LargeList(_) => list_comparator::(left, right, equality), + DataType::FixedSizeList(_, _) => fixed_size_list_comparator(left, right, equality), + DataType::Struct(_) => struct_comparator(left, right, equality), dt => internal_err!("Unsupported type for Spark comparison: {dt}"), } } @@ -72,42 +90,56 @@ where fn list_comparator( left: &dyn Array, right: &dyn Array, + equality: bool, ) -> Result { let (l, r) = (left.as_list::(), right.as_list::()); - let compare = comparator(l.values().as_ref(), r.values().as_ref())?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; let (l, r) = (l.offsets().clone(), r.offsets().clone()); Ok(nulls_first(left, right, move |i, j| { let left = l[i].as_usize()..l[i + 1].as_usize(); let right = r[j].as_usize()..r[j + 1].as_usize(); - lexicographic(&compare, left, right) + lexicographic(&compare, left, right, equality) })) } -fn fixed_size_list_comparator(left: &dyn Array, right: &dyn Array) -> Result { +fn fixed_size_list_comparator( + left: &dyn Array, + right: &dyn Array, + equality: bool, +) -> Result { let (l, r) = (left.as_fixed_size_list(), right.as_fixed_size_list()); - let compare = comparator(l.values().as_ref(), r.values().as_ref())?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; let (l, r) = (l.value_length() as usize, r.value_length() as usize); Ok(nulls_first(left, right, move |i, j| { - lexicographic(&compare, i * l..(i + 1) * l, j * r..(j + 1) * r) + lexicographic(&compare, i * l..(i + 1) * l, j * r..(j + 1) * r, equality) })) } -/// Compares two runs of child values element by element, then by length. -fn lexicographic(compare: &DynComparator, left: Range, right: Range) -> Ordering { +/// Compares two runs of child values element by element, then by length. When only equality +/// matters, runs of different lengths are unequal without comparing any elements. +fn lexicographic( + compare: &DynComparator, + left: Range, + right: Range, + equality: bool, +) -> Ordering { let lengths = left.len().cmp(&right.len()); + if equality && lengths.is_ne() { + return lengths; + } left.zip(right) .map(|(i, j)| compare(i, j)) .find(|ordering| ordering.is_ne()) .unwrap_or(lengths) } -fn struct_comparator(left: &dyn Array, right: &dyn Array) -> Result { +fn struct_comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { let fields = left .as_struct() .columns() .iter() .zip(right.as_struct().columns()) - .map(|(l, r)| comparator(l.as_ref(), r.as_ref())) + .map(|(l, r)| comparator(l.as_ref(), r.as_ref(), equality)) .collect::>>()?; Ok(nulls_first(left, right, move |i, j| { fields @@ -148,6 +180,7 @@ fn nulls_first( #[cfg(test)] mod tests { use super::*; + use crate::float_semantics::{NEGATIVE_NAN, PAYLOAD_NAN}; use arrow::array::{ ArrayRef, FixedSizeListArray, Float64Array, Int32Array, LargeListArray, ListArray, StructArray, @@ -193,7 +226,7 @@ mod tests { fn floats_compare_as_spark_orders_them_at_every_depth() -> Result<()> { let left = vec![ Some(-0.0), - Some(f64::from_bits(0xfff8_0000_0000_0000)), + Some(NEGATIVE_NAN), None, Some(1.0), Some(f64::INFINITY), @@ -202,7 +235,7 @@ mod tests { Some(0.0), Some(f64::INFINITY), None, - Some(f64::from_bits(0x7ff0_0000_0000_0001)), + Some(PAYLOAD_NAN), Some(f64::NEG_INFINITY), Some(-1.0), ]; @@ -216,14 +249,13 @@ mod tests { let nested_right = nestings(Arc::new(Float64Array::from(right.clone()))); for (l_array, r_array) in nested_left.iter().zip(&nested_right) { let compare = spark_comparator(l_array.as_ref(), r_array.as_ref())?; + let equal = spark_equality(l_array.as_ref(), r_array.as_ref())?; for (i, &l) in left.iter().enumerate() { for (j, &r) in right.iter().enumerate() { - assert_eq!( - compare(i, j), - expected(l, r), - "{l:?} vs {r:?} in {}", - l_array.data_type() - ); + let expected = expected(l, r); + let context = format!("{l:?} vs {r:?} in {}", l_array.data_type()); + assert_eq!(compare(i, j), expected, "{context}"); + assert_eq!(equal(i, j), expected.is_eq(), "{context}"); } } } @@ -267,6 +299,11 @@ mod tests { assert_eq!(compare(3, 3), Ordering::Equal); assert_eq!(compare(3, 2), Ordering::Less); assert_eq!(compare(2, 3), Ordering::Greater); + let equal = spark_equality(left.as_ref(), right.as_ref())?; + let pairs = [(0, 0), (1, 1), (2, 2), (3, 3), (3, 2), (2, 3), (0, 1)]; + for (i, j) in pairs { + assert_eq!(equal(i, j), compare(i, j).is_eq(), "({i}, {j})"); + } Ok(()) } @@ -285,10 +322,7 @@ mod tests { ])) }; let left = structs(vec![-0.0, f64::NAN, 1.0], vec![2, 1, 1]); - let right = structs( - vec![0.0, f64::from_bits(0xfff8_0000_0000_0000), 2.0], - vec![1, 1, 0], - ); + let right = structs(vec![0.0, NEGATIVE_NAN, 2.0], vec![1, 1, 0]); let compare = spark_comparator(left.as_ref(), right.as_ref())?; assert_eq!(compare(0, 0), Ordering::Greater); assert_eq!(compare(1, 1), Ordering::Equal); @@ -301,5 +335,6 @@ mod tests { let floats: ArrayRef = Arc::new(Float64Array::from(vec![1.0])); let ints: ArrayRef = Arc::new(Int32Array::from(vec![1])); assert!(spark_comparator(floats.as_ref(), ints.as_ref()).is_err()); + assert!(spark_equality(floats.as_ref(), ints.as_ref()).is_err()); } } diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs index 9ff1b0dc3b8..7d79c3dd315 100644 --- a/native/spark-expr/src/float_semantics/mod.rs +++ b/native/spark-expr/src/float_semantics/mod.rs @@ -33,13 +33,17 @@ //! [`hash_input`]. //! //! A native expression must follow the rule of the Spark function it replaces, so build it from -//! these helpers rather than a local copy. Non-canonical NaNs are not a corner case: on x86-64 +//! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes or hashes the +//! values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's total order +//! agrees with `compareDoubles`. Where Comet compares values itself, or has to return the original +//! bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], +//! [`spark_comparator`] or [`spark_equality`]. Non-canonical NaNs are not a corner case: on x86-64 //! every NaN that arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. mod compare; mod normalize; -pub use compare::spark_comparator; +pub use compare::{spark_comparator, spark_equality}; pub use normalize::{ has_float_leaf, normalize_floats, normalize_nested_floats, NormalizeNaNAndZero, NormalizeNestedFloats, @@ -76,9 +80,30 @@ pub fn canonicalize_nan(v: T) -> T { /// equal, and NaN is greater than every other value, including positive infinity. #[inline] pub fn compare_floats(left: T, right: T) -> Ordering { - // IEEE 754 already treats the two zeros as equal. Only a NaN leaves the values unordered. - left.partial_cmp(&right) - .unwrap_or_else(|| left.is_nan().cmp(&right.is_nan())) + // In this form a caller's `.is_eq()` compiles to the equality test alone, which a + // `partial_cmp` with a NaN fallback does not. + if left == right || (left.is_nan() && right.is_nan()) { + Ordering::Equal + } else if left > right || left.is_nan() { + Ordering::Greater + } else { + Ordering::Less + } +} + +/// Whether `left` sorts before `right` in Spark's SQL ordering, the same as +/// `compare_floats(left, right).is_lt()`. As a single test it compiles to a well-predicted branch +/// in a scan for a minimum, where the three-way comparison is several times slower. +#[inline] +pub fn float_lt(left: T, right: T) -> bool { + left < right || (!left.is_nan() && right.is_nan()) +} + +/// Whether `left` sorts after `right` in Spark's SQL ordering, the same as +/// `compare_floats(left, right).is_gt()`. See [`float_lt`]. +#[inline] +pub fn float_gt(left: T, right: T) -> bool { + left > right || (left.is_nan() && !right.is_nan()) } /// The value that Spark's `Murmur3Hash` and `XxHash64` hash in place of a float. `-0.0` hashes as @@ -96,55 +121,55 @@ pub fn hash_input(v: T) -> T { } } +/// A NaN with the sign bit set, which arithmetic produces on x86-64. +#[cfg(test)] +const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); +/// A signaling NaN with a payload. +#[cfg(test)] +const PAYLOAD_NAN: f64 = f64::from_bits(0x7ff0_0000_0000_0001); + #[cfg(test)] mod tests { use super::*; - /// A NaN with the sign bit set, which arithmetic produces on x86-64. - const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); - /// A signaling NaN with a payload. - const PAYLOAD_NAN: f64 = f64::from_bits(0x7ff0_0000_0000_0001); const NEGATIVE_NAN_F32: f32 = f32::from_bits(0xffc0_0000); #[test] - fn normalize_float_folds_negative_zero_and_nan() { - for (value, expected) in [ - (-0.0, 0.0), - (0.0, 0.0), - (NEGATIVE_NAN, f64::NAN), - (PAYLOAD_NAN, f64::NAN), - (f64::NEG_INFINITY, f64::NEG_INFINITY), - (-1.5, -1.5), - ] { - assert_eq!(normalize_float(value).to_bits(), expected.to_bits()); + fn per_value_rules() { + // Each row: the input, then what `normalize_float`, `canonicalize_nan` and `hash_input` + // return for it. + let rows = [ + (-0.0, 0.0, -0.0, 0.0), + (0.0, 0.0, 0.0, 0.0), + (-1.5, -1.5, -1.5, -1.5), + ( + f64::NEG_INFINITY, + f64::NEG_INFINITY, + f64::NEG_INFINITY, + f64::NEG_INFINITY, + ), + // `hash_input` does not canonicalize NaN yet (#6385). + (NEGATIVE_NAN, f64::NAN, f64::NAN, NEGATIVE_NAN), + (PAYLOAD_NAN, f64::NAN, f64::NAN, PAYLOAD_NAN), + ]; + for (value, normalized, canonical, hashed) in rows { + assert_eq!(normalize_float(value).to_bits(), normalized.to_bits()); + assert_eq!(canonicalize_nan(value).to_bits(), canonical.to_bits()); + assert_eq!(hash_input(value).to_bits(), hashed.to_bits()); } - assert_eq!(normalize_float(-0.0f32).to_bits(), 0.0f32.to_bits()); - assert_eq!( - normalize_float(NEGATIVE_NAN_F32).to_bits(), - f32::NAN.to_bits() - ); - } - - #[test] - fn canonicalize_nan_keeps_the_sign_of_zero() { - for (value, expected) in [ - (-0.0, -0.0), - (0.0, 0.0), - (NEGATIVE_NAN, f64::NAN), - (PAYLOAD_NAN, f64::NAN), - (-1.5, -1.5), - ] { - assert_eq!(canonicalize_nan(value).to_bits(), expected.to_bits()); + let rows = [ + (-0.0, 0.0, -0.0, 0.0), + (NEGATIVE_NAN_F32, f32::NAN, f32::NAN, NEGATIVE_NAN_F32), + ]; + for (value, normalized, canonical, hashed) in rows { + assert_eq!(normalize_float(value).to_bits(), normalized.to_bits()); + assert_eq!(canonicalize_nan(value).to_bits(), canonical.to_bits()); + assert_eq!(hash_input(value).to_bits(), hashed.to_bits()); } - assert_eq!(canonicalize_nan(-0.0f32).to_bits(), (-0.0f32).to_bits()); - assert_eq!( - canonicalize_nan(NEGATIVE_NAN_F32).to_bits(), - f32::NAN.to_bits() - ); } #[test] - fn compare_floats_matches_compare_doubles() { + fn comparisons_match_compare_doubles() { // Ascending under `compareDoubles`, with values that compare equal grouped together. let groups = [ vec![f64::NEG_INFINITY], @@ -156,8 +181,12 @@ mod tests { ]; for (i, left) in groups.iter().enumerate() { for (j, right) in groups.iter().enumerate() { - for (&l, &r) in left.iter().flat_map(|l| right.iter().map(move |r| (l, r))) { - assert_eq!(compare_floats(l, r), i.cmp(&j), "{l:?} vs {r:?}"); + for &l in left { + for &r in right { + assert_eq!(compare_floats(l, r), i.cmp(&j), "{l:?} vs {r:?}"); + assert_eq!(float_lt(l, r), i < j, "{l:?} < {r:?}"); + assert_eq!(float_gt(l, r), i > j, "{l:?} > {r:?}"); + } } } } @@ -167,13 +196,4 @@ mod tests { ); assert_eq!(compare_floats(-0.0f32, 0.0f32), Ordering::Equal); } - - #[test] - fn hash_input_folds_negative_zero() { - assert_eq!(hash_input(-0.0f64).to_bits(), 0); - assert_eq!(hash_input(-0.0f32).to_bits(), 0); - assert_eq!(hash_input(0.0f64).to_bits(), 0); - assert_eq!(hash_input(-1.5f64), -1.5); - assert_eq!(hash_input(f64::NEG_INFINITY), f64::NEG_INFINITY); - } } diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 2ec0388b06a..4c202b14f1d 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -60,8 +60,10 @@ pub mod jvm_udf; mod conditional_funcs; mod conversion_funcs; -pub mod float_semantics; -pub use float_semantics::{NormalizeNaNAndZero, NormalizeNestedFloats}; +mod float_semantics; +pub use float_semantics::{ + canonicalize_nan, normalize_floats, NormalizeNaNAndZero, NormalizeNestedFloats, +}; mod hll_scalar; pub use hll_scalar::spark_hll_sketch_estimate; pub use hll_scalar::spark_hll_union; From a1d999c9aeffec13a4759ecfd8128471b67a189e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 29 Sep 2026 10:20:22 -0600 Subject: [PATCH 3/3] Return an error, not a panic, for mismatched physical types in the comparator `check_types` compares logical types, so a `Dictionary(_, Float64)` passes as `Float64` and then panicked when `float_comparator` downcast it. The comparator now matches on both sides' types and returns the internal error. `mismatched_types_are_rejected` covers a dictionary on either side, at the top level and inside a list. --- .../spark-expr/src/float_semantics/compare.rs | 58 +++++++++++++++---- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/native/spark-expr/src/float_semantics/compare.rs b/native/spark-expr/src/float_semantics/compare.rs index fd7be4bc444..24c01e1e9df 100644 --- a/native/spark-expr/src/float_semantics/compare.rs +++ b/native/spark-expr/src/float_semantics/compare.rs @@ -67,14 +67,18 @@ fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result Ok(float_comparator::(left, right)), - DataType::Float64 => Ok(float_comparator::(left, right)), - DataType::List(_) => list_comparator::(left, right, equality), - DataType::LargeList(_) => list_comparator::(left, right, equality), - DataType::FixedSizeList(_, _) => fixed_size_list_comparator(left, right, equality), - DataType::Struct(_) => struct_comparator(left, right, equality), - dt => internal_err!("Unsupported type for Spark comparison: {dt}"), + match (left.data_type(), right.data_type()) { + (DataType::Float32, DataType::Float32) => Ok(float_comparator::(left, right)), + (DataType::Float64, DataType::Float64) => Ok(float_comparator::(left, right)), + (DataType::List(_), DataType::List(_)) => list_comparator::(left, right, equality), + (DataType::LargeList(_), DataType::LargeList(_)) => { + list_comparator::(left, right, equality) + } + (DataType::FixedSizeList(_, _), DataType::FixedSizeList(_, _)) => { + fixed_size_list_comparator(left, right, equality) + } + (DataType::Struct(_), DataType::Struct(_)) => struct_comparator(left, right, equality), + (l, r) => internal_err!("Unsupported types for Spark comparison: {l} and {r}"), } } @@ -182,8 +186,8 @@ mod tests { use super::*; use crate::float_semantics::{NEGATIVE_NAN, PAYLOAD_NAN}; use arrow::array::{ - ArrayRef, FixedSizeListArray, Float64Array, Int32Array, LargeListArray, ListArray, - StructArray, + ArrayRef, DictionaryArray, FixedSizeListArray, Float64Array, Int32Array, LargeListArray, + ListArray, StructArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::Field; @@ -334,7 +338,37 @@ mod tests { fn mismatched_types_are_rejected() { let floats: ArrayRef = Arc::new(Float64Array::from(vec![1.0])); let ints: ArrayRef = Arc::new(Int32Array::from(vec![1])); - assert!(spark_comparator(floats.as_ref(), ints.as_ref()).is_err()); - assert!(spark_equality(floats.as_ref(), ints.as_ref()).is_err()); + // A dictionary passes the logical type check, so it must not reach a downcast either. + let dictionary: ArrayRef = Arc::new(DictionaryArray::new( + Int32Array::from(vec![0]), + Arc::clone(&floats), + )); + let list = |values: &ArrayRef| -> ArrayRef { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths([1]), + Arc::clone(values), + None, + )) + }; + let (float_list, dictionary_list) = (list(&floats), list(&dictionary)); + for (left, right) in [ + (&floats, &ints), + (&floats, &dictionary), + (&dictionary, &floats), + (&float_list, &dictionary_list), + (&dictionary_list, &float_list), + ] { + let types = format!("{} and {}", left.data_type(), right.data_type()); + assert!( + spark_comparator(left.as_ref(), right.as_ref()).is_err(), + "{types}" + ); + assert!( + spark_equality(left.as_ref(), right.as_ref()).is_err(), + "{types}" + ); + } } }