This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22285-bad2be17427ba550ddb3775a09f11dfd876de6d3 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 5472d00856a1d0f6c36afb415ba89f9defff1aeb Author: Puneet Dixit <[email protected]> AuthorDate: Thu Jun 4 15:07:14 2026 +0530 perf: optimize arrays_zip perfect list zips (#22285) ## Which issue does this PR close? - Closes #22225. ## Rationale for this change `arrays_zip` currently uses the general `MutableArrayData` path even when all regular `ListArray` inputs are already perfectly aligned. In that case, the output list offsets match the inputs and each struct child column can reuse the corresponding input values array instead of copying one row at a time. ## What changes are included in this PR? - Add a fast path for perfect regular `ListArray` zips that reuses the first input's offsets and clones the input values arrays into the output struct children. - Keep the existing general path for ragged inputs, `LargeList`, `FixedSizeList`, `Null` inputs, and null rows that would require padding. - Add unit coverage for offset/value reuse, zero-length null rows, and null rows with hidden values falling back to the general path. - Rename the no-null benchmark case to `arrays_zip_perfect_zip_8192`. ## Are these changes tested? Yes. Local checks run: - `cargo fmt --all` - `cargo test -p datafusion-functions-nested` - `cargo clippy -p datafusion-functions-nested --all-targets --all-features -- -D warnings` - `CARGO_TARGET_DIR=C:\df-target cargo clippy --all-targets --all-features -- -D warnings` - `CARGO_TARGET_DIR=C:\df-target cargo bench -p datafusion-functions-nested --bench arrays_zip -- --warm-up-time 1 --measurement-time 2 --sample-size 10` Latest local benchmark sample: - `arrays_zip_perfect_zip_8192`: `11.234 ?s 11.600 ?s 12.045 ?s` - `arrays_zip_10pct_nulls_8192`: `4.3463 ms 4.5531 ms 4.7898 ms` ## Are there any user-facing changes? No. This is an internal performance optimization with the same `arrays_zip` output semantics. ## Review follow-up validation 2026-05-21: - `cargo fmt --all --check` passed after formatting. - `cargo test -p datafusion-functions-nested perfect_zip_uses_supplied_field_names --lib` could not compile locally because this Windows GNU toolchain is missing `dlltool.exe`. - `cargo clippy -p datafusion-functions-nested --all-targets -- -D warnings` hit the same local `dlltool.exe` toolchain blocker before checking the crate. --------- Co-authored-by: Puneet Dixit <[email protected]> Co-authored-by: Copilot Autofix powered by AI <[email protected]> Co-authored-by: puneetdixit200 <[email protected]> --- datafusion/functions-nested/benches/arrays_zip.rs | 2 +- datafusion/functions-nested/src/arrays_zip.rs | 294 +++++++++++++++++++++- 2 files changed, 290 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-nested/benches/arrays_zip.rs b/datafusion/functions-nested/benches/arrays_zip.rs index bc82b2978c..812e5e3dbe 100644 --- a/datafusion/functions-nested/benches/arrays_zip.rs +++ b/datafusion/functions-nested/benches/arrays_zip.rs @@ -109,7 +109,7 @@ fn bench_arrays_zip(c: &mut Criterion, name: &str, null_density: f64) { } fn criterion_benchmark(c: &mut Criterion) { - bench_arrays_zip(c, "arrays_zip_no_nulls_8192", 0.0); + bench_arrays_zip(c, "arrays_zip_perfect_zip_8192", 0.0); bench_arrays_zip(c, "arrays_zip_10pct_nulls_8192", 0.1); } diff --git a/datafusion/functions-nested/src/arrays_zip.rs b/datafusion/functions-nested/src/arrays_zip.rs index 5f1cb9dedf..76b1b589f4 100644 --- a/datafusion/functions-nested/src/arrays_zip.rs +++ b/datafusion/functions-nested/src/arrays_zip.rs @@ -22,7 +22,7 @@ use arrow::array::{ Array, ArrayRef, Capacities, ListArray, MutableArrayData, NullBufferBuilder, StructArray, new_null_array, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null}; use arrow::datatypes::{DataType, Field, Fields}; use datafusion_common::cast::{ @@ -44,7 +44,7 @@ struct ListColumnView { /// Pre-computed per-row start offsets (length = num_rows + 1). offsets: Vec<usize>, /// Null bitmap from the input array (None means no nulls). - nulls: Option<arrow::buffer::NullBuffer>, + nulls: Option<NullBuffer>, } impl ListColumnView { @@ -130,7 +130,7 @@ impl ScalarUDFImpl for ArraysZip { return exec_err!("arrays_zip expects array arguments, got {dt}"); } }; - fields.push(Field::new(format!("{}", i + 1), element_type, true)); + fields.push(Field::new(arrays_zip_field_name(i), element_type, true)); } Ok(List(Arc::new(Field::new_list_field( @@ -163,8 +163,13 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result<ArrayRef> { return exec_err!("arrays_zip requires at least one argument"); } + let field_names = arrays_zip_field_names(args.len()); let num_rows = args[0].len(); + if let Some(result) = try_perfect_list_zip(args, &field_names)? { + return Ok(result); + } + // Build a type-erased ListColumnView for each argument. // None means the argument is Null-typed (all nulls, no backing data). let mut views: Vec<Option<ListColumnView>> = Vec::with_capacity(args.len()); @@ -225,8 +230,8 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result<ArrayRef> { let struct_fields: Fields = element_types .iter() - .enumerate() - .map(|(i, dt)| Field::new(format!("{}", i + 1), dt.clone(), true)) + .zip(field_names.iter()) + .map(|(dt, name)| Field::new(name.clone(), dt.clone(), true)) .collect::<Vec<_>>() .into(); @@ -327,3 +332,282 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result<ArrayRef> { Ok(Arc::new(result)) } + +fn arrays_zip_field_name(index: usize) -> String { + (index + 1).to_string() +} + +fn arrays_zip_field_names(len: usize) -> Vec<String> { + (0..len).map(arrays_zip_field_name).collect() +} + +/// Fast path for regular List inputs whose existing buffers already match the +/// zipped output: all offsets and values lengths match, and null rows cover no +/// values. This lets us reuse offsets and child values instead of rebuilding. +fn try_perfect_list_zip( + args: &[ArrayRef], + field_names: &[String], +) -> Result<Option<ArrayRef>> { + debug_assert_eq!(args.len(), field_names.len()); + + let mut list_arrays = Vec::with_capacity(args.len()); + let mut struct_fields = Vec::with_capacity(args.len()); + + for (arg, field_name) in args.iter().zip(field_names) { + let arr = match arg.data_type() { + List(field) => { + struct_fields.push(Field::new( + field_name.clone(), + field.data_type().clone(), + true, + )); + as_list_array(arg)? + } + _ => return Ok(None), + }; + + list_arrays.push(arr); + } + + let first = list_arrays[0]; + let num_rows = first.len(); + let offsets = first.offsets().clone(); + let values_len = first.values().len(); + + // Reusing the child arrays is only valid when every list uses the exact + // same row boundaries and exposes the same total number of child values. + for arr in &list_arrays { + if arr.values().len() != values_len || arr.offsets() != &offsets { + return Ok(None); + } + } + + let nulls = if list_arrays.iter().any(|arr| arr.null_count() != 0) { + let first_nulls = first.nulls(); + if list_arrays.iter().all(|arr| arr.nulls() == first_nulls) { + first_nulls.cloned() + } else { + // Match the general path: arrays_zip only marks an output row null + // when every concrete input list is null. Mixed null and non-null + // empty lists still produce a non-null empty list, but mixed null + // rows with values must fall back to preserve field-level nulls. + let mut null_builder = NullBufferBuilder::new(num_rows); + for row_idx in 0..num_rows { + let mut all_null = true; + + for arr in &list_arrays { + if arr.is_null(row_idx) { + if arr.offsets()[row_idx + 1] != arr.offsets()[row_idx] { + return Ok(None); + } + } else { + all_null = false; + } + } + + if all_null { + null_builder.append_null(); + } else { + null_builder.append_non_null(); + } + } + + null_builder.finish() + } + } else { + None + }; + + let struct_columns = list_arrays + .iter() + .map(|arr| Arc::clone(arr.values())) + .collect::<Vec<_>>(); + let struct_array = + StructArray::try_new(Fields::from(struct_fields), struct_columns, None)?; + let result = ListArray::try_new( + Arc::new(Field::new_list_field( + struct_array.data_type().clone(), + true, + )), + offsets, + Arc::new(struct_array), + nulls, + )?; + + Ok(Some(Arc::new(result))) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::buffer::NullBuffer; + + fn list(values: Vec<i64>, offsets: Vec<i32>) -> Arc<ListArray> { + list_with_validity(values, offsets, None) + } + + fn list_with_validity( + values: Vec<i64>, + offsets: Vec<i32>, + valid: Option<Vec<bool>>, + ) -> Arc<ListArray> { + Arc::new( + ListArray::try_new( + Arc::new(Field::new_list_field(DataType::Int64, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(Int64Array::from(values)), + valid.map(NullBuffer::from), + ) + .unwrap(), + ) + } + + #[test] + fn perfect_zip_reuses_input_values_and_offsets() { + let left = list(vec![1, 2, 3, 4, 5, 6], vec![0, 2, 3, 6]); + let right = list(vec![10, 20, 30, 40, 50, 60], vec![0, 2, 3, 6]); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(Arc::ptr_eq(values.column(0), left.values())); + assert!(Arc::ptr_eq(values.column(1), right.values())); + } + + #[test] + fn perfect_zip_uses_supplied_field_names() { + let left = list(vec![1, 2, 3], vec![0, 1, 3]); + let right = list(vec![10, 20, 30], vec![0, 1, 3]); + let field_names = vec!["left".to_string(), "right".to_string()]; + + let result = try_perfect_list_zip( + &[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ], + &field_names, + ) + .unwrap() + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + let names = values + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::<Vec<_>>(); + + assert_eq!(names, vec!["left", "right"]); + } + + #[test] + fn perfect_zip_reuses_zero_length_null_rows() { + let left = list_with_validity( + vec![1, 2, 3, 4], + vec![0, 2, 2, 4], + Some(vec![true, false, true]), + ); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 2, 4], + Some(vec![true, false, true]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(result.is_null(1)); + } + + #[test] + fn perfect_zip_preserves_mixed_null_empty_rows() { + let left = + list_with_validity(vec![], vec![0, 0, 0, 0], Some(vec![false, true, false])); + let right = + list_with_validity(vec![], vec![0, 0, 0, 0], Some(vec![true, false, false])); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(!result.is_null(0)); + assert!(!result.is_null(1)); + assert!(result.is_null(2)); + } + + #[test] + fn perfect_zip_reuses_null_rows_with_hidden_values() { + let left = + list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false])); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 4], + Some(vec![true, false]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert_eq!(result.value_offsets(), &[0, 2, 4]); + assert!(result.is_null(1)); + } + + #[test] + fn mixed_null_row_with_hidden_values_uses_general_path() { + let left = + list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false])); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 4], + Some(vec![true, true]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::<ListArray>().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + + assert!(!result.offsets().ptr_eq(left.offsets())); + assert_eq!(result.value_offsets(), &[0, 2, 4]); + assert!(values.column(0).is_null(2)); + assert!(values.column(0).is_null(3)); + assert!(!values.column(1).is_null(2)); + assert!(!values.column(1).is_null(3)); + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
