Skip to content

fix: update PrimitiveGroupValueBuilder to match NaN correctly in scalar equal_to #17979

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rluvaton merged 1 commit into from
Oct 9, 2025
Merged

fix: update PrimitiveGroupValueBuilder to match NaN correctly in scalar equal_to #17979

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
56 changes: 33 additions & 23 deletions datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
// Otherwise, we need to check their values
}

self.group_values[lhs_row] == array.as_primitive::<T>().value(rhs_row)
self.group_values[lhs_row].is_eq(array.as_primitive::<T>().value(rhs_row))
}

fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> {
Expand Down Expand Up @@ -217,22 +217,22 @@ mod tests {
use std::sync::Arc;

use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder;
use arrow::array::{ArrayRef, Int64Array, NullBufferBuilder};
use arrow::datatypes::{DataType, Int64Type};
use arrow::array::{ArrayRef, Float32Array, Int64Array, NullBufferBuilder};
use arrow::datatypes::{DataType, Float32Type, Int64Type};

use super::GroupColumn;

#[test]
fn test_nullable_primitive_equal_to() {
let append = |builder: &mut PrimitiveGroupValueBuilder<Int64Type, true>,
let append = |builder: &mut PrimitiveGroupValueBuilder<Float32Type, true>,
Copy link
Contributor

@alamb alamb Oct 9, 2025

Choose a reason for hiding this comment

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

why is this test changed?

Copy link
Member Author

@rluvaton rluvaton Oct 9, 2025

Choose a reason for hiding this comment

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

I changed the helper method type to use float32 instead of int64 so all tests in the future that uses the helper will test that as well

builder_array: &ArrayRef,
append_rows: &[usize]| {
for &index in append_rows {
builder.append_val(builder_array, index).unwrap();
}
};

let equal_to = |builder: &PrimitiveGroupValueBuilder<Int64Type, true>,
let equal_to = |builder: &PrimitiveGroupValueBuilder<Float32Type, true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
Expand All @@ -248,15 +248,15 @@ mod tests {

#[test]
fn test_nullable_primitive_vectorized_equal_to() {
let append = |builder: &mut PrimitiveGroupValueBuilder<Int64Type, true>,
let append = |builder: &mut PrimitiveGroupValueBuilder<Float32Type, true>,
Copy link
Contributor

@alamb alamb Oct 9, 2025

Choose a reason for hiding this comment

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

likewise it seems unecessary to change this test

Copy link
Member Author

@rluvaton rluvaton Oct 9, 2025

Choose a reason for hiding this comment

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

I changed the helper method type to use float32 instead of int64 so all tests in the future that uses the helper will test that as well

builder_array: &ArrayRef,
append_rows: &[usize]| {
builder
.vectorized_append(builder_array, append_rows)
.unwrap();
};

let equal_to = |builder: &PrimitiveGroupValueBuilder<Int64Type, true>,
let equal_to = |builder: &PrimitiveGroupValueBuilder<Float32Type, true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
Expand All @@ -274,9 +274,9 @@ mod tests {

fn test_nullable_primitive_equal_to_internal<A, E>(mut append: A, mut equal_to: E)
where
A: FnMut(&mut PrimitiveGroupValueBuilder<Int64Type, true>, &ArrayRef, &[usize]),
A: FnMut(&mut PrimitiveGroupValueBuilder<Float32Type, true>, &ArrayRef, &[usize]),
Copy link
Contributor

@alamb alamb Oct 9, 2025

Choose a reason for hiding this comment

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

I verified this test fails without the code change in this PR

assertion failed: equal_to_results[5]
thread 'aggregates::group_values::multi_group_by::primitive::tests::test_nullable_primitive_equal_to' panicked at datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs:346:9:
assertion failed: equal_to_results[5]
stack backtrace:
   0: __rustc::rust_begin_unwind
             at /rustc/1159e78c4747b02ef996e55082b704c09b970588/library/std/src/panicking.rs:697:5
   1: core::panicking::panic_fmt
             at /rustc/1159e78c4747b02ef996e55082b704c09b970588/library/core/src/panicking.rs:75:14
   2: core::panicking::panic
             at /rustc/1159e78c4747b02ef996e55082b704c09b970588/library/core/src/panicking.rs:145:5
   3: datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::tests::test_nullable_primitive_equal_to_internal
             at ./src/aggregates/group_values/multi_group_by/primitive.rs:346:9
   4: datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::tests::test_nullable_primitive_equal_to
             at ./src/aggregates/group_values/multi_group_by/primitive.rs:246:9
   5: datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::tests::test_nullable_primitive_equal_to::{{closure}}
             at ./src/aggregates/group_values/multi_group_by/primitive.rs:226:42
   6: core::ops::function::FnOnce::call_once
             at /Users/andrewlamb/.rustup/toolchains/1.90.0-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/ops/function.rs:253:5
   7: core::ops::function::FnOnce::call_once
             at /rustc/1159e78c4747b02ef996e55082b704c09b970588/library/core/src/ops/function.rs:253:5
note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace.

E: FnMut(
&PrimitiveGroupValueBuilder<Int64Type, true>,
&PrimitiveGroupValueBuilder<Float32Type, true>,
&[usize],
&ArrayRef,
&[usize],
Expand All @@ -293,48 +293,58 @@ mod tests {

// Define PrimitiveGroupValueBuilder
let mut builder =
PrimitiveGroupValueBuilder::<Int64Type, true>::new(DataType::Int64);
let builder_array = Arc::new(Int64Array::from(vec![
PrimitiveGroupValueBuilder::<Float32Type, true>::new(DataType::Float32);
let builder_array = Arc::new(Float32Array::from(vec![
None,
None,
None,
Some(1),
Some(2),
Some(3),
Some(1.0),
Some(2.0),
Some(f32::NAN),
Some(3.0),
])) as ArrayRef;
append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]);
append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5, 6]);

// Define input array
let (_, values, _nulls) =
Int64Array::from(vec![Some(1), Some(2), None, None, Some(1), Some(3)])
.into_parts();
let (_, values, _nulls) = Float32Array::from(vec![
Some(1.0),
Some(2.0),
None,
Some(1.0),
None,
Some(f32::NAN),
None,
])
.into_parts();

// explicitly build a null buffer where one of the null values also happens to match
let mut nulls = NullBufferBuilder::new(6);
nulls.append_non_null();
nulls.append_null(); // this sets Some(2) to null above
nulls.append_null();
nulls.append_null();
nulls.append_non_null();
nulls.append_null();
nulls.append_non_null();
let input_array = Arc::new(Int64Array::new(values, nulls.finish())) as ArrayRef;
nulls.append_null();
let input_array = Arc::new(Float32Array::new(values, nulls.finish())) as ArrayRef;

// Check
let mut equal_to_results = vec![true; builder.len()];
equal_to(
&builder,
&[0, 1, 2, 3, 4, 5],
&[0, 1, 2, 3, 4, 5, 6],
&input_array,
&[0, 1, 2, 3, 4, 5],
&[0, 1, 2, 3, 4, 5, 6],
&mut equal_to_results,
);

assert!(!equal_to_results[0]);
assert!(equal_to_results[1]);
assert!(equal_to_results[2]);
assert!(!equal_to_results[3]);
assert!(equal_to_results[3]);
assert!(!equal_to_results[4]);
assert!(equal_to_results[5]);
assert!(!equal_to_results[6]);
}

#[test]
Expand Down