From 64ab2be8cf3f1942cdfb991bb0bcdb56ca108ebb Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 21:01:39 +0800 Subject: [PATCH 1/8] Trying to batch and using MultipleArrayData to optimize interleave_list --- arrow-select/src/interleave.rs | 47 +++++++++++++++++++++++++--------- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 13d28747dfcf..211dc5222ed8 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -380,6 +380,7 @@ fn interleave_list( ) -> Result { let interleaved = Interleave::<'_, GenericListArray>::new(values, indices); + // Step 1: compute output offsets and total child capacity let mut capacity = 0usize; let mut offsets = Vec::with_capacity(indices.len() + 1); offsets.push(O::from_usize(0).unwrap()); @@ -392,27 +393,49 @@ fn interleave_list( ); } - let mut child_indices = Vec::with_capacity(capacity); - for (array, row) in indices { - let list = interleaved.arrays[*array]; - let start = list.value_offsets()[*row].as_usize(); - let end = list.value_offsets()[*row + 1].as_usize(); - child_indices.extend((start..end).map(|i| (*array, i))); - } - - let child_arrays: Vec<&dyn Array> = interleaved + // Step 2: use MutableArrayData to directly copy child ranges, + // merging adjacent contiguous ranges from the same source array. + let child_data: Vec<_> = interleaved .arrays .iter() - .map(|list| list.values().as_ref()) + .map(|list| list.values().to_data()) .collect(); + let child_data_refs: Vec<_> = child_data.iter().collect(); + // nulls are derived from children. + let mut mutable_child = MutableArrayData::new(child_data_refs, /*use_nulls=*/false, capacity); - let interleaved_values = interleave(&child_arrays, &child_indices)?; + let first_offsets = interleaved.arrays[indices[0].0].value_offsets(); + let mut cur_array = indices[0].0; + let mut cur_start = first_offsets[indices[0].1].as_usize(); + let mut cur_end = first_offsets[indices[0].1 + 1].as_usize(); + + for &(array, row) in &indices[1..] { + let o = interleaved.arrays[array].value_offsets(); + let row_start = o[row].as_usize(); + let row_end = o[row + 1].as_usize(); + if array == cur_array && row_start == cur_end { + // Extend current contiguous range + cur_end = row_end; + } else { + // Flush accumulated range + if cur_end > cur_start { + mutable_child.extend(cur_array, cur_start, cur_end); + } + cur_array = array; + cur_start = row_start; + cur_end = row_end; + } + } + // Flush final range + if cur_end > cur_start { + mutable_child.extend(cur_array, cur_start, cur_end); + } let offsets = OffsetBuffer::new(offsets.into()); let list_array = GenericListArray::::new( field.clone(), offsets, - interleaved_values, + make_array(mutable_child.freeze()), interleaved.nulls, ); From e98e758eb0142ece1c96aa8492dbe9876e609630 Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 21:09:43 +0800 Subject: [PATCH 2/8] Handles basic --- arrow-select/src/interleave.rs | 34 +++++++--------------------------- 1 file changed, 7 insertions(+), 27 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 211dc5222ed8..028eaa87f2aa 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -393,43 +393,23 @@ fn interleave_list( ); } - // Step 2: use MutableArrayData to directly copy child ranges, - // merging adjacent contiguous ranges from the same source array. + // Step 2: use MutableArrayData to directly copy child ranges. let child_data: Vec<_> = interleaved .arrays .iter() .map(|list| list.values().to_data()) .collect(); let child_data_refs: Vec<_> = child_data.iter().collect(); - // nulls are derived from children. - let mut mutable_child = MutableArrayData::new(child_data_refs, /*use_nulls=*/false, capacity); - - let first_offsets = interleaved.arrays[indices[0].0].value_offsets(); - let mut cur_array = indices[0].0; - let mut cur_start = first_offsets[indices[0].1].as_usize(); - let mut cur_end = first_offsets[indices[0].1 + 1].as_usize(); + let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity); - for &(array, row) in &indices[1..] { + for &(array, row) in indices { let o = interleaved.arrays[array].value_offsets(); - let row_start = o[row].as_usize(); - let row_end = o[row + 1].as_usize(); - if array == cur_array && row_start == cur_end { - // Extend current contiguous range - cur_end = row_end; - } else { - // Flush accumulated range - if cur_end > cur_start { - mutable_child.extend(cur_array, cur_start, cur_end); - } - cur_array = array; - cur_start = row_start; - cur_end = row_end; + let start = o[row].as_usize(); + let end = o[row + 1].as_usize(); + if end > start { + mutable_child.extend(array, start, end); } } - // Flush final range - if cur_end > cur_start { - mutable_child.extend(cur_array, cur_start, cur_end); - } let offsets = OffsetBuffer::new(offsets.into()); let list_array = GenericListArray::::new( From 99761a0dc03fdd03a5fc92dec007050968f3560a Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 21:28:20 +0800 Subject: [PATCH 3/8] Revert when not primitive --- arrow-select/src/interleave.rs | 56 +++++++++++++++++++++++----------- 1 file changed, 39 insertions(+), 17 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 028eaa87f2aa..3c271d440d75 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -393,29 +393,51 @@ fn interleave_list( ); } - // Step 2: use MutableArrayData to directly copy child ranges. - let child_data: Vec<_> = interleaved - .arrays - .iter() - .map(|list| list.values().to_data()) - .collect(); - let child_data_refs: Vec<_> = child_data.iter().collect(); - let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity); - - for &(array, row) in indices { - let o = interleaved.arrays[array].value_offsets(); - let start = o[row].as_usize(); - let end = o[row + 1].as_usize(); - if end > start { - mutable_child.extend(array, start, end); + // Step 2: build child values. + // For primitive child types, use MutableArrayData to directly memcpy contiguous + // ranges, avoiding the intermediate child_indices Vec allocation. + // For complex child types (nested lists, structs, views, dictionaries, etc.), + // use recursive interleave to benefit from type-specific optimizations. + let child_values = if field.data_type().primitive_width().is_some() { + let child_data: Vec<_> = interleaved + .arrays + .iter() + .map(|list| list.values().to_data()) + .collect(); + let child_data_refs: Vec<_> = child_data.iter().collect(); + let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity); + + for &(array, row) in indices { + let o = interleaved.arrays[array].value_offsets(); + let start = o[row].as_usize(); + let end = o[row + 1].as_usize(); + if end > start { + mutable_child.extend(array, start, end); + } + } + make_array(mutable_child.freeze()) + } else { + let mut child_indices = Vec::with_capacity(capacity); + for (array, row) in indices { + let list = interleaved.arrays[*array]; + let start = list.value_offsets()[*row].as_usize(); + let end = list.value_offsets()[*row + 1].as_usize(); + child_indices.extend((start..end).map(|i| (*array, i))); } - } + + let child_arrays: Vec<&dyn Array> = interleaved + .arrays + .iter() + .map(|list| list.values().as_ref()) + .collect(); + interleave(&child_arrays, &child_indices)? + }; let offsets = OffsetBuffer::new(offsets.into()); let list_array = GenericListArray::::new( field.clone(), offsets, - make_array(mutable_child.freeze()), + child_values, interleaved.nulls, ); From 161d29fa2846cc4eac2b54e6da7317cf2c91065c Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 21:59:18 +0800 Subject: [PATCH 4/8] Type specific optimize --- arrow-select/src/interleave.rs | 126 ++++++++++++++++++++++----------- 1 file changed, 86 insertions(+), 40 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 3c271d440d75..5c988dd0efec 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -373,6 +373,64 @@ fn interleave_struct( Ok(Arc::new(struct_array)) } +/// Specialized interleave for list child arrays that are primitive. +/// Directly copies typed value slices and null bit ranges without +/// going through MutableArrayData's function pointer indirection. +fn interleave_list_primitive_child( + interleaved: &Interleave<'_, GenericListArray>, + indices: &[(usize, usize)], + capacity: usize, +) -> ArrayRef { + let child_arrays: Vec<&PrimitiveArray> = interleaved + .arrays + .iter() + .map(|list| list.values().as_primitive::()) + .collect(); + + let has_child_nulls = child_arrays.iter().any(|a| a.null_count() > 0); + + // Build values buffer by copying contiguous slices + let mut values: Vec = Vec::with_capacity(capacity); + for &(array, row) in indices { + let o = interleaved.arrays[array].value_offsets(); + let start = o[row].as_usize(); + let end = o[row + 1].as_usize(); + if end > start { + values.extend_from_slice(&child_arrays[array].values()[start..end]); + } + } + + // Build null buffer. Pre-allocate capacity so appends never resize. + // For sources without nulls: append_n sets bits to 1 (byte-level fill). + // For sources with nulls: append_packed_range copies the validity bits. + let nulls = if has_child_nulls { + let mut builder = BooleanBufferBuilder::new(capacity); + for &(array, row) in indices { + let o = interleaved.arrays[array].value_offsets(); + let start = o[row].as_usize(); + let end = o[row + 1].as_usize(); + let len = end - start; + if len > 0 { + match child_arrays[array].nulls() { + Some(null_buffer) => { + let offset = null_buffer.offset(); + builder.append_packed_range( + offset + start..offset + end, + null_buffer.validity(), + ); + } + None => builder.append_n(len, true), + } + } + } + Some(NullBuffer::new(builder.finish())) + } else { + None + }; + + Arc::new(PrimitiveArray::::new(values.into(), nulls)) +} + fn interleave_list( values: &[&dyn Array], indices: &[(usize, usize)], @@ -394,52 +452,40 @@ fn interleave_list( } // Step 2: build child values. - // For primitive child types, use MutableArrayData to directly memcpy contiguous - // ranges, avoiding the intermediate child_indices Vec allocation. - // For complex child types (nested lists, structs, views, dictionaries, etc.), - // use recursive interleave to benefit from type-specific optimizations. - let child_values = if field.data_type().primitive_width().is_some() { - let child_data: Vec<_> = interleaved - .arrays - .iter() - .map(|list| list.values().to_data()) - .collect(); - let child_data_refs: Vec<_> = child_data.iter().collect(); - let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity); + macro_rules! list_primitive_helper { + ($t:ty) => { + interleave_list_primitive_child::(&interleaved, indices, capacity) + }; + } - for &(array, row) in indices { - let o = interleaved.arrays[array].value_offsets(); - let start = o[row].as_usize(); - let end = o[row + 1].as_usize(); - if end > start { - mutable_child.extend(array, start, end); + let child_values = downcast_primitive! { + // For primitive child types, directly copy typed value slices and null bit + // ranges, avoiding both the intermediate child_indices Vec allocation and + // MutableArrayData's function pointer indirection. + field.data_type() => (list_primitive_helper), + _ => { + // For complex child types (nested lists, structs, views, dictionaries, etc.), + // use recursive interleave to benefit from type-specific optimizations. + let mut child_indices = Vec::with_capacity(capacity); + for (array, row) in indices { + let list = interleaved.arrays[*array]; + let start = list.value_offsets()[*row].as_usize(); + let end = list.value_offsets()[*row + 1].as_usize(); + child_indices.extend((start..end).map(|i| (*array, i))); } - } - make_array(mutable_child.freeze()) - } else { - let mut child_indices = Vec::with_capacity(capacity); - for (array, row) in indices { - let list = interleaved.arrays[*array]; - let start = list.value_offsets()[*row].as_usize(); - let end = list.value_offsets()[*row + 1].as_usize(); - child_indices.extend((start..end).map(|i| (*array, i))); - } - let child_arrays: Vec<&dyn Array> = interleaved - .arrays - .iter() - .map(|list| list.values().as_ref()) - .collect(); - interleave(&child_arrays, &child_indices)? + let child_arrays: Vec<&dyn Array> = interleaved + .arrays + .iter() + .map(|list| list.values().as_ref()) + .collect(); + interleave(&child_arrays, &child_indices)? + } }; let offsets = OffsetBuffer::new(offsets.into()); - let list_array = GenericListArray::::new( - field.clone(), - offsets, - child_values, - interleaved.nulls, - ); + let list_array = + GenericListArray::::new(field.clone(), offsets, child_values, interleaved.nulls); Ok(Arc::new(list_array)) } From f08e703f46d13138867c6e6fcaa08c9e7984e2ec Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 22:12:27 +0800 Subject: [PATCH 5/8] Optimize building nulls --- arrow-select/src/interleave.rs | 33 ++++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 5c988dd0efec..db964f4d78c9 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -23,6 +23,8 @@ use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; +use arrow_buffer::bit_mask::set_bits; +use arrow_buffer::bit_util; use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer}; use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; @@ -400,11 +402,15 @@ fn interleave_list_primitive_child( } } - // Build null buffer. Pre-allocate capacity so appends never resize. - // For sources without nulls: append_n sets bits to 1 (byte-level fill). - // For sources with nulls: append_packed_range copies the validity bits. + // Build null buffer. Pre-allocate with 0x00 (all null), then: + // - Sources with nulls: set_bits ORs in valid bits from source. + // - Sources without nulls: set the bit range to all 1s directly. let nulls = if has_child_nulls { - let mut builder = BooleanBufferBuilder::new(capacity); + let null_byte_len = bit_util::ceil(capacity, 8); + let mut null_buf = MutableBuffer::new(null_byte_len); + null_buf.resize(null_byte_len, 0); + + let mut offset_write = 0; for &(array, row) in indices { let o = interleaved.arrays[array].value_offsets(); let start = o[row].as_usize(); @@ -413,17 +419,26 @@ fn interleave_list_primitive_child( if len > 0 { match child_arrays[array].nulls() { Some(null_buffer) => { - let offset = null_buffer.offset(); - builder.append_packed_range( - offset + start..offset + end, + set_bits( + null_buf.as_slice_mut(), null_buffer.validity(), + offset_write, + null_buffer.offset() + start, + len, ); } - None => builder.append_n(len, true), + None => { + // Slow path. For a non-nullable source, set the bit range to all 1s directly. + let buf = null_buf.as_slice_mut(); + (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i)); + } } } + offset_write += len; } - Some(NullBuffer::new(builder.finish())) + + let bool_buf = BooleanBuffer::new(null_buf.into(), 0, capacity); + Some(NullBuffer::new(bool_buf)) } else { None }; From 270524bdf9acde41e4d4e3826d8d59103b099d02 Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 22:17:00 +0800 Subject: [PATCH 6/8] Remove some rubbish --- arrow-select/src/interleave.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index db964f4d78c9..e449e98d81d4 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -375,9 +375,6 @@ fn interleave_struct( Ok(Arc::new(struct_array)) } -/// Specialized interleave for list child arrays that are primitive. -/// Directly copies typed value slices and null bit ranges without -/// going through MutableArrayData's function pointer indirection. fn interleave_list_primitive_child( interleaved: &Interleave<'_, GenericListArray>, indices: &[(usize, usize)], From 03a5c156a815949124de98142e1c8155d92c1745 Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 27 May 2026 22:58:29 +0800 Subject: [PATCH 7/8] Add optimization for null_count --- arrow-select/src/interleave.rs | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index e449e98d81d4..8a297d9bcbc8 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -379,6 +379,7 @@ fn interleave_list_primitive_child( interleaved: &Interleave<'_, GenericListArray>, indices: &[(usize, usize)], capacity: usize, + data_type: &DataType, ) -> ArrayRef { let child_arrays: Vec<&PrimitiveArray> = interleaved .arrays @@ -408,6 +409,7 @@ fn interleave_list_primitive_child( null_buf.resize(null_byte_len, 0); let mut offset_write = 0; + let mut null_count = 0usize; for &(array, row) in indices { let o = interleaved.arrays[array].value_offsets(); let start = o[row].as_usize(); @@ -416,7 +418,7 @@ fn interleave_list_primitive_child( if len > 0 { match child_arrays[array].nulls() { Some(null_buffer) => { - set_bits( + null_count += set_bits( null_buf.as_slice_mut(), null_buffer.validity(), offset_write, @@ -425,7 +427,7 @@ fn interleave_list_primitive_child( ); } None => { - // Slow path. For a non-nullable source, set the bit range to all 1s directly. + // For a non-nullable source, set the bit range to all 1s directly. let buf = null_buf.as_slice_mut(); (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i)); } @@ -434,13 +436,18 @@ fn interleave_list_primitive_child( offset_write += len; } - let bool_buf = BooleanBuffer::new(null_buf.into(), 0, capacity); - Some(NullBuffer::new(bool_buf)) + if null_count > 0 { + let bool_buf = BooleanBuffer::new(null_buf.into(), 0, capacity); + // SAFETY: null_count is accumulated from set_bits which correctly counts unset bits + Some(unsafe { NullBuffer::new_unchecked(bool_buf, null_count) }) + } else { + None + } } else { None }; - Arc::new(PrimitiveArray::::new(values.into(), nulls)) + Arc::new(PrimitiveArray::::new(values.into(), nulls).with_data_type(data_type.clone())) } fn interleave_list( @@ -466,7 +473,12 @@ fn interleave_list( // Step 2: build child values. macro_rules! list_primitive_helper { ($t:ty) => { - interleave_list_primitive_child::(&interleaved, indices, capacity) + interleave_list_primitive_child::( + &interleaved, + indices, + capacity, + field.data_type(), + ) }; } From 63c76e48f6b3a129aa0602dba09544789b8d1656 Mon Sep 17 00:00:00 2001 From: mwish Date: Thu, 28 May 2026 20:49:47 +0800 Subject: [PATCH 8/8] Apply suggestion --- arrow-select/src/interleave.rs | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 8a297d9bcbc8..4b366e3ae1fd 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -401,15 +401,14 @@ fn interleave_list_primitive_child( } // Build null buffer. Pre-allocate with 0x00 (all null), then: - // - Sources with nulls: set_bits ORs in valid bits from source. + // - Sources with nulls: set_bits copies the source validity bits into the destination range. // - Sources without nulls: set the bit range to all 1s directly. let nulls = if has_child_nulls { let null_byte_len = bit_util::ceil(capacity, 8); - let mut null_buf = MutableBuffer::new(null_byte_len); - null_buf.resize(null_byte_len, 0); + let mut output_null_buf = MutableBuffer::from_len_zeroed(null_byte_len); let mut offset_write = 0; - let mut null_count = 0usize; + let mut output_null_count = 0usize; for &(array, row) in indices { let o = interleaved.arrays[array].value_offsets(); let start = o[row].as_usize(); @@ -418,8 +417,8 @@ fn interleave_list_primitive_child( if len > 0 { match child_arrays[array].nulls() { Some(null_buffer) => { - null_count += set_bits( - null_buf.as_slice_mut(), + output_null_count += set_bits( + output_null_buf.as_slice_mut(), null_buffer.validity(), offset_write, null_buffer.offset() + start, @@ -428,7 +427,7 @@ fn interleave_list_primitive_child( } None => { // For a non-nullable source, set the bit range to all 1s directly. - let buf = null_buf.as_slice_mut(); + let buf = output_null_buf.as_slice_mut(); (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i)); } } @@ -436,10 +435,10 @@ fn interleave_list_primitive_child( offset_write += len; } - if null_count > 0 { - let bool_buf = BooleanBuffer::new(null_buf.into(), 0, capacity); + if output_null_count > 0 { + let bool_buf = BooleanBuffer::new(output_null_buf.into(), 0, capacity); // SAFETY: null_count is accumulated from set_bits which correctly counts unset bits - Some(unsafe { NullBuffer::new_unchecked(bool_buf, null_count) }) + Some(unsafe { NullBuffer::new_unchecked(bool_buf, output_null_count) }) } else { None }