From 97a3d565c72e35b50549d34a34a65ae789460024 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 13:06:12 -0400 Subject: [PATCH 1/2] rebuild lists over only the elements valid lists reference Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/mod.rs | 2 + vortex-array/src/arrays/list/rebuild.rs | 109 ++++++++++++++++++++ vortex-array/src/arrays/list/tests.rs | 48 +++++++++ vortex-array/src/arrays/listview/rebuild.rs | 109 ++++++++++++++++++++ vortex-array/src/validity.rs | 18 ++++ 5 files changed, 286 insertions(+) create mode 100644 vortex-array/src/arrays/list/rebuild.rs diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index bfd43cdb401..7c604852b59 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -12,6 +12,8 @@ pub use vtable::ListArray; pub(crate) mod compute; +mod rebuild; + mod vtable; pub use vtable::List; diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs new file mode 100644 index 00000000000..622631b8c18 --- /dev/null +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ListArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::dtype::IntegerPType; +use crate::match_each_integer_ptype; + +impl ListArray { + /// Rebuilds the [`ListArray`] over only the elements that valid lists reference: nothing before + /// the first list, after the last or under a null list. + /// + /// The elements are sliced when only the ends go, and copied a run of valid lists at a time + /// when a null list holds elements. An array with nothing to drop is returned as it is. + pub fn rebuild_with_referenced_valid_elements( + &self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let offsets = self.offsets().clone().execute::(ctx)?; + let nulls = self.list_validity().execute_null_runs(self.len(), ctx)?; + match_each_integer_ptype!(offsets.ptype(), |O| { + self.rebuild_without_null_runs(offsets.as_slice::(), &nulls) + }) + } + + fn rebuild_without_null_runs( + &self, + offsets: &[O], + nulls: &[(usize, usize)], + ) -> VortexResult { + let kept = referenced_valid_ranges(offsets, nulls); + if kept.iter().map(|(start, end)| end - start).sum::() == self.elements().len() { + return Ok(self.clone()); + } + let elements = match kept.as_slice() { + [] => self.elements().slice(0..0)?, + [(start, end)] => self.elements().slice(*start..*end)?, + _ => { + let mut bits = BitBufferMut::new_unset(self.elements().len()); + for &(start, end) in &kept { + bits.fill_range(start, end, true); + } + // With the ranges cached on the mask, the filter copies one range at a time. + self.elements() + .filter(Mask::from_buffer_with_slices(bits.freeze(), kept))? + } + }; + let offsets = shift_past_null_runs(offsets, nulls); + // SAFETY: the elements are the ranges of the runs of valid lists, in order, and every offset + // moved down by the elements dropped before it, so each valid list keeps its elements and + // each null list is empty. + Ok(unsafe { + ListArray::new_unchecked(elements, offsets.into_array(), self.list_validity()) + }) + } +} + +/// The element ranges of the runs of valid lists, with touching ranges merged. +fn referenced_valid_ranges( + offsets: &[O], + nulls: &[(usize, usize)], +) -> Vec<(usize, usize)> { + let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(nulls.len() + 1); + let mut keep = |start: O, end: O| { + let (start, end): (usize, usize) = (start.as_(), end.as_()); + match ranges.last_mut() { + _ if start == end => {} + Some(last) if last.1 == start => last.1 = end, + _ => ranges.push((start, end)), + } + }; + let mut valid_from = 0; + for &(null_start, null_end) in nulls { + keep(offsets[valid_from], offsets[null_start]); + valid_from = null_end; + } + keep(offsets[valid_from], offsets[offsets.len() - 1]); + ranges +} + +/// The offsets with every offset moved down by the elements dropped before it: those before the +/// first list and those under null lists. A null list then starts and ends where the valid lists +/// before it end. +fn shift_past_null_runs(offsets: &[O], nulls: &[(usize, usize)]) -> BufferMut { + let mut shifted = BufferMut::copy_from(offsets); + let mut dropped = offsets[0]; + let mut valid_from = 0; + for &(null_start, null_end) in nulls { + for offset in &mut shifted[valid_from..null_start] { + *offset = *offset - dropped; + } + shifted[null_start..null_end].fill(offsets[null_start] - dropped); + dropped += offsets[null_end] - offsets[null_start]; + valid_from = null_end; + } + for offset in &mut shifted[valid_from..] { + *offset = *offset - dropped; + } + shifted +} diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 08792fc251a..07a24e80e87 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use std::sync::LazyLock; +use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -1058,3 +1059,50 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } + +#[rstest] +#[case::nothing_to_drop(&[1, 2, 3], &[0, 1, 1, 3], &[true, false, true], &[1, 2, 3], &[0, 1, 1, 3])] +#[case::leftovers_at_the_ends(&[9, 1, 2, 9], &[1, 2, 3], &[true, true], &[1, 2], &[0, 1, 2])] +#[case::null_lists_holding_elements( + &[9, 9, 9, 1, 9, 9, 2, 3, 9, 9], + &[1, 3, 4, 6, 8, 9], + &[false, true, false, true, false], + &[1, 2, 3], + &[0, 0, 1, 1, 3, 3] +)] +#[case::only_the_last_null_list_holds_elements( + &[1, 2, 9], + &[0, 1, 1, 2, 3], + &[true, false, true, false], + &[1, 2], + &[0, 1, 1, 2, 2] +)] +#[case::every_list_null(&[9, 9], &[0, 1, 2], &[false, false], &[], &[0, 0, 0])] +fn rebuild_with_referenced_valid_elements( + #[case] elements: &[i32], + #[case] offsets: &[u32], + #[case] validity: &[bool], + #[case] expected_elements: &[i32], + #[case] expected_offsets: &[u32], +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let list = ListArray::try_new( + PrimitiveArray::from_iter(elements.iter().copied()).into_array(), + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + Validity::from_iter(validity.iter().copied()), + &mut ctx, + )?; + let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + assert_arrays_eq!( + rebuilt.elements().clone(), + PrimitiveArray::from_iter(expected_elements.iter().copied()), + &mut ctx + ); + assert_arrays_eq!( + rebuilt.offsets().clone(), + PrimitiveArray::from_iter(expected_offsets.iter().copied()), + &mut ctx + ); + assert_arrays_eq!(rebuilt, list, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 0fe73d4f0ad..53d0651da27 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -18,6 +18,8 @@ use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::listview::list_from_list_view; +use crate::arrays::listview::list_view_from_list; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; @@ -121,6 +123,47 @@ impl ListViewArray { } } + /// Rebuilds the [`ListViewArray`] over only the elements that valid lists reference: nothing + /// before the first list, after the last, between lists or under a null list. + /// + /// Lists that are zero-copy to a [`ListArray`] keep their layout, with their ends sliced off, + /// unless a null list holds elements. Other lists are rebuilt back to back. An array with + /// nothing to drop is returned as it is. + /// + /// [`ListArray`]: crate::arrays::ListArray + pub fn rebuild_with_referenced_valid_elements( + &self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !self.is_zero_copy_to_list() { + // The rebuild lays out only the valid lists' elements, and leaves null lists empty. + return self.rebuild(ListViewRebuildMode::MakeExact, ctx); + } + if self.is_empty() { + return self.rebuild(ListViewRebuildMode::TrimElements, ctx); + } + let nulls = self.validity()?.execute_null_runs(self.len(), ctx)?; + let sizes = self.sizes().clone().execute::(ctx)?; + let null_lists_empty = match_each_integer_ptype!(sizes.ptype(), |S| { + let sizes = sizes.as_slice::(); + nulls + .iter() + .all(|&(start, end)| sizes[start..end].iter().all(|size| *size == S::default())) + }); + if !null_lists_empty { + // The lists are back to back, so they are a `ListArray`'s, without copying. + let list = list_from_list_view(self.clone(), ctx)? + .rebuild_with_referenced_valid_elements(ctx)?; + return list_view_from_list(list, ctx); + } + let (start, end) = self.referenced_element_bounds(ctx)?; + if start == 0 && end == self.elements().len() { + return Ok(self.clone()); + } + // SAFETY: the bounds are the referenced elements' bounds. + unsafe { self.trim_elements(start, end) } + } + /// Rebuilds a [`ListViewArray`], removing all data overlaps and creating a flattened layout. /// /// This is useful when the `elements` child array of the [`ListViewArray`] might have @@ -470,6 +513,72 @@ mod tests { use crate::dtype::Nullability; use crate::validity::Validity; + #[rstest] + #[case::nothing_to_drop(&[1, 2], &[(0, 1), (1, 1)], &[true, true], true, &[1, 2], &[(0, 1), (1, 1)])] + #[case::leftovers_at_the_ends( + &[9, 1, 2, 9], + &[(1, 1), (2, 1)], + &[true, true], + true, + &[1, 2], + &[(0, 1), (1, 1)] + )] + #[case::null_list_holding_elements( + &[9, 1, 9, 2], + &[(1, 1), (2, 1), (3, 1)], + &[true, false, true], + true, + &[1, 2], + &[(0, 1), (1, 0), (1, 1)] + )] + #[case::gap_between_lists(&[1, 9, 2], &[(0, 1), (2, 1)], &[true, true], false, &[1, 2], &[(0, 1), (1, 1)])] + #[case::lists_out_of_order( + &[1, 2, 3, 9], + &[(2, 1), (0, 2)], + &[true, true], + false, + &[3, 1, 2], + &[(0, 1), (1, 2)] + )] + fn rebuild_with_referenced_valid_elements( + #[case] elements: &[i32], + #[case] views: &[(u32, u32)], + #[case] validity: &[bool], + #[case] zero_copy_to_list: bool, + #[case] expected_elements: &[i32], + #[case] expected_views: &[(u32, u32)], + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let view = ListViewArray::try_new( + PrimitiveArray::from_iter(elements.iter().copied()).into_array(), + PrimitiveArray::from_iter(views.iter().map(|(offset, _)| *offset)).into_array(), + PrimitiveArray::from_iter(views.iter().map(|(_, size)| *size)).into_array(), + Validity::from_iter(validity.iter().copied()), + &mut ctx, + )?; + // SAFETY: the cases marked zero-copy to a list have sorted lists with nothing between them. + let view = unsafe { view.with_zero_copy_to_list(zero_copy_to_list) }; + let rebuilt = view.rebuild_with_referenced_valid_elements(&mut ctx)?; + assert!(rebuilt.is_zero_copy_to_list()); + assert_arrays_eq!( + rebuilt.elements().clone(), + PrimitiveArray::from_iter(expected_elements.iter().copied()), + &mut ctx + ); + assert_arrays_eq!( + rebuilt.offsets().clone(), + PrimitiveArray::from_iter(expected_views.iter().map(|(offset, _)| *offset)), + &mut ctx + ); + assert_arrays_eq!( + rebuilt.sizes().clone(), + PrimitiveArray::from_iter(expected_views.iter().map(|(_, size)| *size)), + &mut ctx + ); + assert_arrays_eq!(rebuilt, view, &mut ctx); + Ok(()) + } + #[test] fn test_rebuild_flatten_removes_overlaps() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 449f66b07d9..8041b9ba3ef 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -17,6 +17,7 @@ use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_mask::MaskValues; @@ -272,6 +273,23 @@ impl Validity { } } + /// Execute this validity into the runs of null rows of the given length, as sorted + /// `(start, end)` ranges. + pub(crate) fn execute_null_runs( + &self, + length: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if length == 0 || matches!(self, Self::NonNullable | Self::AllValid) { + return Ok(Vec::new()); + } + Ok(match (!&self.execute_mask(length, ctx)?).slices() { + AllOr::All => vec![(0, length)], + AllOr::None => Vec::new(), + AllOr::Some(runs) => runs.to_vec(), + }) + } + /// Compare the logical masks of two Validity values of the given length, executing them /// into [`Mask`]s if necessary. pub fn mask_eq( From a0ac766d5ae0d936de5794438ceed4815e016642 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 14:46:00 -0400 Subject: [PATCH 2/2] walk runs of valid lists instead of null lists Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/rebuild.rs | 47 +++++++++------------ vortex-array/src/arrays/listview/rebuild.rs | 16 +++++-- vortex-array/src/validity.rs | 18 -------- 3 files changed, 34 insertions(+), 47 deletions(-) diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index 622631b8c18..21845676422 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -26,18 +26,20 @@ impl ListArray { ctx: &mut ExecutionCtx, ) -> VortexResult { let offsets = self.offsets().clone().execute::(ctx)?; - let nulls = self.list_validity().execute_null_runs(self.len(), ctx)?; + let validity = self.list_validity().execute_mask(self.len(), ctx)?; + let all = [(0, self.len())]; + let valid = validity.slices().unwrap_or_else(|| &all, || &[]); match_each_integer_ptype!(offsets.ptype(), |O| { - self.rebuild_without_null_runs(offsets.as_slice::(), &nulls) + self.rebuild_with_valid_runs(offsets.as_slice::(), valid) }) } - fn rebuild_without_null_runs( + fn rebuild_with_valid_runs( &self, offsets: &[O], - nulls: &[(usize, usize)], + valid: &[(usize, usize)], ) -> VortexResult { - let kept = referenced_valid_ranges(offsets, nulls); + let kept = referenced_valid_ranges(offsets, valid); if kept.iter().map(|(start, end)| end - start).sum::() == self.elements().len() { return Ok(self.clone()); } @@ -54,7 +56,7 @@ impl ListArray { .filter(Mask::from_buffer_with_slices(bits.freeze(), kept))? } }; - let offsets = shift_past_null_runs(offsets, nulls); + let offsets = shift_past_null_lists(offsets, valid); // SAFETY: the elements are the ranges of the runs of valid lists, in order, and every offset // moved down by the elements dropped before it, so each valid list keeps its elements and // each null list is empty. @@ -67,43 +69,36 @@ impl ListArray { /// The element ranges of the runs of valid lists, with touching ranges merged. fn referenced_valid_ranges( offsets: &[O], - nulls: &[(usize, usize)], + valid: &[(usize, usize)], ) -> Vec<(usize, usize)> { - let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(nulls.len() + 1); - let mut keep = |start: O, end: O| { - let (start, end): (usize, usize) = (start.as_(), end.as_()); + let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(valid.len()); + for &(first, end) in valid { + let (start, end): (usize, usize) = (offsets[first].as_(), offsets[end].as_()); match ranges.last_mut() { _ if start == end => {} Some(last) if last.1 == start => last.1 = end, _ => ranges.push((start, end)), } - }; - let mut valid_from = 0; - for &(null_start, null_end) in nulls { - keep(offsets[valid_from], offsets[null_start]); - valid_from = null_end; } - keep(offsets[valid_from], offsets[offsets.len() - 1]); ranges } /// The offsets with every offset moved down by the elements dropped before it: those before the /// first list and those under null lists. A null list then starts and ends where the valid lists /// before it end. -fn shift_past_null_runs(offsets: &[O], nulls: &[(usize, usize)]) -> BufferMut { +fn shift_past_null_lists(offsets: &[O], valid: &[(usize, usize)]) -> BufferMut { let mut shifted = BufferMut::copy_from(offsets); let mut dropped = offsets[0]; - let mut valid_from = 0; - for &(null_start, null_end) in nulls { - for offset in &mut shifted[valid_from..null_start] { + let mut null_from = 0; + for &(valid_start, valid_end) in valid { + shifted[null_from..valid_start].fill(offsets[null_from] - dropped); + dropped += offsets[valid_start] - offsets[null_from]; + for offset in &mut shifted[valid_start..valid_end] { *offset = *offset - dropped; } - shifted[null_start..null_end].fill(offsets[null_start] - dropped); - dropped += offsets[null_end] - offsets[null_start]; - valid_from = null_end; - } - for offset in &mut shifted[valid_from..] { - *offset = *offset - dropped; + null_from = valid_end; } + // The null lists after the last valid one, and the end of the last list. + shifted[null_from..].fill(offsets[null_from] - dropped); shifted } diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 53d0651da27..9d97d36beb7 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -142,13 +142,23 @@ impl ListViewArray { if self.is_empty() { return self.rebuild(ListViewRebuildMode::TrimElements, ctx); } - let nulls = self.validity()?.execute_null_runs(self.len(), ctx)?; + let validity = self.validity()?.execute_mask(self.len(), ctx)?; + let all = [(0, self.len())]; + let valid = validity.slices().unwrap_or_else(|| &all, || &[]); let sizes = self.sizes().clone().execute::(ctx)?; let null_lists_empty = match_each_integer_ptype!(sizes.ptype(), |S| { let sizes = sizes.as_slice::(); - nulls + let mut null_from = 0; + valid .iter() - .all(|&(start, end)| sizes[start..end].iter().all(|size| *size == S::default())) + .chain([&(self.len(), self.len())]) + .all(|&(valid_start, valid_end)| { + let empty = sizes[null_from..valid_start] + .iter() + .all(|size| *size == S::default()); + null_from = valid_end; + empty + }) }); if !null_lists_empty { // The lists are back to back, so they are a `ListArray`'s, without copying. diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 8041b9ba3ef..449f66b07d9 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -17,7 +17,6 @@ use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_mask::MaskValues; @@ -273,23 +272,6 @@ impl Validity { } } - /// Execute this validity into the runs of null rows of the given length, as sorted - /// `(start, end)` ranges. - pub(crate) fn execute_null_runs( - &self, - length: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - if length == 0 || matches!(self, Self::NonNullable | Self::AllValid) { - return Ok(Vec::new()); - } - Ok(match (!&self.execute_mask(length, ctx)?).slices() { - AllOr::All => vec![(0, length)], - AllOr::None => Vec::new(), - AllOr::Some(runs) => runs.to_vec(), - }) - } - /// Compare the logical masks of two Validity values of the given length, executing them /// into [`Mask`]s if necessary. pub fn mask_eq(