Skip to content

Avoid normalizing hidden values in sliced list set operations #25300

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
neilconway merged 4 commits into from
Sep 17, 2026
+245 −93
Merged

Avoid normalizing hidden values in sliced list set operations #25300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
df3211c
fix: normalize only visible list values
yinli-systems Sep 14, 2026
e5a4c59
Merge branch 'main' into fix/sliced-set-op-normalization
yinli-systems Sep 17, 2026
284a8cf
address review
yinli-systems Sep 17, 2026
800a042
cleanup set operations
yinli-systems Sep 17, 2026
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
105 changes: 60 additions & 45 deletions datafusion/functions-nested/benches/array_set_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.

use arrow::array::{ArrayRef, Int64Array, ListArray};
use arrow::array::{ArrayRef, Float64Array, Int64Array, ListArray};
use arrow::buffer::OffsetBuffer;
use arrow::datatypes::{DataType, Field};
use criterion::{
Expand Down Expand Up @@ -47,6 +47,7 @@ fn criterion_benchmark(c: &mut Criterion) {
bench_array_union_sliced(c);
bench_array_intersect_sliced(c);
bench_array_distinct_sliced(c);
bench_array_distinct_sliced_float(c);
bench_array_except_sliced(c);
}

Expand All @@ -69,6 +70,23 @@ fn invoke_udf(udf: &impl ScalarUDFImpl, array1: &ArrayRef, array2: &ArrayRef) {
);
}

fn invoke_unary_udf(
udf: &impl ScalarUDFImpl,
array: &ArrayRef,
number_rows: usize,
) -> ColumnarValue {
black_box(
udf.invoke_with_args(ScalarFunctionArgs {
args: vec![ColumnarValue::Array(array.clone())],
arg_fields: vec![Field::new("arr", array.data_type().clone(), false).into()],
number_rows,
return_field: Field::new("result", array.data_type().clone(), false).into(),
config_options: Arc::new(ConfigOptions::default()),
})
.unwrap(),
)
}

fn bench_array_union(c: &mut Criterion) {
let mut group = c.benchmark_group("array_union");
let udf = ArrayUnion::new();
Expand Down Expand Up @@ -139,28 +157,7 @@ fn bench_array_distinct(c: &mut Criterion) {
group.bench_with_input(
BenchmarkId::new(*duplicate_label, array_size),
&array_size,
|b, _| {
b.iter(|| {
black_box(
udf.invoke_with_args(ScalarFunctionArgs {
args: vec![ColumnarValue::Array(array.clone())],
arg_fields: vec![
Field::new("arr", array.data_type().clone(), false)
.into(),
],
number_rows: NUM_ROWS,
return_field: Field::new(
"result",
array.data_type().clone(),
false,
)
.into(),
config_options: Arc::new(ConfigOptions::default()),
})
.unwrap(),
)
})
},
|b, _| b.iter(|| invoke_unary_udf(&udf, &array, NUM_ROWS)),
);
}
}
Expand Down Expand Up @@ -358,32 +355,50 @@ fn bench_array_distinct_sliced(c: &mut Criterion) {
group.bench_with_input(
BenchmarkId::from_parameter(array_size),
&array_size,
|b, _| {
b.iter(|| {
black_box(
udf.invoke_with_args(ScalarFunctionArgs {
args: vec![ColumnarValue::Array(array.clone())],
arg_fields: vec![
Field::new("arr", array.data_type().clone(), false)
.into(),
],
number_rows: NUM_ROWS,
return_field: Field::new(
"result",
array.data_type().clone(),
false,
)
.into(),
config_options: Arc::new(ConfigOptions::default()),
})
.unwrap(),
)
})
},
|b, _| b.iter(|| invoke_unary_udf(&udf, &array, NUM_ROWS)),
);
}
group.finish();
}

fn create_sliced_float_array(backing_values: usize, visible_values: usize) -> ArrayRef {
assert!(visible_values > 0 && visible_values <= backing_values);

let values = Float64Array::from(
(0..backing_values)
.map(|i| if i.is_multiple_of(2) { -0.0 } else { 0.0 })
.collect::<Vec<_>>(),
);
let left_padding = (backing_values - visible_values) / 2;
let offsets = vec![
0,
left_padding as i32,
(left_padding + visible_values) as i32,
backing_values as i32,
];
let array = ListArray::try_new(
Arc::new(Field::new("item", DataType::Float64, true)),
OffsetBuffer::new(offsets.into()),
Arc::new(values),
None,
)
.unwrap();

Arc::new(array.slice(1, 1))
}

/// Keep the visible list fixed at one `-0.0` and one `0.0` inside a much larger
/// backing values buffer to catch regressions that process values outside the slice.
fn bench_array_distinct_sliced_float(c: &mut Criterion) {
let mut group = c.benchmark_group("array_distinct_sliced_float");
let udf = ArrayDistinct::new();
let array = create_sliced_float_array(1024 * 1024, 2);

group.bench_function("backing_values_1048576_visible_values_2", |b| {
b.iter(|| invoke_unary_udf(&udf, &array, 1))
});
group.finish();
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
57 changes: 39 additions & 18 deletions datafusion/functions-nested/src/except.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

//! [`ScalarUDFImpl`] definition for array_except function.

use crate::set_ops::normalize_visible_values;
use crate::utils::{check_datatypes, make_scalar_function};
use arrow::array::new_null_array;
use arrow::array::{
Expand All @@ -27,7 +28,7 @@ use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::compute::take;
use arrow::datatypes::{DataType, FieldRef};
use arrow::row::{RowConverter, SortField};
use datafusion_common::utils::{ListCoercion, normalize_float_zero, take_function_args};
use datafusion_common::utils::{ListCoercion, take_function_args};
use datafusion_common::{HashSet, Result, internal_err};
use datafusion_expr::{
ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
Expand Down Expand Up @@ -171,24 +172,18 @@ fn general_except<OffsetSize: OffsetSizeTrait>(

// Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) groups
// ±0 together for both the rhs lookup set and the lhs probe.
let l_values_norm = normalize_float_zero(l.values());
let r_values_norm = normalize_float_zero(r.values());

// Only convert the visible portion of the values array. For sliced
// ListArrays, values() returns the full underlying array but only
// elements between the first and last offset are referenced.
let l_first = l.offsets()[0].as_usize();
let l_len = l.offsets()[l.len()].as_usize() - l_first;
let l_values = converter.convert_columns(&[l_values_norm.slice(l_first, l_len)])?;
let l_values_norm = normalize_visible_values(l);
let l_rows = converter.convert_columns(&[Arc::clone(&l_values_norm)])?;

let r_first = r.offsets()[0].as_usize();
let r_len = r.offsets()[r.len()].as_usize() - r_first;
let r_values = converter.convert_columns(&[r_values_norm.slice(r_first, r_len)])?;
let r_values_norm = normalize_visible_values(r);
let r_rows = converter.convert_columns(&[Arc::clone(&r_values_norm)])?;

let mut offsets = Vec::<OffsetSize>::with_capacity(l.len() + 1);
offsets.push(OffsetSize::usize_as(0));

let mut indices: Vec<usize> = Vec::with_capacity(l_values.num_rows());
let mut indices: Vec<usize> = Vec::with_capacity(l_rows.num_rows());
let mut dedup = HashSet::new();

let nulls = NullBuffer::union(l.nulls(), r.nulls());
Expand All @@ -207,13 +202,13 @@ fn general_except<OffsetSize: OffsetSizeTrait>(
}

for element_index in r_start.as_usize() - r_first..r_end.as_usize() - r_first {
let right_row = r_values.row(element_index);
let right_row = r_rows.row(element_index);
dedup.insert(right_row);
}
for element_index in l_start.as_usize() - l_first..l_end.as_usize() - l_first {
let left_row = l_values.row(element_index);
let left_row = l_rows.row(element_index);
if dedup.insert(left_row) {
indices.push(element_index + l_first);
indices.push(element_index);
}
}

Expand Down Expand Up @@ -245,9 +240,9 @@ fn general_except<OffsetSize: OffsetSizeTrait>(

#[cfg(test)]
mod tests {
use super::ArrayExcept;
use arrow::array::{Array, AsArray, Int32Array, ListArray};
use arrow::datatypes::{Field, Int32Type};
use super::{ArrayExcept, general_except};
use arrow::array::{Array, AsArray, Int32Array, LargeListArray, ListArray};
use arrow::datatypes::{DataType, Field, Float64Type, Int32Type};
use datafusion_common::{Result, config::ConfigOptions};
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
use std::sync::Arc;
Expand Down Expand Up @@ -302,4 +297,30 @@ mod tests {

Ok(())
}

#[test]
fn test_array_except_sliced_float_large_lists() -> Result<()> {
let l = LargeListArray::from_iter_primitive::<Float64Type, _, _>(vec![
Some(vec![Some(99.0)]),
Some(vec![Some(-0.0), Some(0.0), Some(1.0), None, None]),
Some(vec![Some(-99.0)]),
])
.slice(1, 1);
let r = LargeListArray::from_iter_primitive::<Float64Type, _, _>(vec![
Some(vec![Some(98.0)]),
Some(vec![Some(0.0), None]),
Some(vec![Some(-98.0)]),
])
.slice(1, 1);
let DataType::LargeList(field) = l.data_type() else {
unreachable!()
};

let result = general_except::<i64>(&l, &r, field)?;
let values = result.value(0);
let values = values.as_primitive::<Float64Type>();
assert_eq!(values.len(), 1);
assert_eq!(values.value(0), 1.0);
Ok(())
}
}
Loading