Skip to content

Commit 6ea6184

Browse files
perf: use null buffer union for aligned list nulls
1 parent b184e03 commit 6ea6184

1 file changed

Lines changed: 60 additions & 25 deletions

File tree

‎datafusion/functions-nested/src/arrays_zip.rs‎

Lines changed: 60 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use arrow::array::{
2222
Array, ArrayRef, Capacities, ListArray, MutableArrayData, NullBufferBuilder,
2323
StructArray, new_null_array,
2424
};
25-
use arrow::buffer::OffsetBuffer;
25+
use arrow::buffer::{NullBuffer, OffsetBuffer};
2626
use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
2727
use arrow::datatypes::{DataType, Field, Fields};
2828
use datafusion_common::cast::{
@@ -44,7 +44,7 @@ struct ListColumnView {
4444
/// Pre-computed per-row start offsets (length = num_rows + 1).
4545
offsets: Vec<usize>,
4646
/// Null bitmap from the input array (None means no nulls).
47-
nulls: Option<arrow::buffer::NullBuffer>,
47+
nulls: Option<NullBuffer>,
4848
}
4949

5050
impl ListColumnView {
@@ -383,32 +383,37 @@ fn try_perfect_list_zip(
383383
}
384384

385385
let nulls = if list_arrays.iter().any(|arr| arr.null_count() != 0) {
386-
// Match the general path: arrays_zip only marks an output row null
387-
// when every concrete input list is null. Mixed null and non-null
388-
// empty lists still produce a non-null empty list, so this cannot use
389-
// NullBuffer::union_many, which would make rows null if any input is.
390-
let mut null_builder = NullBufferBuilder::new(num_rows);
391-
for row_idx in 0..num_rows {
392-
let mut all_null = true;
393-
394-
for arr in &list_arrays {
395-
if arr.is_null(row_idx) {
396-
if arr.offsets()[row_idx + 1] != arr.offsets()[row_idx] {
397-
return Ok(None);
386+
let first_nulls = first.nulls();
387+
if list_arrays.iter().all(|arr| arr.nulls() == first_nulls) {
388+
NullBuffer::union_many(list_arrays.iter().map(|arr| arr.nulls()))
389+
} else {
390+
// Match the general path: arrays_zip only marks an output row null
391+
// when every concrete input list is null. Mixed null and non-null
392+
// empty lists still produce a non-null empty list, but mixed null
393+
// rows with values must fall back to preserve field-level nulls.
394+
let mut null_builder = NullBufferBuilder::new(num_rows);
395+
for row_idx in 0..num_rows {
396+
let mut all_null = true;
397+
398+
for arr in &list_arrays {
399+
if arr.is_null(row_idx) {
400+
if arr.offsets()[row_idx + 1] != arr.offsets()[row_idx] {
401+
return Ok(None);
402+
}
403+
} else {
404+
all_null = false;
398405
}
406+
}
407+
408+
if all_null {
409+
null_builder.append_null();
399410
} else {
400-
all_null = false;
411+
null_builder.append_non_null();
401412
}
402413
}
403414

404-
if all_null {
405-
null_builder.append_null();
406-
} else {
407-
null_builder.append_non_null();
408-
}
415+
null_builder.finish()
409416
}
410-
411-
null_builder.finish()
412417
} else {
413418
None
414419
};
@@ -555,7 +560,7 @@ mod tests {
555560
}
556561

557562
#[test]
558-
fn null_row_with_hidden_values_uses_general_path() {
563+
fn perfect_zip_reuses_null_rows_with_hidden_values() {
559564
let left =
560565
list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false]));
561566
let right = list_with_validity(
@@ -571,8 +576,38 @@ mod tests {
571576
.unwrap();
572577
let result = result.as_any().downcast_ref::<ListArray>().unwrap();
573578

574-
assert!(!result.offsets().ptr_eq(left.offsets()));
575-
assert_eq!(result.value_offsets(), &[0, 2, 2]);
579+
assert!(result.offsets().ptr_eq(left.offsets()));
580+
assert_eq!(result.value_offsets(), &[0, 2, 4]);
576581
assert!(result.is_null(1));
577582
}
583+
584+
#[test]
585+
fn mixed_null_row_with_hidden_values_uses_general_path() {
586+
let left =
587+
list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false]));
588+
let right = list_with_validity(
589+
vec![10, 20, 30, 40],
590+
vec![0, 2, 4],
591+
Some(vec![true, true]),
592+
);
593+
594+
let result = arrays_zip_inner(&[
595+
Arc::clone(&left) as ArrayRef,
596+
Arc::clone(&right) as ArrayRef,
597+
])
598+
.unwrap();
599+
let result = result.as_any().downcast_ref::<ListArray>().unwrap();
600+
let values = result
601+
.values()
602+
.as_any()
603+
.downcast_ref::<StructArray>()
604+
.unwrap();
605+
606+
assert!(!result.offsets().ptr_eq(left.offsets()));
607+
assert_eq!(result.value_offsets(), &[0, 2, 4]);
608+
assert!(values.column(0).is_null(2));
609+
assert!(values.column(0).is_null(3));
610+
assert!(!values.column(1).is_null(2));
611+
assert!(!values.column(1).is_null(3));
612+
}
578613
}

0 commit comments

Comments
 (0)