Treat signed zeros as equal in IN LIST filters

Author: geoffreyclaude

datafusion/common/src/utils/mod.rs

@@ -1407,9 +1407,9 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result<Int32Array>
     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::<Float32Type>();
             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_values() -> 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::<Int8Type>();
+        assert_eq!(dictionary.keys(), &keys);
+        let values = dictionary.values().as_primitive::<Float64Type>();
+        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<ArrayRef> = vec![

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,
 };
@@ -82,6 +83,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<dyn PhysicalExpr>],
@@ -370,12 +392,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<dyn PhysicalExpr>| -> Result<BooleanArray> {
-                    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())
@@ -3364,6 +3389,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.
@@ -3889,6 +4019,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 =
+            <Float64Type as branchless_filter::BranchlessFilterType>::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)]);