kosiew commented on code in PR #25754:
URL: https://github.com/apache/datafusion/pull/25754#discussion_r4282017165
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -529,6 +579,59 @@ pub(crate) mod tests {
use arrow::datatypes::Int64Type;
use datafusion_common::utils::SingleRowListArrayBuilder;
+ /// Compare `prepare_list_filter` with a per-row reference on a sliced list
+ /// with the given row lengths and validity, returning the input and output
+ /// bitmaps so callers can check whether the input was reused.
+ fn check_list_filter<O: OffsetSizeTrait>(
+ lengths: [usize; 8],
+ valid: [bool; 8],
+ ) -> (BooleanBuffer, BooleanBuffer) {
+ let list = GenericListArray::<O>::new(
+ Arc::new(Field::new_list_field(DataType::Int32, true)),
+ OffsetBuffer::from_lengths(lengths),
+ // The helper must permit selecting NULL elements.
+ arrow::array::new_null_array(&DataType::Int32,
lengths.iter().sum()),
+ Some(NullBuffer::from(valid.to_vec())),
+ )
+ .slice(1, 7);
+ let (first, len) = offset_span(list.offsets());
+ // Use an offset that is not a multiple of the bitmap's three-bit
period.
+ let keep = BooleanBuffer::collect_bool(len + 1, |i| i % 3 !=
0).slice(1, len);
+ let mut expected = keep.iter().collect::<Vec<_>>();
+ let mut expected_offsets = vec![0];
+ for (row, window) in list.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first;
+ let end = window[1].as_usize() - first;
+ if list.is_null(row) {
+ expected[start..end].fill(false);
+ }
+ expected_offsets.push(expected[..end].iter().filter(|&&v|
v).count());
+ }
+ let (offsets, mask) = prepare_list_filter(&list, &keep);
+ assert_eq!(mask.iter().collect::<Vec<_>>(), expected);
+ assert_eq!(
+ offsets.iter().map(|o| o.as_usize()).collect::<Vec<_>>(),
+ expected_offsets
+ );
+ (keep, mask)
+ }
+
+ #[test]
+ fn list_filter_mask_handles_slices_and_hidden_elements() {
+ let valid = [true, false, true, false, false, true, true, false];
+ // After slicing off the first row: an empty NULL row, a NULL row
storing
+ // an element after a valid row, consecutive NULL rows, and a trailing
one.
+ let hidden = [2, 0, 2, 1, 2, 0, 3, 2];
Review Comment:
Could we add a test case that crosses 64-bit bitmap boundaries for both
offset types? For example, lengths `[2, 0, 63, 65, 64, 0, 65, 2]` with the
existing validity array would exercise sliced bitmaps and hidden elements
across word boundaries. The current cases cover only short bitmaps, so this
would give us more confidence in the packed range copying.
##########
datafusion/functions-nested/benches/array_filter.rs:
##########
@@ -0,0 +1,371 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::{hint::black_box, sync::Arc};
+
+use arrow::array::{
+ Array, ArrayRef, Int32Array, ListArray, RecordBatch, StringArray,
StringViewArray,
+ UInt64Array,
+};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::compute::{cast, take};
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{Criterion, SamplingMode, criterion_group, criterion_main};
+use datafusion_common::{DFSchema, ScalarValue};
+use datafusion_expr::{
+ Expr, col, execution_props::ExecutionProps, expr::LambdaVariable, lambda,
lit,
+ physical_planning_context::PhysicalPlanningContext,
+};
+use datafusion_functions_nested::expr_fn::array_filter;
+use datafusion_physical_expr::create_physical_expr;
+use rand::{Rng, SeedableRng, rngs::StdRng};
+
+// Inputs are built by hand rather than from SQL literals so the bench can
cover
+// sliced arrays and elements stored under NULL list rows.
+// Defaults: 256 rows of 16 elements, seeded random values, no NULL list rows.
+struct Case {
+ kind: DataType,
+ rows: usize,
+ width: usize,
+ element_nulls: f64,
+ row_nulls: f64,
+ // NULL list rows keep their elements, which evaluation must strip.
+ hidden_storage: bool,
+ predicate: Predicate,
+ variable: bool,
+ sliced: bool,
+ large_list: bool,
+}
+
+#[derive(Clone, Copy)]
+enum Predicate {
+ LessThan(i32),
+ Scalar(Option<bool>),
+ IsNull,
+ Captured,
+}
+
+impl Default for Case {
+ fn default() -> Self {
+ Self {
+ kind: DataType::Int32,
+ rows: 256,
+ width: 16,
+ element_nulls: 0.0,
+ row_nulls: 0.0,
+ hidden_storage: false,
+ predicate: Predicate::LessThan(50),
+ variable: false,
+ sliced: false,
+ large_list: false,
+ }
+ }
+}
+
+fn string_value(value: i32) -> String {
+ format!("{value:03}-array-filter-value")
+}
+
+// Evaluate a planned projection, including lambda evaluation, on an in-memory
+// batch. Constructing inputs and planning are outside the timed loop.
+fn bench_case(c: &mut Criterion, name: &str, case: &Case) {
+ let mut rng = StdRng::seed_from_u64(42);
+ let row_valid = (0..case.rows)
+ .map(|_| !rng.random_bool(case.row_nulls))
+ .collect::<Vec<_>>();
+ let lengths = (0..case.rows)
+ .map(|row| {
+ if !row_valid[row] && !case.hidden_storage {
+ 0
+ } else if case.variable {
+ (row * 17) % (2 * case.width + 1)
+ } else {
+ case.width
+ }
+ })
+ .collect::<Vec<_>>();
+ let values = (0..lengths.iter().sum())
+ .map(|_| {
+ let value = rng.random_range(0..100);
+ (!rng.random_bool(case.element_nulls)).then_some(value)
+ })
+ .collect::<Vec<_>>();
+ let elements: ArrayRef = match case.kind {
+ DataType::Int32 => Arc::new(Int32Array::from(values.clone())),
+ DataType::Utf8 => Arc::new(StringArray::from_iter(
+ values.iter().map(|v| v.map(string_value)),
+ )),
+ DataType::Utf8View => Arc::new(StringViewArray::from_iter(
+ values.iter().map(|v| v.map(string_value)),
+ )),
+ _ => unreachable!(),
+ };
+ let row_nulls = (case.row_nulls != 0.0).then(||
NullBuffer::from(row_valid.clone()));
+ let field = Arc::new(Field::new_list_field(case.kind.clone(), true));
+ let list = ListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::from_lengths(lengths),
+ Arc::clone(&elements),
+ row_nulls,
+ );
+ let list = if case.sliced {
+ list.slice(1, case.rows - 2)
+ } else {
+ list
+ };
+ let first_row = usize::from(case.sliced);
+ let thresholds = (first_row..first_row + list.len())
+ .map(|row| (row % 100) as i32)
+ .collect::<Vec<_>>();
+
+ // Independently check logical output before timing, including NULL
elements
+ // deliberately retained by IS NULL and values hidden by NULL list rows.
+ let mut selected = Vec::new();
+ let mut result_lengths = Vec::new();
+ for (row, offsets) in list.offsets().windows(2).enumerate() {
+ let before = selected.len();
+ if row_valid[first_row + row] {
+ let (start, end) = (offsets[0] as usize, offsets[1] as usize);
+ for (index, value) in (start..).zip(&values[start..end]) {
+ let keep = match case.predicate {
+ Predicate::LessThan(limit) => value.is_some_and(|v| v <
limit),
+ Predicate::Scalar(value) => value == Some(true),
+ Predicate::IsNull => value.is_none(),
+ Predicate::Captured => value.is_some_and(|v| v <
thresholds[row]),
+ };
+ if keep {
+ selected.push(index as u64);
+ }
+ }
+ }
+ result_lengths.push(selected.len() - before);
+ }
+ let expected = ListArray::new(
+ Arc::clone(&field),
+ OffsetBuffer::<i32>::from_lengths(result_lengths),
+ take(elements.as_ref(), &UInt64Array::from(selected), None).unwrap(),
+ list.nulls().cloned(),
+ );
+ let mut input = Arc::new(list) as ArrayRef;
+ let mut expected = Arc::new(expected) as ArrayRef;
+ if case.large_list {
+ let data_type = DataType::LargeList(Arc::clone(&field));
+ input = cast(&input, &data_type).unwrap();
+ expected = cast(&expected, &data_type).unwrap();
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("list", input.data_type().clone(), true),
+ Field::new("threshold", DataType::Int32, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![input, Arc::new(Int32Array::from(thresholds))],
+ )
+ .unwrap();
+ let v = Expr::LambdaVariable(LambdaVariable::new("v".into(), Some(field)));
+ let body = match case.predicate {
+ Predicate::LessThan(limit) => {
+ let limit = if case.kind == DataType::Int32 {
+ lit(limit)
+ } else {
+ lit(ScalarValue::new_utf8(string_value(limit))
+ .cast_to(&case.kind)
+ .unwrap())
+ };
+ v.lt(limit)
+ }
+ Predicate::Scalar(value) => lit(ScalarValue::Boolean(value)),
+ Predicate::IsNull => v.is_null(),
+ Predicate::Captured => v.lt(col("threshold")),
+ };
+ let expression = create_physical_expr(
+ &array_filter(col("list"), lambda(["v"], body)),
+ &DFSchema::try_from(schema).unwrap(),
+ &ExecutionProps::new(),
+ &PhysicalPlanningContext::default(),
+ )
+ .unwrap();
+ let output = expression
+ .evaluate(&batch)
+ .unwrap()
+ .into_array(batch.num_rows())
+ .unwrap();
+ assert_eq!(output.as_ref(), expected.as_ref(), "{name}");
+ output.to_data().validate_full().unwrap();
+
+ let mut group = c.benchmark_group("array_filter");
+ group.sampling_mode(SamplingMode::Flat);
+ group.bench_function(name, |b| {
+ b.iter(|| black_box(expression.evaluate(black_box(&batch)).unwrap()))
+ });
+ group.finish();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ for (kind, name) in [
+ (DataType::Int32, "int32"),
+ (DataType::Utf8, "utf8"),
+ (DataType::Utf8View, "utf8_view"),
+ ] {
+ for keep in [10, 50, 90] {
+ bench_case(
+ c,
+ &format!("{name}/keep{keep}/no_nulls"),
+ &Case {
+ kind: kind.clone(),
+ predicate: Predicate::LessThan(keep),
+ ..Case::default()
+ },
+ );
+ }
+ // 25% NULL elements make comparison results nullable. Half of the
+ // remaining elements pass, so approximately 37.5% of all slots
survive.
+ bench_case(
+ c,
+ &format!("{name}/keep50/null_elements25"),
+ &Case {
+ kind: kind.clone(),
+ element_nulls: 0.25,
+ ..Case::default()
+ },
+ );
+ // IS NULL is cheap to evaluate, so list reconstruction dominates.
+ bench_case(
+ c,
+ &format!("{name}/keep_null_elements"),
+ &Case {
+ kind: kind.clone(),
+ predicate: Predicate::IsNull,
+ element_nulls: 0.25,
+ ..Case::default()
+ },
+ );
+ for percent in [50, 99] {
+ bench_case(
+ c,
+ &format!("{name}/keep50/null_rows{percent}"),
+ &Case {
+ kind: kind.clone(),
+ row_nulls: f64::from(percent) / 100.0,
+ element_nulls: 0.25,
+ ..Case::default()
+ },
+ );
+ }
+ }
+ for (name, case) in [
Review Comment:
Could we also benchmark short Int32 lists, such as 2,048 rows with 2
elements each and some NULL elements? The current numeric benchmarks use width
16, so a narrower case would help measure per-row bookkeeping and popcount
overhead when there are many rows relative to the total number of elements.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
