From 49f5351e7d88f564385fb6a0ef16ba627637de79 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 30 Jan 2026 13:41:18 +0000 Subject: [PATCH 1/2] performance: Use push_unchecked when building indices for element take in fixedsizedlist take Signed-off-by: Robert Kruszewski --- .../arrays/fixed_size_list/compute/take.rs | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 68432e2d798..a55314a8fae 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; use vortex_dtype::IntegerPType; -use vortex_dtype::Nullability; use vortex_dtype::match_each_integer_ptype; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -16,8 +16,6 @@ use crate::ToCanonical; use crate::arrays::FixedSizeListArray; use crate::arrays::FixedSizeListVTable; use crate::arrays::PrimitiveArray; -use crate::builders::ArrayBuilder; -use crate::builders::PrimitiveBuilder; use crate::compute::TakeKernel; use crate::compute::TakeKernelAdapter; use crate::compute::{self}; @@ -92,8 +90,7 @@ fn take_non_nullable_fsl( let new_len = indices.len(); // Build the element indices directly without validity tracking. - let mut elements_indices = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, new_len * list_size); + let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); // Build the element indices for each list. for data_idx in indices { @@ -106,14 +103,17 @@ fn take_non_nullable_fsl( // Expand the list into individual element indices. for i in list_start..list_end { - elements_indices.append_value(I::from_usize(i).vortex_expect("i < list_end")) + unsafe { + elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end")) + }; } } - let elements_indices = elements_indices.finish(); + let elements_indices = elements_indices.freeze(); debug_assert_eq!(elements_indices.len(), new_len * list_size); - let new_elements = compute::take(array.elements(), elements_indices.as_ref())?; + let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); + let new_elements = compute::take(array.elements(), elements_indices_array.as_ref())?; debug_assert_eq!(new_elements.len(), new_len * list_size); // Both inputs are non-nullable, so the result is non-nullable. @@ -142,8 +142,7 @@ fn take_nullable_fsl( // We must use placeholder zeros for null lists to maintain the array length without // propagating nullability to the element array's take operation. - let mut elements_indices = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, new_len * list_size); + let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); let mut new_validity_builder = BitBufferMut::with_capacity(new_len); // Build the element indices while tracking which lists are null. @@ -158,7 +157,7 @@ fn take_nullable_fsl( if !is_index_valid || !array_validity.value(data_idx) { // Append placeholder zeros for null lists. These will be masked by the validity array. // We cannot use append_nulls here as explained above. - elements_indices.append_zeros(list_size); + unsafe { elements_indices.push_n_unchecked(I::zero(), list_size) }; new_validity_builder.append(false); } else { // Append the actual element indices for this list. @@ -167,17 +166,20 @@ fn take_nullable_fsl( // Expand the list into individual element indices. for i in list_start..list_end { - elements_indices.append_value(I::from_usize(i).vortex_expect("i < list_end")) + unsafe { + elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end")) + }; } new_validity_builder.append(true); } } - let elements_indices = elements_indices.finish(); + let elements_indices = elements_indices.freeze(); debug_assert_eq!(elements_indices.len(), new_len * list_size); - let new_elements = compute::take(array.elements(), elements_indices.as_ref())?; + let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); + let new_elements = compute::take(array.elements(), elements_indices_array.as_ref())?; debug_assert_eq!(new_elements.len(), new_len * list_size); // At least one input was nullable, so the result is nullable. From 4ab43a1b025bfe126eda9241db107d3ce4b0874e Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sun, 1 Feb 2026 21:39:49 -0500 Subject: [PATCH 2/2] comment Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/fixed_size_list/compute/take.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index a55314a8fae..f4ae9ad9b21 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -103,6 +103,7 @@ fn take_non_nullable_fsl( // Expand the list into individual element indices. for i in list_start..list_end { + // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. unsafe { elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end")) }; @@ -166,6 +167,7 @@ fn take_nullable_fsl( // Expand the list into individual element indices. for i in list_start..list_end { + // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. unsafe { elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end")) };