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

Open
yinli-systems wants to merge 1 commit into
base: main
Choose a base branch
from
+299 −92
Open

Avoid normalizing hidden values in sliced list set operations #25300

Show file tree
Hide file tree
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
142 changes: 97 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 All @@ -38,6 +38,11 @@ const SEED: u64 = 42;
/// Extra rows on each side when building sliced arrays, so the underlying
/// values buffer is much larger than the visible portion.
const SLICE_PADDING: usize = 5000;
/// Keep the visible slice at two values while varying the backing array from
/// 8 KiB to 8 MiB.
const SLICED_FLOAT_BACKING_VALUES: &[usize] = &[1024, 1024 * 1024];
/// Keep the backing array at 8 MiB while varying the visible slice.
const SLICED_FLOAT_VISIBLE_VALUES: &[usize] = &[2, 2048];

fn criterion_benchmark(c: &mut Criterion) {
bench_array_union(c);
Expand All @@ -47,6 +52,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 +75,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 +162,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,30 +360,80 @@ 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))
}

fn create_unsliced_float_array(value_count: usize) -> ArrayRef {
let values = Float64Array::from(
(0..value_count)
.map(|i| if i.is_multiple_of(2) { -0.0 } else { 0.0 })
.collect::<Vec<_>>(),
);
Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Float64, true)),
OffsetBuffer::new(vec![0, value_count as i32].into()),
Arc::new(values),
None,
))
}

/// Keep the visible list fixed at one `-0.0` and one `0.0` while increasing
/// the backing values buffer. This isolates work outside the logical slice.
fn bench_array_distinct_sliced_float(c: &mut Criterion) {
let mut group = c.benchmark_group("array_distinct_sliced_float");
let udf = ArrayDistinct::new();

for &backing_values in SLICED_FLOAT_BACKING_VALUES {
let array = create_sliced_float_array(backing_values, 2);
group.bench_with_input(
BenchmarkId::new("backing_values", backing_values),
&backing_values,
|b, _| b.iter(|| invoke_unary_udf(&udf, &array, 1)),
);
}

for &visible_values in SLICED_FLOAT_VISIBLE_VALUES {
let array = create_sliced_float_array(1024 * 1024, visible_values);
group.bench_with_input(
BenchmarkId::new("visible_values", visible_values),
&visible_values,
|b, _| b.iter(|| invoke_unary_udf(&udf, &array, 1)),
);
}

let array = create_unsliced_float_array(1024);
group.bench_function("unsliced_values_1024", |b| {
b.iter(|| invoke_unary_udf(&udf, &array, 1))
});
group.finish();
}

Expand Down
66 changes: 49 additions & 17 deletions datafusion/functions-nested/src/except.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,26 +169,32 @@ fn general_except<OffsetSize: OffsetSizeTrait>(
) -> Result<GenericListArray<OffsetSize>> {
let converter = RowConverter::new(vec![SortField::new(l.value_type())])?;

// 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.
// ListArray::values() returns the full underlying array for sliced lists.
// Slice first so normalization only scans and, when -0.0 is present,
// allocates for values referenced by the logical array.
// Normalization keeps SQL signed-zero equality when rows are encoded.
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 = if l_first == 0 && l_len == l.values().len() {
normalize_float_zero(l.values())
} else {
normalize_float_zero(&l.values().slice(l_first, l_len))
};
let l_rows = converter.convert_columns(&[l_values_norm.slice(0, l_len)])?;

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 = if r_first == 0 && r_len == r.values().len() {
normalize_float_zero(r.values())
} else {
normalize_float_zero(&r.values().slice(r_first, r_len))
};
let r_rows = converter.convert_columns(&[r_values_norm.slice(0, r_len)])?;

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 +213,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 +251,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 +308,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