diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index b97a1c6083a..140076312eb 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -120,3 +120,15 @@ any version other than 4.2.0 falls back even if it includes SPARK-54918. A vendo reports 4.2.0 but includes SPARK-59602 still runs natively; set `spark.comet.expression.ArrayDistinct.enabled=false` and `spark.comet.expression.ArrayUnion.enabled=false` on such a build. + +## `arrays_overlap` + +`arrays_overlap` runs natively on every Spark version and follows that version's rule for `FLOAT` +and `DOUBLE` elements in a flat array. Spark 4.0.5+, 4.1.4+, and 4.2+ treat `-0.0` and `0.0` as one +value (SPARK-54918, SPARK-59602), and earlier releases keep them apart. Every NaN representation is +one value on every version, and nested arrays and structs compare their elements in Spark's order, +in which the two zeros are equal. + +Comet chooses the rule by the Spark version number. On a vendor build that reports an older +version but includes SPARK-59602, Comet returns `false` for two flat arrays whose only common value +is a zero of opposite sign, where Spark returns `true`. diff --git a/native/spark-expr/src/array_funcs/arrays_overlap.rs b/native/spark-expr/src/array_funcs/arrays_overlap.rs index 542687c7d16..52aab72891c 100644 --- a/native/spark-expr/src/array_funcs/arrays_overlap.rs +++ b/native/spark-expr/src/array_funcs/arrays_overlap.rs @@ -28,6 +28,12 @@ //! - true if any non-null element appears in both arrays //! - null if no definite overlap but either array contains null elements //! - false if no overlap and neither array contains null elements +//! +//! Spark looks up flat elements in a hash set of boxed values, so every NaN is one value, while +//! `-0.0` and `0.0` are two. `spark_arrays_overlap` matches that. `spark_arrays_overlap_normalized` +//! normalizes flat `FLOAT` and `DOUBLE` elements before the lookup, which makes the zeros one value +//! too. `CometArraysOverlap.nativeFunction` picks the variant for the Spark version. Nested +//! elements compare in Spark's SQL ordering in both variants, where the zeros are equal. use arrow::array::{ make_comparator, Array, ArrayRef, AsArray, BooleanArray, GenericListArray, GenericStringArray, @@ -46,16 +52,20 @@ use datafusion::common::{exec_err, utils::take_function_args, HashSet, Result, S use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; +use num::Float; use std::cmp::Ordering; 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::math_funcs::internal::normalize_float; #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkArraysOverlap { signature: Signature, + /// Whether `-0.0` and `0.0` are one value in a flat `FLOAT` or `DOUBLE` array. + normalize_floats: bool, } impl Default for SparkArraysOverlap { @@ -68,13 +78,27 @@ impl SparkArraysOverlap { pub fn new() -> Self { Self { signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable), + normalize_floats: false, + } + } + + /// Normalizes flat `FLOAT` and `DOUBLE` elements before matching them, so `-0.0` and `0.0` are + /// one value. Every NaN is one value in both variants. + pub fn normalized() -> Self { + Self { + normalize_floats: true, + ..Self::new() } } } impl ScalarUDFImpl for SparkArraysOverlap { fn name(&self) -> &str { - "spark_arrays_overlap" + if self.normalize_floats { + "spark_arrays_overlap_normalized" + } else { + "spark_arrays_overlap" + } } fn signature(&self) -> &Signature { @@ -117,15 +141,15 @@ impl ScalarUDFImpl for SparkArraysOverlap { (DataType::List(_), DataType::List(_)) => arrays_overlap_list::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), + self.normalize_floats, )?, (DataType::LargeList(_), DataType::LargeList(_)) => arrays_overlap_list::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), + self.normalize_floats, )?, (l, r) => { - return exec_err!( - "spark_arrays_overlap does not support types '{l}' and '{r}'" - ) + return exec_err!("{} does not support types '{l}' and '{r}'", self.name()) } }; Ok(ColumnarValue::Array(result)) @@ -138,15 +162,15 @@ impl ScalarUDFImpl for SparkArraysOverlap { (DataType::List(_), DataType::List(_)) => arrays_overlap_list::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), + self.normalize_floats, )?, (DataType::LargeList(_), DataType::LargeList(_)) => arrays_overlap_list::( left_arr.as_any().downcast_ref().unwrap(), right_arr.as_any().downcast_ref().unwrap(), + self.normalize_floats, )?, (l, r) => { - return exec_err!( - "spark_arrays_overlap does not support types '{l}' and '{r}'" - ) + return exec_err!("{} does not support types '{l}' and '{r}'", self.name()) } }; let scalar = ScalarValue::try_from_array(&result, 0)?; @@ -163,9 +187,12 @@ impl ScalarUDFImpl for SparkArraysOverlap { /// - true if any non-null element appears in both arrays /// - null if no definite overlap but either array contains null elements /// - false otherwise +/// +/// With `normalize_floats`, `-0.0` and `0.0` are one value in a flat `FLOAT` or `DOUBLE` array. fn arrays_overlap_list( left: &GenericListArray, right: &GenericListArray, + normalize_floats: bool, ) -> Result { let left_values = left.values(); let right_values = right.values(); @@ -189,6 +216,18 @@ fn arrays_overlap_list( ) }; } + macro_rules! float_fast_path { + ($t:ty) => { + if normalize_floats { + flat_fast_path!( + NormalizedFloats(left_values.as_primitive::<$t>()), + NormalizedFloats(right_values.as_primitive::<$t>()) + ) + } else { + primitive_fast_path!($t) + } + }; + } match left_values.data_type() { DataType::Boolean => flat_fast_path!(left_values.as_boolean(), right_values.as_boolean()), @@ -200,8 +239,8 @@ fn arrays_overlap_list( DataType::UInt16 => primitive_fast_path!(UInt16Type), DataType::UInt32 => primitive_fast_path!(UInt32Type), DataType::UInt64 => primitive_fast_path!(UInt64Type), - DataType::Float32 => primitive_fast_path!(Float32Type), - DataType::Float64 => primitive_fast_path!(Float64Type), + DataType::Float32 => float_fast_path!(Float32Type), + DataType::Float64 => float_fast_path!(Float64Type), DataType::Date32 => primitive_fast_path!(Date32Type), DataType::Date64 => primitive_fast_path!(Date64Type), DataType::Decimal128(_, _) => primitive_fast_path!(Decimal128Type), @@ -356,6 +395,30 @@ impl<'a, S: OffsetSizeTrait> KeyedValues<'a> for &'a GenericStringArray { } } +/// Flat `FLOAT` or `DOUBLE` elements keyed after [`normalize_float`], so `-0.0` and `0.0` share a +/// key as well as every NaN. +struct NormalizedFloats<'a, T: ArrowPrimitiveType>(&'a PrimitiveArray); + +// Derived impls would require `T: Copy`, which Arrow's type markers don't implement. +impl Clone for NormalizedFloats<'_, T> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for NormalizedFloats<'_, T> {} + +impl<'a, T: ArrowPrimitiveType> KeyedValues<'a> for NormalizedFloats<'a, T> +where + T::Native: OverlapKey + Float, +{ + type Key = ::Key; + + fn key_at(&self, i: usize) -> Option { + (!self.0.is_null(i)).then(|| normalize_float(self.0.value(i)).overlap_key()) + } +} + /// Above this many pairwise comparisons a hash probe beats the nested scan. const NESTED_SCAN_BUDGET: usize = 256; @@ -415,9 +478,8 @@ fn normalize_list_element_floats( /// Fallback for nested and otherwise unhandled element types. /// -/// note: Spark's flat arrays_overlap (HashSet) treats -0.0 and 0.0 as different, -/// only the nested path here treats them as equal. this normalization can't move into the -/// flat fast path in arrays_overlap_list without breaking that difference. +/// Spark compares nested elements in its SQL ordering on every version, so this path normalizes +/// their floats whether or not the flat fast path in arrays_overlap_list does. fn arrays_overlap_list_generic( left: &GenericListArray, right: &GenericListArray, @@ -546,11 +608,12 @@ fn needs_comparator(dt: &DataType) -> bool { mod tests { use super::*; use arrow::array::{ - Float64Array, Float64Builder, Int32Array, Int32Builder, ListArray, ListBuilder, - StructArray, StructBuilder, + Float64Array, Float64Builder, Int32Array, Int32Builder, LargeListArray, ListArray, + ListBuilder, StructArray, StructBuilder, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Fields}; + use datafusion::config::ConfigOptions; fn make_list_array( values: &Int32Array, @@ -571,7 +634,7 @@ mod tests { let left = make_list_array(&Int32Array::from(vec![1, 2, 3]), &[0, 3], None); let right = make_list_array(&Int32Array::from(vec![3, 4, 5]), &[0, 3], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); assert!(result.is_valid(0)); @@ -584,7 +647,7 @@ mod tests { let left = make_list_array(&Int32Array::from(vec![1, 2]), &[0, 2], None); let right = make_list_array(&Int32Array::from(vec![3, 4]), &[0, 2], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(!result.value(0)); assert!(result.is_valid(0)); @@ -613,6 +676,10 @@ mod tests { .map(|value| Some(value as f32)) .chain([Some(-0.0)]) .collect::>(); + let hash_sign_right = (17..=32) + .map(|value| Some(value as f32)) + .chain([Some(-1.0)]) + .collect::>(); let left = ListArray::from_iter_primitive::([ Some(vec![Some(positive_nan)]), @@ -623,6 +690,8 @@ mod tests { Some(vec![Some(positive_nan), None]), Some(vec![Some(0.0), None]), Some(hash_nan_left), + Some(hash_zero_left.clone()), + Some(vec![Some(-1.0)]), Some(hash_zero_left), ]); let right = ListArray::from_iter_primitive::([ @@ -635,11 +704,13 @@ mod tests { Some(vec![Some(-0.0)]), Some(hash_nan_right), Some(hash_zero_right), + Some(vec![Some(1.0)]), + Some(hash_sign_right), ]); - let result = arrays_overlap_list::(&left, &right)?; - let result = result.as_any().downcast_ref::().unwrap(); - let expected = BooleanArray::from(vec![ + // Every NaN is one value in both variants, and only the normalized one also makes the + // zeros one value. The last two rows check that other values keep their sign. + let boxed = BooleanArray::from(vec![ Some(true), Some(true), Some(true), @@ -649,8 +720,30 @@ mod tests { None, Some(true), Some(false), + Some(false), + Some(false), ]); - assert_eq!(result, &expected); + let normalized = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(false), + Some(false), + ]); + for (normalize_floats, expected) in [(false, boxed), (true, normalized)] { + let result = arrays_overlap_list::(&left, &right, normalize_floats)?; + assert_eq!( + result.as_boolean(), + &expected, + "normalize_floats: {normalize_floats}" + ); + } Ok(()) } @@ -676,6 +769,10 @@ mod tests { .map(|value| Some(value as f64)) .chain([Some(-0.0)]) .collect::>(); + let hash_sign_right = (17..=32) + .map(|value| Some(value as f64)) + .chain([Some(-1.0)]) + .collect::>(); let left = ListArray::from_iter_primitive::([ Some(vec![Some(positive_nan)]), @@ -686,6 +783,8 @@ mod tests { Some(vec![Some(positive_nan), None]), Some(vec![Some(0.0), None]), Some(hash_nan_left), + Some(hash_zero_left.clone()), + Some(vec![Some(-1.0)]), Some(hash_zero_left), ]); let right = ListArray::from_iter_primitive::([ @@ -698,11 +797,13 @@ mod tests { Some(vec![Some(-0.0)]), Some(hash_nan_right), Some(hash_zero_right), + Some(vec![Some(1.0)]), + Some(hash_sign_right), ]); - let result = arrays_overlap_list::(&left, &right)?; - let result = result.as_any().downcast_ref::().unwrap(); - let expected = BooleanArray::from(vec![ + // Every NaN is one value in both variants, and only the normalized one also makes the + // zeros one value. The last two rows check that other values keep their sign. + let boxed = BooleanArray::from(vec![ Some(true), Some(true), Some(true), @@ -712,18 +813,140 @@ mod tests { None, Some(true), Some(false), + Some(false), + Some(false), + ]); + let normalized = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(false), + Some(false), ]); - assert_eq!(result, &expected); + for (normalize_floats, expected) in [(false, boxed), (true, normalized)] { + let result = arrays_overlap_list::(&left, &right, normalize_floats)?; + assert_eq!( + result.as_boolean(), + &expected, + "normalize_floats: {normalize_floats}" + ); + } Ok(()) } + /// Runs `udf` over two arguments, giving one value per row. + fn invoke(udf: &SparkArraysOverlap, args: [ColumnarValue; 2], rows: usize) -> Result { + let arg_fields = args + .iter() + .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true))) + .collect(); + udf.invoke_with_args(ScalarFunctionArgs { + args: args.into(), + arg_fields, + number_rows: rows, + return_field: Arc::new(Field::new("result", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + })? + .into_array(rows) + } + + #[test] + fn test_normalized_variant_makes_flat_signed_zeros_one_value() -> Result<()> { + // Row 0 holds the two zeros, and row 1 two values that differ only in sign. + let left = [Some(vec![Some(0.0)]), Some(vec![Some(1.0)])]; + let right = [Some(vec![Some(-0.0)]), Some(vec![Some(-1.0)])]; + let lists: [(ArrayRef, ArrayRef); 2] = [ + ( + Arc::new(ListArray::from_iter_primitive::( + left.clone(), + )), + Arc::new(ListArray::from_iter_primitive::( + right.clone(), + )), + ), + ( + Arc::new(LargeListArray::from_iter_primitive::( + left, + )), + Arc::new(LargeListArray::from_iter_primitive::( + right, + )), + ), + ]; + + for (udf, name, zeros_match) in [ + (SparkArraysOverlap::new(), "spark_arrays_overlap", false), + ( + SparkArraysOverlap::normalized(), + "spark_arrays_overlap_normalized", + true, + ), + ] { + assert_eq!(udf.name(), name); + for (left, right) in &lists { + let columns = [ + ColumnarValue::Array(Arc::clone(left)), + ColumnarValue::Array(Arc::clone(right)), + ]; + let result = invoke(&udf, columns, 2)?; + assert_eq!( + result.as_boolean(), + &BooleanArray::from(vec![zeros_match, false]), + "{name} over {} columns", + left.data_type() + ); + + let scalars = [ + ColumnarValue::Scalar(ScalarValue::try_from_array(left, 0)?), + ColumnarValue::Scalar(ScalarValue::try_from_array(right, 0)?), + ]; + let result = invoke(&udf, scalars, 1)?; + assert_eq!( + result.as_boolean(), + &BooleanArray::from(vec![zeros_match]), + "{name} over {} scalars", + left.data_type() + ); + } + } + Ok(()) + } + + #[test] + fn test_unsupported_types_error_names_the_variant() { + for udf in [SparkArraysOverlap::new(), SparkArraysOverlap::normalized()] { + let column = Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef; + let arrays = [ + ColumnarValue::Array(Arc::clone(&column)), + ColumnarValue::Array(column), + ]; + let scalars = [ + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + for (args, rows) in [(arrays, 2), (scalars, 1)] { + let message = invoke(&udf, args, rows).unwrap_err().to_string(); + assert!( + message.contains(&format!("{} does not support types", udf.name())), + "{message}" + ); + } + } + } + #[test] fn test_null_only_overlap() -> Result<()> { // [1, NULL] vs [NULL, 2] => null (no definite overlap, but nulls present) let left = make_list_array(&Int32Array::from(vec![Some(1), None]), &[0, 2], None); let right = make_list_array(&Int32Array::from(vec![None, Some(2)]), &[0, 2], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_null(0)); Ok(()) @@ -735,7 +958,7 @@ mod tests { let left = make_list_array(&Int32Array::from(vec![Some(1), None]), &[0, 2], None); let right = make_list_array(&Int32Array::from(vec![1, 2]), &[0, 2], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); assert!(result.is_valid(0)); @@ -752,7 +975,7 @@ mod tests { ); let right = make_list_array(&Int32Array::from(Vec::::new()), &[0, 0], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(!result.value(0)); assert!(result.is_valid(0)); @@ -769,7 +992,7 @@ mod tests { ); let right = make_list_array(&Int32Array::from(vec![1, 2]), &[0, 2], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_null(0)); Ok(()) @@ -781,7 +1004,7 @@ mod tests { let left = make_list_array(&Int32Array::from(vec![None::]), &[0, 1], None); let right = make_list_array(&Int32Array::from(vec![None::]), &[0, 1], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_null(0)); Ok(()) @@ -807,7 +1030,7 @@ mod tests { None, ); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!( result.is_null(0), @@ -823,7 +1046,7 @@ mod tests { let left = make_list_array(&Int32Array::from(vec![Some(3), None]), &[0, 2], None); let right = make_list_array(&Int32Array::from(vec![1, 2]), &[0, 2], None); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_null(0)); Ok(()) @@ -872,13 +1095,13 @@ mod tests { // NaN matches itself, and signed zeros are equal, matching Spark. let left = make_nested_float_list(&[&[f64::NAN]]); let right = make_nested_float_list(&[&[f64::NAN]]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); let left = make_nested_float_list(&[&[0.0]]); let right = make_nested_float_list(&[&[-0.0]]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); Ok(()) @@ -889,7 +1112,7 @@ mod tests { // [[-NaN]] vs [[NaN]] => true let left = make_nested_float_list(&[&[-f64::NAN]]); let right = make_nested_float_list(&[&[f64::NAN]]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); Ok(()) @@ -907,7 +1130,7 @@ mod tests { Some(vec![Some(5), Some(6)]), ]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -920,7 +1143,7 @@ mod tests { let left = make_nested_list(vec![Some(vec![Some(1), Some(2)])]); let right = make_nested_list(vec![Some(vec![Some(3), Some(4)])]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(!result.value(0)); @@ -933,7 +1156,7 @@ mod tests { let left = make_nested_list(vec![Some(vec![Some(1), None])]); let right = make_nested_list(vec![Some(vec![Some(1), None])]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -949,7 +1172,7 @@ mod tests { Some(vec![Some(3), Some(4)]), ]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(!result.value(0)); @@ -962,7 +1185,7 @@ mod tests { let left = make_nested_list(vec![Some(vec![None])]); let right = make_nested_list(vec![Some(vec![None])]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!( result.is_valid(0), @@ -984,7 +1207,7 @@ mod tests { ]); let right = make_nested_list(vec![Some(vec![Some(1), Some(2)])]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -1032,7 +1255,7 @@ mod tests { let left = make_struct_list(vec![Some((Some(1), Some(2))), Some((Some(3), Some(4)))]); let right = make_struct_list(vec![Some((Some(3), Some(4))), Some((Some(5), Some(6)))]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -1045,7 +1268,7 @@ mod tests { let left = make_struct_list(vec![Some((Some(1), Some(2)))]); let right = make_struct_list(vec![Some((Some(3), Some(4)))]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(!result.value(0)); @@ -1058,7 +1281,7 @@ mod tests { let left = make_struct_list(vec![Some((Some(1), None))]); let right = make_struct_list(vec![Some((Some(1), None))]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -1071,7 +1294,7 @@ mod tests { let left = make_struct_list(vec![Some((Some(1), Some(2))), Some((Some(1), None))]); let right = make_struct_list(vec![Some((Some(1), Some(2)))]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -1125,7 +1348,7 @@ mod tests { let left = make_list(vec![1.0, 2.0]); let right = make_list(vec![1.0, 2.0]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.value(0)); assert!(result.value(1)); @@ -1138,7 +1361,7 @@ mod tests { let left = make_struct_float_list(vec![Some(-0.0)]); let right = make_struct_float_list(vec![Some(0.0)]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_valid(0)); assert!(result.value(0)); @@ -1151,7 +1374,7 @@ mod tests { let left = make_struct_list(vec![None]); let right = make_struct_list(vec![Some((Some(1), Some(2)))]); - let result = arrays_overlap_list::(&left, &right)?; + let result = arrays_overlap_list::(&left, &right, false)?; let result = result.as_any().downcast_ref::().unwrap(); assert!(result.is_null(0)); Ok(()) diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 738c8705bfb..531a7af8825 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -321,6 +321,7 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkArrayPositionFunc::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraySlice::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraysOverlap::default())), + Arc::new(ScalarUDF::new_from_impl(SparkArraysOverlap::normalized())), Arc::new(ScalarUDF::new_from_impl(SparkContains::default())), Arc::new(ScalarUDF::new_from_impl(SparkDateDiff::default())), Arc::new(ScalarUDF::new_from_impl(SparkDateFromUnixDate::default())), diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 8b6a9e507ba..a9f54cdb345 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -289,6 +289,16 @@ object CometArraysOverlap override def getSupportLevel(expr: ArraysOverlap): SupportLevel = ArrayElementEqualitySupport.getSupportLevel("arrays_overlap", expr.left.dataType) + /** + * The native variant for a Spark version. Where ArraySetSupport.normalizesFloats holds, Spark + * normalizes flat FLOAT and DOUBLE elements, so -0.0 and 0.0 are one value, and the + * `_normalized` variant matches them that way. Earlier releases keep the zeros apart. Both + * variants treat every NaN as one value and compare nested elements in Spark's SQL ordering. + */ + def nativeFunction(version: String): String = + if (ArraySetSupport.normalizesFloats(version)) "spark_arrays_overlap_normalized" + else "spark_arrays_overlap" + override def convert( expr: ArraysOverlap, inputs: Seq[Attribute], @@ -297,7 +307,7 @@ object CometArraysOverlap val rightArrayExprProto = exprToProtoInternal(expr.right, inputs, binding) val arraysOverlapScalarExpr = scalarFunctionExprToProtoWithReturnType( - "spark_arrays_overlap", + nativeFunction(SPARK_VERSION), BooleanType, false, leftArrayExprProto, @@ -555,6 +565,19 @@ private[comet] object ArraySetSupport { def normalizesArgumentsInPlan(version: String): Boolean = Utils.majorMinorPatchVersion(version).contains((4, 2, 0)) + // Spark 4.2.0 normalizes the arguments of arrays_overlap in the plan (SPARK-54918), and 4.0.5, + // 4.1.4 and 4.2.1 normalize while evaluating it (SPARK-59602). Either way, Spark treats -0.0 + // and 0.0 as one value in a flat FLOAT or DOUBLE array, which CometArraysOverlap.nativeFunction + // matches. Earlier releases keep them apart. The check reads the version rather than a + // KnownFloatingPointNormalized marker, because SPARK-59602 adds no marker, and SPARK-54918 + // normalizes CreateArray, If, CaseWhen and Coalesce without wrapping the resulting array. + def normalizesFloats(version: String): Boolean = + Utils.majorMinorPatchVersion(version).exists { + case (4, 0, patch) => patch >= 5 + case (4, 1, patch) => patch >= 4 + case (major, minor, _) => major > 4 || (major == 4 && minor >= 2) + } + def supportLevel(dataType: DataType): SupportLevel = { if (SupportLevel.containsType(dataType, classOf[FloatType], classOf[DoubleType]) && !normalizesArgumentsInPlan(SPARK_VERSION)) { diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 4fd924be70f..e3a6286e378 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -23,7 +23,7 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayDistinct, ArrayExcept, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayRepeat, ArrayUnion} +import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayDistinct, ArrayExcept, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayRepeat, ArraysOverlap, ArrayUnion} import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayRemove} import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Cast, CreateArray, ElementAt, Literal, MonotonicallyIncreasingID, Sequence} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{ArrayType, DoubleType, FloatType, IntegerType import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.serde.{ArraySetSupport, CometArrayDistinct, CometArrayExcept, CometArrayJoin, CometArrayRemove, CometArrayReverse, CometArrayUnion, CometFlatten, CometSequence, Compatible, ExprOuterClass, Incompatible, SparkFixes} +import org.apache.comet.serde.{ArraySetSupport, CometArrayDistinct, CometArrayExcept, CometArrayJoin, CometArrayRemove, CometArrayReverse, CometArraysOverlap, CometArrayUnion, CometFlatten, CometSequence, Compatible, ExprOuterClass, Incompatible, SparkFixes} import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -731,6 +731,29 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp assert(nullableDelimiter.isDefined && !nullableDelimiter.get.hasIf) } + test("arrays_overlap native variant by Spark version") { + // 4.2.0 normalizes the arguments in the plan (SPARK-54918), and 4.0.5+, 4.1.4+ and 4.2.1+ + // normalize during evaluation (SPARK-59602). Where Spark normalizes, -0.0 and 0.0 are one + // value in a flat FLOAT or DOUBLE array, which no pinned Spark profile can show: 4.2.0 + // normalizes the arguments before the native kernel sees them. + Seq("3.4.3", "3.5.9", "4.0.0", "4.0.4", "4.0.4-SNAPSHOT", "4.1.0", "4.1.3").foreach { + version => + assert(CometArraysOverlap.nativeFunction(version) == "spark_arrays_overlap", version) + } + Seq("4.0.5", "4.0.10", "4.1.4", "4.2.0", "4.2.1", "4.3.0", "5.0.0").foreach { version => + assert( + CometArraysOverlap.nativeFunction(version) == "spark_arrays_overlap_normalized", + version) + } + // convert asks for the variant of the Spark version it runs on. + val array = AttributeReference("a", ArrayType(DoubleType))() + val proto = + CometArraysOverlap.convert(ArraysOverlap(array, array), Seq(array), binding = false) + assert( + proto.map(_.getScalarFunc.getFunc) == + Some(CometArraysOverlap.nativeFunction(org.apache.spark.SPARK_VERSION))) + } + test("arrays_overlap") { Seq(true, false).foreach { dictionaryEnabled => withTempDir { dir =>