Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 102 additions & 22 deletions datafusion/common/src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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 {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only NEG_ZERO_F32_BITS needs a change, so compare directly against it

0.0_f32
} else {
v
}
});
Arc::new(normalized)
}
DataType::Float64 => {
Expand All @@ -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 {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only NEG_ZERO_F64_BITS needs a change, so compare directly against it

0.0_f64
} else {
v
}
});
Arc::new(normalized)
}
DataType::Float16 => {
Expand All @@ -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
}
Expand All @@ -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)]
Expand All @@ -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;
Expand Down Expand Up @@ -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::<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![
Expand Down
157 changes: 155 additions & 2 deletions datafusion/physical-expr/src/expressions/in_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down Expand Up @@ -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<dyn PhysicalExpr>],
Expand Down Expand Up @@ -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<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())
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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 =
<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)]);
Expand Down
Loading
Loading