diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index e047db39a5740..2da9c5e6d40f2 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -1407,9 +1407,9 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result Ok(PrimitiveArray::new(rows_number.into(), None)) } -/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array. -/// For non-float arrays returns the input unchanged. NaN payloads are -/// preserved. +/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array, +/// including dictionary-wrapped floats. For other arrays, returns the input +/// unchanged. NaN payloads are preserved. /// /// Arrow's comparison kernels (`arrow::compute::kernels::cmp::eq` etc.) and /// row-encoding (`arrow::row::RowConverter`) use IEEE 754 totalOrder @@ -1430,6 +1430,17 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { const NEG_ZERO_F32_BITS: u32 = (-0.0_f32).to_bits(); const NEG_ZERO_F64_BITS: u64 = (-0.0_f64).to_bits(); match array.data_type() { + DataType::Dictionary(_, value_type) + if is_float_or_dictionary_float(value_type) => + { + let dictionary = array.as_any_dictionary(); + let values = normalize_float_zero(dictionary.values()); + if Arc::ptr_eq(&values, dictionary.values()) { + Arc::clone(array) + } else { + dictionary.with_values(values) + } + } DataType::Float32 => { let arr: &Float32Array = array.as_primitive::(); if !arr @@ -1439,8 +1450,13 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { { return Arc::clone(array); } - let normalized: Float32Array = - arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f32 } else { v }); + let normalized: Float32Array = arr.unary(|v| { + if v.to_bits() == NEG_ZERO_F32_BITS { + 0.0_f32 + } else { + v + } + }); Arc::new(normalized) } DataType::Float64 => { @@ -1452,8 +1468,13 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { { return Arc::clone(array); } - let normalized: Float64Array = - arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f64 } else { v }); + let normalized: Float64Array = arr.unary(|v| { + if v.to_bits() == NEG_ZERO_F64_BITS { + 0.0_f64 + } else { + v + } + }); Arc::new(normalized) } DataType::Float16 => { @@ -1466,8 +1487,8 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { return Arc::clone(array); } let normalized: Float16Array = arr.unary(|v| { - if v.to_bits() << 1 == 0 { - half::f16::from_bits(0) + if v.to_bits() == NEG_ZERO_F16_BITS { + half::f16::ZERO } else { v } @@ -1478,22 +1499,38 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { } } +fn is_float_or_dictionary_float(mut data_type: &DataType) -> bool { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type.is_floating() +} + /// Replace `-0.0` with `+0.0` in `Float16`, `Float32`, or `Float64` scalar -/// values. Other variants are returned unchanged. See [`normalize_float_zero`] -/// for context. -pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue { - match scalar { - ScalarValue::Float32(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float32(Some(0.0)) +/// values, including dictionary-wrapped floats. Other variants are returned +/// unchanged. See [`normalize_float_zero`] for context. +pub fn normalize_float_zero_scalar(mut scalar: ScalarValue) -> ScalarValue { + let mut value = &mut scalar; + while let ScalarValue::Dictionary(_, dictionary_value) = value { + value = dictionary_value.as_mut(); + } + + match value { + ScalarValue::Float32(Some(value)) if value.to_bits() == (-0.0_f32).to_bits() => { + *value = 0.0 } - ScalarValue::Float64(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float64(Some(0.0)) + ScalarValue::Float64(Some(value)) if value.to_bits() == (-0.0_f64).to_bits() => { + *value = 0.0 } - ScalarValue::Float16(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float16(Some(half::f16::from_bits(0))) + ScalarValue::Float16(Some(value)) + if value.to_bits() == half::f16::NEG_ZERO.to_bits() => + { + *value = half::f16::ZERO; } - other => other, + _ => {} } + + scalar } #[cfg(test)] @@ -1503,9 +1540,9 @@ mod tests { use super::*; use crate::ScalarValue::Null; use arrow::{ - array::{Float64Array, Int32Array}, + array::{DictionaryArray, Float64Array, Int8Array, Int32Array}, buffer::NullBuffer, - datatypes::Int32Type, + datatypes::{Float64Type, Int8Type, Int32Type}, }; #[cfg(feature = "sql")] use sqlparser::ast::Ident; @@ -1534,6 +1571,49 @@ mod tests { } } + #[test] + fn normalize_float_zero_in_dictionary_arrays_and_scalars() -> Result<()> { + let nan = f64::from_bits(0x7ff8_0000_0000_0001); + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let array: ArrayRef = Arc::new(DictionaryArray::try_new( + keys.clone(), + Arc::new(Float64Array::from(vec![-0.0, nan, 1.0])), + )?); + + let normalized = normalize_float_zero(&array); + assert!(!Arc::ptr_eq(&normalized, &array)); + let dictionary = normalized.as_dictionary::(); + assert_eq!(dictionary.keys(), &keys); + let values = dictionary.values().as_primitive::(); + assert_eq!(values.value(0).to_bits(), 0.0_f64.to_bits()); + assert_eq!(values.value(1).to_bits(), nan.to_bits()); + assert_eq!(values.value(2), 1.0); + + let without_negative_zero: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1]), + Arc::new(Float64Array::from(vec![0.0, nan])), + )?); + assert!(Arc::ptr_eq( + &normalize_float_zero(&without_negative_zero), + &without_negative_zero + )); + + let scalar = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Float64(Some(-0.0))), + ); + let ScalarValue::Dictionary(_, value) = normalize_float_zero_scalar(scalar) + else { + unreachable!() + }; + let ScalarValue::Float64(Some(value)) = *value else { + unreachable!() + }; + assert_eq!(value.to_bits(), 0.0_f64.to_bits()); + + Ok(()) + } + #[test] fn test_bisect_linear_left_and_right() -> Result<()> { let arrays: Vec = vec![ diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 0a4ad0b804f0c..d7bd29be462de 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -31,6 +31,7 @@ use arrow::compute::kernels::boolean::{not, or_kleene}; use arrow::compute::kernels::cmp::eq as arrow_eq; use arrow::datatypes::*; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; use datafusion_common::{ DFSchema, Result, ScalarValue, assert_or_internal_err, exec_err, }; @@ -81,6 +82,27 @@ fn supports_arrow_eq(dt: &DataType) -> bool { } } +fn normalize_in_list_float_zero_value(value: ColumnarValue) -> ColumnarValue { + match value { + ColumnarValue::Array(array) + if is_float_or_dictionary_float(array.data_type()) => + { + ColumnarValue::Array(normalize_float_zero(&array)) + } + ColumnarValue::Scalar(scalar) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(scalar)) + } + value => value, + } +} + +fn is_float_or_dictionary_float(mut data_type: &DataType) -> bool { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type.is_floating() +} + /// Evaluates the list of expressions into an array, flattening any dictionaries fn evaluate_list( list: &[Arc], @@ -369,12 +391,15 @@ impl PhysicalExpr for InListExpr { // Use Arrow's vectorized eq kernel for types it supports (primitive, // boolean, string, binary, dictionary), falling back to row-by-row // comparator for unsupported types (nested, RunEndEncoded, etc.). - let value = value.into_array(num_rows)?; + // Normalize the left side once for the whole list. Doing this + // outside `compare_one` avoids rescanning it for every item. + let value = + normalize_in_list_float_zero_value(value).into_array(num_rows)?; let lhs_supports_arrow_eq = supports_arrow_eq(value.data_type()); // Helper: compare value against a single list expression let compare_one = |expr: &Arc| -> Result { - match expr.evaluate(batch)? { + match normalize_in_list_float_zero_value(expr.evaluate(batch)?) { ColumnarValue::Array(array) => { if lhs_supports_arrow_eq && supports_arrow_eq(array.data_type()) @@ -3363,6 +3388,111 @@ mod tests { Ok(()) } + #[test] + fn test_in_list_with_columns_float_signed_zero() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + ]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0])), + Arc::new(Float64Array::from(vec![-0.0, 0.0, 2.0])), + ], + )?; + + let expr = make_in_list_with_columns( + col("a", &schema)?, + vec![col("b", &schema)?], + false, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + Ok(()) + } + + #[test] + fn test_in_list_with_columns_float_scalar_signed_zero() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Float32Array::from(vec![0.0, -0.0, 1.0]))], + )?; + let list = vec![lit(ScalarValue::Float32(Some(-0.0)))]; + + for (negated, expected) in [ + (false, BooleanArray::from(vec![true, true, false])), + (true, BooleanArray::from(vec![false, false, true])), + ] { + let expr = + make_in_list_with_columns(col("a", &schema)?, list.clone(), negated); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(as_boolean_array(&result), &expected); + } + + // A scalar left-hand side is normalized before it is broadcast. + let expr = make_in_list_with_columns( + lit(ScalarValue::Float32(Some(-0.0))), + vec![col("a", &schema)?], + false, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + + Ok(()) + } + + #[test] + fn test_in_list_with_columns_dictionary_float_signed_zero() -> Result<()> { + let left: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1, 2]), + Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0])), + )?); + let right: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1, 2]), + Arc::new(Float64Array::from(vec![-0.0, 0.0, 2.0])), + )?); + let data_type = left.data_type().clone(); + let schema = Schema::new(vec![ + Field::new("a", data_type.clone(), false), + Field::new("b", data_type, false), + ]); + let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![left, right])?; + + for (negated, expected) in [ + (false, BooleanArray::from(vec![true, true, false])), + (true, BooleanArray::from(vec![false, false, true])), + ] { + let expr = make_in_list_with_columns( + col("a", &schema)?, + vec![col("b", &schema)?], + negated, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(as_boolean_array(&result), &expected); + } + + let scalar = lit(ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Float64(Some(-0.0))), + )); + let expr = make_in_list_with_columns(col("a", &schema)?, vec![scalar], false); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + + Ok(()) + } + /// Tests that short-circuit evaluation produces correct results. /// When all rows match after the first list item, remaining items /// should be skipped without affecting correctness. @@ -3871,6 +4001,29 @@ mod tests { Ok(()) } + #[test] + fn test_try_new_from_array_dict_haystack_float64_signed_zero() -> Result<()> { + // One value beyond the branchless limit selects the hash-set strategy. + let list_len = + ::MAX_LIST_LEN + 1; + let mut list_values = vec![Some(-0.0)]; + list_values.extend((1..list_len).map(|value| Some(value as f64))); + let haystack = make_f64_dict_array(list_values); + let needles: ArrayRef = Arc::new(Float64Array::from(vec![0.0, -0.0, -1.0])); + let expected = BooleanArray::from(vec![true, true, false]); + + assert_eq!( + eval_in_list_from_array(Arc::clone(&needles), Arc::clone(&haystack))?, + expected + ); + assert_eq!( + eval_in_list_from_array(wrap_in_dict(needles), haystack)?, + expected + ); + + Ok(()) + } + #[test] fn test_try_new_from_array_type_mismatch_rejects() -> Result<()> { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); diff --git a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs index 86cea37b3a98d..aa19bc44dadc8 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -29,16 +29,17 @@ //! //! # How does it work? //! -//! When the filter is built, it stores the non-null list values and chooses a -//! comparison function for that list length. Only this small function is -//! specialized for each length. The rest of [`BranchlessFilter`] is shared, -//! which keeps the generated code small. +//! When the filter is built, it stores the values needed by the comparison +//! chain and chooses a function for that chain's length. Only this small +//! function is specialized for each length. The rest of [`BranchlessFilter`] +//! is shared, which keeps the generated code small. //! //! Some Arrow types share the same in-memory representation. For example, a //! `Float32` and a `UInt32` both use four bytes per value. The filter compares //! those stored bits through an unsigned type of the same size, without copying //! the value buffer. A bit pattern is simply the bytes Arrow uses to store a -//! value. Comparing it preserves details such as `0.0` versus `-0.0` and +//! value. Float comparisons treat `0.0` and `-0.0` as equal under SQL +//! semantics, while preserving distinctions between //! different NaN values. [`BranchlessFilterType`] defines these safe, //! same-sized mappings and checks their sizes at compile time. //! @@ -75,6 +76,7 @@ use arrow::buffer::{BooleanBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; +use half::f16; use super::result::build_result_from_contains; use super::static_filter::StaticFilter; @@ -117,16 +119,20 @@ const BRANCHLESS_MAX_16B: usize = 4; /// /// `T` is the logical Arrow type accepted by the filter. `CompareType` is the /// same-width type used for the fixed comparison chain. Signed integers, -/// floats, and temporal values use an unsigned comparison type so they compare -/// by their raw bit pattern. +/// floats, and temporal values use an unsigned comparison type. Most values +/// compare by raw bit pattern; floats additionally treat both signed-zero +/// patterns as equal. pub(super) trait BranchlessFilterType: - ArrowPrimitiveType + Send + Sync + 'static + ArrowPrimitiveType + Send + Sync + Sized + 'static { type CompareType: ArrowPrimitiveType + Send + Sync + 'static; /// Maximum number of non-null IN-list values to handle with /// [`BranchlessFilter`] for this primitive type. const MAX_LIST_LEN: usize; + + /// The two signed-zero encodings for float types. + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = None; } macro_rules! branchless_filter_type { @@ -151,18 +157,33 @@ branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); -branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); +impl BranchlessFilterType for Float16Type { + type CompareType = UInt16Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_2B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([f16::ZERO.to_bits(), f16::NEG_ZERO.to_bits()]); +} branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); +impl BranchlessFilterType for Float32Type { + type CompareType = UInt32Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_4B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([0.0_f32.to_bits(), (-0.0_f32).to_bits()]); +} branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); +impl BranchlessFilterType for Float64Type { + type CompareType = UInt64Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_8B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([0.0_f64.to_bits(), (-0.0_f64).to_bits()]); +} branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); @@ -189,9 +210,9 @@ type MembershipCheck = fn(in_list_values: &[C], input_values: &[C]) -> Boolea /// `T::MAX_LIST_LEN` values. /// /// The filter stores the non-null `IN`-list values in a slice and chooses a -/// comparison function for that length. Keeping the length out of -/// `BranchlessFilter` avoids generating a full copy of the filter for every -/// supported length. +/// comparison function for that length. If a float zero occurs, both signed +/// encodings are stored once. Keeping the length out of `BranchlessFilter` +/// avoids generating a full copy of the filter for every supported length. pub(super) struct BranchlessFilter { expected_data_type: DataType, null_count: usize, @@ -218,8 +239,10 @@ where } let all_values = branchless_values::(in_array); - let mut in_list_values = Vec::with_capacity(non_null_count); - + // Float zero may add its other signed encoding. + let mut in_list_values = Vec::with_capacity( + non_null_count + usize::from(T::SIGNED_ZERO_BITS.is_some()), + ); match in_array.nulls() { None => { in_list_values.extend(all_values.iter().copied()); @@ -232,8 +255,9 @@ where } } } + materialize_signed_zero_encodings::(&mut in_list_values); - debug_assert_eq!(in_list_values.len(), non_null_count); + debug_assert!(in_list_values.len() <= non_null_count + 1); let in_list_values = in_list_values.into_boxed_slice(); let check_values = membership_check_for_len::(in_list_values.len()); @@ -246,6 +270,27 @@ where } } +fn materialize_signed_zero_encodings(values: &mut Vec>) +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq, +{ + let Some([positive_zero, negative_zero]) = T::SIGNED_ZERO_BITS else { + return; + }; + + // This list has at most 32 entries, so checking both SQL-equivalent + // encodings once is cheap and keeps per-row lookup branch-free. + match ( + values.contains(&positive_zero), + values.contains(&negative_zero), + ) { + (true, false) => values.push(negative_zero), + (false, true) => values.push(positive_zero), + _ => {} + } +} + impl StaticFilter for BranchlessFilter where T: BranchlessFilterType, @@ -299,6 +344,18 @@ where }; } + // A single float zero expands to its two physical encodings. Only these + // three extra lengths can therefore exceed the logical list-size limit. + if T::SIGNED_ZERO_BITS.is_some() && len > T::MAX_LIST_LEN { + debug_assert_eq!(len, T::MAX_LIST_LEN + 1); + return match T::MAX_LIST_LEN { + 8 => check_values::, 9>, + 16 => check_values::, 17>, + 32 => check_values::, 33>, + _ => unreachable!("signed-zero expansion exceeded an unsupported limit"), + }; + } + // Avoid creating checks for lengths a type does not support. match T::MAX_LIST_LEN { 4 => choose!(0, 1, 2, 3, 4), @@ -451,11 +508,11 @@ mod tests { assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![None, Some(true), Some(true), None, None]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), None, None]) ); assert_eq!( filter.contains(&needles, true)?, - BooleanArray::from(vec![None, Some(false), Some(false), None, None]) + BooleanArray::from(vec![Some(false), Some(false), Some(false), None, None]) ); let wrong_type = UInt16Array::from(vec![Some(0x8000), Some(0x7e01)]); @@ -466,20 +523,26 @@ mod tests { } #[test] - fn branchless_filter_floats_use_bit_equality() -> Result<()> { + fn branchless_filter_floats_use_sql_zero_equality() -> Result<()> { let nan_a = f32::from_bits(0x7fc0_0001); let nan_b = f32::from_bits(0x7fc0_0002); let haystack: ArrayRef = - Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + Arc::new(Float32Array::from(vec![Some(0.0), Some(nan_a)])); let filter = BranchlessFilter::::try_new(&haystack)?; let needles = Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]) ); + // A list containing both encodings is not expanded. + let zero_only: ArrayRef = + Arc::new(Float32Array::from(vec![Some(0.0), Some(-0.0)])); + let filter = BranchlessFilter::::try_new(&zero_only)?; + assert_eq!(filter.in_list_values.len(), 2); + let nan_a = f64::from_bits(0x7ff8_0000_0000_0001); let nan_b = f64::from_bits(0x7ff8_0000_0000_0002); let haystack: ArrayRef = @@ -490,7 +553,7 @@ mod tests { assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]) ); Ok(()) diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index ceb9bd525b965..564b96ae36028 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -27,6 +27,7 @@ use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; +use half::f16; use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; use super::result::build_in_list_result; @@ -199,6 +200,12 @@ trait BitmapFilterType: ArrowPrimitiveType + Send + Sync + 'static { /// Returns the index in the bitmap to check for this value. fn index(value: Self::Native) -> usize; + + /// Adds one value to the bitmap. + #[inline] + fn insert(bitmap: &mut Self::Storage, value: Self::Native) { + bitmap.set_bit(Self::index(value)); + } } /// `Int8` has 256 possible bit patterns, so four `u64` words cover the full domain. @@ -254,6 +261,17 @@ impl BitmapFilterType for Float16Type { fn index(value: Self::Native) -> usize { value.to_bits() as usize } + + fn insert(bitmap: &mut Self::Storage, value: Self::Native) { + if value == f16::ZERO { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + bitmap.set_bit(f16::ZERO.to_bits() as usize); + bitmap.set_bit(f16::NEG_ZERO.to_bits() as usize); + } else { + bitmap.set_bit(Self::index(value)); + } + } } /// `IN` filter backed by one bit per possible value. @@ -280,14 +298,14 @@ where match prim_array.nulls() { None => { for &v in values { - bits.set_bit(T::index(v)); + T::insert(&mut bits, v); } } Some(nulls) => { for i in BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) { - bits.set_bit(T::index(values[i])); + T::insert(&mut bits, values[i]); } } } @@ -332,6 +350,16 @@ where } } +/// Converts native values to hash keys and inserts any additional physical +/// encodings that belong to the same SQL equality class. +trait HashSetKey: From + Eq + Hash + Sized { + fn insert(values: &mut HashSet, value: V) { + values.insert(Self::from(value)); + } +} + +impl HashSetKey for T where T: Copy + Eq + Hash {} + /// Wrapper for f32 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -357,6 +385,19 @@ impl From for OrderedFloat32 { } } +impl HashSetKey for OrderedFloat32 { + fn insert(values: &mut HashSet, value: f32) { + if value == 0.0 { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + values.insert(Self(0.0)); + values.insert(Self(-0.0)); + } else { + values.insert(Self(value)); + } + } +} + /// Wrapper for f64 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -382,6 +423,19 @@ impl From for OrderedFloat64 { } } +impl HashSetKey for OrderedFloat64 { + fn insert(values: &mut HashSet, value: f64) { + if value == 0.0 { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + values.insert(Self(0.0)); + values.insert(Self(-0.0)); + } else { + values.insert(Self(value)); + } + } +} + /// Hash-set membership for primitive types. /// /// `K` defaults to the Arrow type's native value. Floats use an ordered wrapper @@ -399,7 +453,7 @@ impl PrimitiveHashSetFilter where T: ArrowPrimitiveType, T::Native: Copy, - K: From + Eq + Hash, + K: HashSetKey, { fn try_new(in_array: &ArrayRef) -> Result { let in_array = in_array.as_primitive_opt::().ok_or_else(|| { @@ -410,7 +464,7 @@ where })?; let mut values = HashSet::with_capacity(in_array.len() - in_array.null_count()); for value in in_array.iter().flatten() { - values.insert(K::from(value)); + K::insert(&mut values, value); } Ok(Self { @@ -425,7 +479,7 @@ impl StaticFilter for PrimitiveHashSetFilter where T: ArrowPrimitiveType + Send + Sync + 'static, T::Native: Copy + Send + Sync, - K: From + Eq + Hash + Send + Sync + 'static, + K: HashSetKey + Send + Sync + 'static, { fn null_count(&self) -> usize { self.null_count @@ -461,7 +515,6 @@ mod tests { DictionaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, UInt8Array, UInt16Array, UInt32Array, }; - use half::f16; use super::super::dictionary_filter::DictionaryFilter; @@ -507,6 +560,45 @@ mod tests { Ok(()) } + #[test] + fn branchless_float_zero_expansion_handles_max_list_len() -> Result<()> { + fn assert_routed_filter(haystack: ArrayRef, needles: &dyn Array) -> Result<()> { + let filter = instantiate_primitive_filter(&haystack)? + .expect("top-level float arrays always use a primitive filter"); + assert_contains( + filter.as_ref(), + needles, + vec![Some(true), Some(true), Some(false)], + ) + } + + // Adding the mirror zero to a full logical list exercises the extra + // generated comparison length for each branchless size class. + let len = ::MAX_LIST_LEN; + let mut values = vec![f16::NEG_ZERO]; + values.extend((1..len).map(|value| f16::from_f32(value as f32))); + let haystack: ArrayRef = Arc::new(Float16Array::from(values)); + let needles = + Float16Array::from(vec![f16::ZERO, f16::NEG_ZERO, f16::from_f32(100.0)]); + assert_routed_filter(haystack, &needles)?; + + let len = ::MAX_LIST_LEN; + let mut values = vec![-0.0_f32]; + values.extend((1..len).map(|value| value as f32)); + let haystack: ArrayRef = Arc::new(Float32Array::from(values)); + let needles = Float32Array::from(vec![0.0, -0.0, 100.0]); + assert_routed_filter(haystack, &needles)?; + + let len = ::MAX_LIST_LEN; + let mut values = vec![-0.0_f64]; + values.extend((1..len).map(|value| value as f64)); + let haystack: ArrayRef = Arc::new(Float64Array::from(values)); + let needles = Float64Array::from(vec![0.0, -0.0, 100.0]); + assert_routed_filter(haystack, &needles)?; + + Ok(()) + } + #[test] fn primitive_hash_filter_handles_float_keys() -> Result<()> { let nan32 = f32::NAN; @@ -524,14 +616,14 @@ mod tests { assert_contains( &filter, &needles, - vec![Some(true), Some(false), Some(true), Some(false), None], + vec![Some(true), Some(true), Some(true), Some(false), None], )?; let nan64 = f64::NAN; - let haystack: ArrayRef = Arc::new(Float64Array::from(vec![1.0, nan64])); + let haystack: ArrayRef = Arc::new(Float64Array::from(vec![-0.0, nan64])); let filter = PrimitiveHashSetFilter::::try_new(&haystack)?; - let needles = Float64Array::from(vec![Some(1.0), Some(nan64), Some(2.0)]); + let needles = Float64Array::from(vec![Some(0.0), Some(nan64), Some(2.0)]); assert_contains(&filter, &needles, vec![Some(true), Some(true), Some(false)]) } @@ -660,21 +752,22 @@ mod tests { ); let filter = BitmapFilter::::try_new(&haystack)?; let needles = Float16Array::from(vec![ + Some(f16::from_f32(9.0)), Some(f16::from_f32(0.0)), Some(f16::from_f32(-0.0)), Some(nan_a), Some(nan_b), None, ]) - .slice(1, 4); + .slice(1, 5); assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(true), Some(true), None, None]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), None, None]) ); assert_eq!( filter.contains(&needles, true)?, - BooleanArray::from(vec![Some(false), Some(false), None, None]) + BooleanArray::from(vec![Some(false), Some(false), Some(false), None, None]) ); Ok(()) diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt index 8ea1122880e14..c2c7cbb16e323 100644 --- a/datafusion/sqllogictest/test_files/negative_zero.slt +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -47,6 +47,72 @@ SELECT 0.0 IS DISTINCT FROM -0.0 AS is_distinct; ---- false +##### +## IN / NOT IN predicates +##### + +statement ok +CREATE TABLE negative_zero_in_list AS +SELECT arrow_cast(0.0, 'Float16') AS positive_f16, + arrow_cast(0.0, 'Float32') AS positive_f32, + arrow_cast(0.0, 'Float64') AS positive_f64, + arrow_cast(-0.0, 'Float16') AS negative_f16, + arrow_cast(-0.0, 'Float32') AS negative_f32, + arrow_cast(-0.0, 'Float64') AS negative_f64, + arrow_cast(-0.0, 'Dictionary(Int32, Float64)') AS negative_dict_f64; + +# A three-item list is rewritten to comparisons, while a four-item list remains +# an InList. All paths must implement the same signed-zero equality. +query BBBB +SELECT + positive_f64 IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64')), + positive_f16 IN (arrow_cast(-0.0, 'Float16'), arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'), arrow_cast(3.0, 'Float16')), + positive_f32 IN (arrow_cast(-0.0, 'Float32'), arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float32'), arrow_cast(3.0, 'Float32')), + positive_f64 IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true true true + +# Check the mirror direction for every floating-point width. +query BBB +SELECT + negative_f16 IN (arrow_cast(0.0, 'Float16'), arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'), arrow_cast(3.0, 'Float16')), + negative_f32 IN (arrow_cast(0.0, 'Float32'), arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float32'), arrow_cast(3.0, 'Float32')), + negative_f64 IN (arrow_cast(0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true true + +# NOT IN uses the same membership result before negation. +query B +SELECT + positive_f64 NOT IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +false + +# Dictionary encoding must not change the result on either side of the rewrite +# threshold. A one-item list becomes equality; a four-item list stays InList. +query BB +SELECT + negative_dict_f64 IN (arrow_cast(0.0, 'Float64')), + negative_dict_f64 IN (arrow_cast(0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true + +# A column-valued list item forces non-static evaluation. Check both directions. +query BB +SELECT + positive_f64 IN (negative_f64, arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')), + negative_f64 IN (positive_f64, arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true + +statement ok +DROP TABLE negative_zero_in_list; + ##### ## SELECT DISTINCT with +0.0 / -0.0 (Float64) #####