From 06eaba283936a6b675b00a1ee61799878a4a67b4 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 13:06:12 -0400 Subject: [PATCH 01/10] 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 437c0bbf09797ecd0c2c89da1df9a07ee0fb5932 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 23 Oct 2026 14:46:00 -0400 Subject: [PATCH 02/10] 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( From a6a24d698216656fb57ab27167ca17fac35b5622 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 10:46:20 +0000 Subject: [PATCH 03/10] check whether a null list holds elements before collecting runs Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/mod.rs | 1 + vortex-array/src/arrays/list/rebuild.rs | 136 +++++++++++--------- vortex-array/src/arrays/list/tests.rs | 8 +- vortex-array/src/arrays/listview/rebuild.rs | 17 +-- 4 files changed, 86 insertions(+), 76 deletions(-) diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 7c604852b59..370ad4af599 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -13,6 +13,7 @@ pub use vtable::ListArray; pub(crate) mod compute; mod rebuild; +pub(crate) use rebuild::null_lists_hold_elements; mod vtable; pub use vtable::List; diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index 21845676422..4bc2fd99df7 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ExecutionCtx; @@ -19,86 +21,98 @@ 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. + /// A list with no null lists and nothing before the first list or after the last is returned + /// as it is, after reading two offsets. pub fn rebuild_with_referenced_valid_elements( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { + let validity = self.list_validity(); + if validity.definitely_no_nulls() + && self.offset_at(0, ctx)? == 0 + && self.offset_at(self.len(), ctx)? == self.elements().len() + { + return Ok(self.clone()); + } + let offsets = self.offsets().clone().execute::(ctx)?; - let validity = self.list_validity().execute_mask(self.len(), ctx)?; - let all = [(0, self.len())]; - let valid = validity.slices().unwrap_or_else(|| &all, || &[]); + let valid = validity.execute_mask(self.len(), ctx)?; match_each_integer_ptype!(offsets.ptype(), |O| { - self.rebuild_with_valid_runs(offsets.as_slice::(), valid) + self.rebuild_with_valid(offsets.as_slice::(), &valid) }) } - fn rebuild_with_valid_runs( + fn rebuild_with_valid( &self, offsets: &[O], - valid: &[(usize, usize)], + valid: &Mask, ) -> VortexResult { - let kept = referenced_valid_ranges(offsets, valid); - if kept.iter().map(|(start, end)| end - start).sum::() == self.elements().len() { + let all = [(0, self.len())]; + let (first, last): (usize, usize) = (offsets[0].as_(), offsets[self.len()].as_()); + // When no null list holds elements, all the lists are kept as one run: the null lists in it + // stay empty. + let runs = if null_lists_hold_elements(valid, |index| offsets[index] != offsets[index + 1]) + { + valid.slices().unwrap_or_else(|| &all, || &[]) + } else if first == 0 && last == 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))? - } + } else { + &all }; - 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. - 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], - valid: &[(usize, usize)], -) -> Vec<(usize, usize)> { - 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)), + // Keep the elements of each run, merging runs whose elements touch, and move each offset + // down by the elements dropped before it, so that null lists are empty. + let mut kept: Vec<(usize, usize)> = Vec::with_capacity(runs.len()); + let mut ends = BufferMut::copy_from(offsets); + let mut dropped = offsets[0]; + let mut null_from = 0; + for &(run_start, run_end) in runs { + let (start, stop): (usize, usize) = (offsets[run_start].as_(), offsets[run_end].as_()); + match kept.last_mut() { + Some(range) if range.1 == start => range.1 = stop, + _ if start < stop => kept.push((start, stop)), + _ => {} + } + ends[null_from..run_start].fill(offsets[null_from] - dropped); + dropped += offsets[run_start] - offsets[null_from]; + for end in &mut ends[run_start..run_end] { + *end = *end - dropped; + } + null_from = run_end; + } + ends[null_from..].fill(offsets[null_from] - dropped); + + let mut bits = BitBufferMut::new_unset(self.elements().len()); + for &(start, stop) in &kept { + bits.fill_range(start, stop, true); } + // With the ranges cached on the mask, the filter slices one range and copies several range + // by range. + let elements = self + .elements() + .filter(Mask::from_buffer_with_slices(bits.freeze(), kept))?; + // SAFETY: the elements are the runs' elements 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, ends.freeze().into_array(), self.list_validity()) + }) } - 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_lists(offsets: &[O], valid: &[(usize, usize)]) -> BufferMut { - let mut shifted = BufferMut::copy_from(offsets); - let mut dropped = offsets[0]; - 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; +/// Whether a null list holds elements, given which lists hold elements. +pub(crate) fn null_lists_hold_elements( + valid: &Mask, + mut holds_elements: impl FnMut(usize) -> bool, +) -> bool { + match valid.bit_buffer() { + AllOr::All => false, + AllOr::None => (0..valid.len()).any(holds_elements), + AllOr::Some(valid) => { + BitBuffer::collect_bool(valid.len(), &mut holds_elements) + .bitand_not(valid) + .true_count() + > 0 } - 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/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 07a24e80e87..fdcff3160df 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -1062,7 +1062,9 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { #[rstest] #[case::nothing_to_drop(&[1, 2, 3], &[0, 1, 1, 3], &[true, false, true], &[1, 2, 3], &[0, 1, 1, 3])] +#[case::exact(&[1, 2, 3], &[0, 1, 3], &[true, true], &[1, 2, 3], &[0, 1, 3])] #[case::leftovers_at_the_ends(&[9, 1, 2, 9], &[1, 2, 3], &[true, true], &[1, 2], &[0, 1, 2])] +#[case::leftover_at_the_end(&[1, 2, 9], &[0, 1, 2], &[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], @@ -1089,7 +1091,11 @@ fn rebuild_with_referenced_valid_elements( 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()), + if validity.iter().all(|valid| *valid) { + Validity::NonNullable + } else { + Validity::from_iter(validity.iter().copied()) + }, &mut ctx, )?; let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 9d97d36beb7..ae688faca6f 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -16,6 +16,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; +use crate::arrays::list::null_lists_hold_elements; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::list_from_list_view; @@ -142,23 +143,11 @@ impl ListViewArray { if self.is_empty() { return self.rebuild(ListViewRebuildMode::TrimElements, ctx); } - let validity = self.validity()?.execute_mask(self.len(), ctx)?; - let all = [(0, self.len())]; - let valid = validity.slices().unwrap_or_else(|| &all, || &[]); + let valid = self.validity()?.execute_mask(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::(); - let mut null_from = 0; - valid - .iter() - .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 - }) + !null_lists_hold_elements(&valid, |index| sizes[index] != S::default()) }); if !null_lists_empty { // The lists are back to back, so they are a `ListArray`'s, without copying. From 14aa30ae32c02c3d9d3e03729217fba928ec9d4e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 11:05:00 +0000 Subject: [PATCH 04/10] split the list rebuild into its mask and its offsets Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/rebuild.rs | 80 ++++++++++++++----------- vortex-array/src/arrays/list/tests.rs | 7 +++ 2 files changed, 53 insertions(+), 34 deletions(-) diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index 4bc2fd99df7..fb5dcbf121e 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::AllOr; @@ -60,44 +61,55 @@ impl ListArray { &all }; - // Keep the elements of each run, merging runs whose elements touch, and move each offset - // down by the elements dropped before it, so that null lists are empty. - let mut kept: Vec<(usize, usize)> = Vec::with_capacity(runs.len()); - let mut ends = BufferMut::copy_from(offsets); - let mut dropped = offsets[0]; - let mut null_from = 0; - for &(run_start, run_end) in runs { - let (start, stop): (usize, usize) = (offsets[run_start].as_(), offsets[run_end].as_()); - match kept.last_mut() { - Some(range) if range.1 == start => range.1 = stop, - _ if start < stop => kept.push((start, stop)), - _ => {} - } - ends[null_from..run_start].fill(offsets[null_from] - dropped); - dropped += offsets[run_start] - offsets[null_from]; - for end in &mut ends[run_start..run_end] { - *end = *end - dropped; - } - null_from = run_end; + let elements = + self.elements() + .filter(elements_of_runs(offsets, runs, self.elements().len()))?; + let offsets = offsets_over_runs(offsets, runs); + // SAFETY: the elements are the runs' elements in order, and the offsets moved down by the + // elements dropped before them, so each list in a run keeps its elements and each list + // between runs is empty. + Ok(unsafe { + ListArray::new_unchecked(elements, offsets.into_array(), self.list_validity()) + }) + } +} + +/// The elements of the runs of lists, as a mask over `len` elements. +fn elements_of_runs(offsets: &[O], runs: &[(usize, usize)], len: usize) -> Mask { + let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(runs.len()); + for &(run_start, run_end) in runs { + let (start, stop): (usize, usize) = (offsets[run_start].as_(), offsets[run_end].as_()); + match ranges.last_mut() { + Some(range) if range.1 == start => range.1 = stop, + _ if start < stop => ranges.push((start, stop)), + _ => {} } - ends[null_from..].fill(offsets[null_from] - dropped); + } + let mut bits = BitBufferMut::new_unset(len); + for &(start, stop) in &ranges { + bits.fill_range(start, stop, true); + } + // With the ranges cached on the mask, a filter slices one range and copies several range by + // range. + Mask::from_buffer_with_slices(bits.freeze(), ranges) +} - let mut bits = BitBufferMut::new_unset(self.elements().len()); - for &(start, stop) in &kept { - bits.fill_range(start, stop, true); +/// The offsets once only the runs' elements remain: each offset moves down by the elements dropped +/// before it, so that the lists between runs are empty. +fn offsets_over_runs(offsets: &[O], runs: &[(usize, usize)]) -> Buffer { + let mut shifted = BufferMut::copy_from(offsets); + let mut dropped = offsets[0]; + let mut between_from = 0; + for &(run_start, run_end) in runs { + shifted[between_from..run_start].fill(offsets[between_from] - dropped); + dropped += offsets[run_start] - offsets[between_from]; + for offset in &mut shifted[run_start..run_end] { + *offset = *offset - dropped; } - // With the ranges cached on the mask, the filter slices one range and copies several range - // by range. - let elements = self - .elements() - .filter(Mask::from_buffer_with_slices(bits.freeze(), kept))?; - // SAFETY: the elements are the runs' elements 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, ends.freeze().into_array(), self.list_validity()) - }) + between_from = run_end; } + shifted[between_from..].fill(offsets[between_from] - dropped); + shifted.freeze() } /// Whether a null list holds elements, given which lists hold elements. diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index fdcff3160df..9af246e002a 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -1079,6 +1079,13 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { &[1, 2], &[0, 1, 1, 2, 2] )] +#[case::an_empty_valid_list_between_null_lists( + &[9, 9, 1], + &[0, 1, 1, 2, 3], + &[false, true, false, true], + &[1], + &[0, 0, 0, 0, 1] +)] #[case::every_list_null(&[9, 9], &[0, 1, 2], &[false, false], &[], &[0, 0, 0])] fn rebuild_with_referenced_valid_elements( #[case] elements: &[i32], From 03809727a40a15c65f0f1f02efed843a37c094c0 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 13:59:57 +0000 Subject: [PATCH 05/10] walk null runs, slice the ends, rebuild zero-copy list views in place Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/mod.rs | 3 +- vortex-array/src/arrays/list/rebuild.rs | 175 +++++++----- vortex-array/src/arrays/list/tests.rs | 74 +++++ vortex-array/src/arrays/listview/rebuild.rs | 289 ++++++++++++++++++-- 4 files changed, 441 insertions(+), 100 deletions(-) diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 370ad4af599..bbbba9563ca 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -13,7 +13,8 @@ pub use vtable::ListArray; pub(crate) mod compute; mod rebuild; -pub(crate) use rebuild::null_lists_hold_elements; +pub(crate) use rebuild::drop_unreferenced_elements; +pub(crate) use rebuild::null_runs_holding_elements; mod vtable; pub use vtable::List; diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index fb5dcbf121e..0231620944f 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -1,29 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBuffer; -use vortex_buffer::BitBufferMut; -use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; +use crate::ArrayRef; 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::arrays::primitive::PrimitiveArrayExt; use crate::dtype::IntegerPType; -use crate::match_each_integer_ptype; +use crate::dtype::PType; +use crate::match_each_unsigned_integer_ptype; +use crate::validity::Validity; 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. /// - /// A list with no null lists and nothing before the first list or after the last is returned - /// as it is, after reading two offsets. + /// Null lists come out empty, the offsets start at 0, end at the length of the elements and + /// keep their type, and the validity is kept. An array with nothing to drop is returned as it + /// is. pub fn rebuild_with_referenced_valid_elements( &self, ctx: &mut ExecutionCtx, @@ -37,94 +41,119 @@ impl ListArray { } let offsets = self.offsets().clone().execute::(ctx)?; + let offsets_ptype = offsets.ptype(); + // Offsets are never negative, so they are rebuilt unsigned and reinterpreted back. + let offsets = offsets.reinterpret_cast(offsets_ptype.to_unsigned()); let valid = validity.execute_mask(self.len(), ctx)?; - match_each_integer_ptype!(offsets.ptype(), |O| { - self.rebuild_with_valid(offsets.as_slice::(), &valid) + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + self.rebuild_with_valid(offsets.as_slice::(), offsets_ptype, &valid, ctx) }) } fn rebuild_with_valid( &self, offsets: &[O], + offsets_ptype: PType, valid: &Mask, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let all = [(0, self.len())]; - let (first, last): (usize, usize) = (offsets[0].as_(), offsets[self.len()].as_()); - // When no null list holds elements, all the lists are kept as one run: the null lists in it - // stay empty. - let runs = if null_lists_hold_elements(valid, |index| offsets[index] != offsets[index + 1]) - { - valid.slices().unwrap_or_else(|| &all, || &[]) - } else if first == 0 && last == self.elements().len() { + let (starts, end) = (&offsets[..self.len()], offsets[self.len()].as_()); + let null_runs = null_runs_holding_elements(starts, end, valid); + if null_runs.is_empty() && offsets[0].as_() == 0 && end == self.elements().len() { return Ok(self.clone()); - } else { - &all - }; + } - let elements = - self.elements() - .filter(elements_of_runs(offsets, runs, self.elements().len()))?; - let offsets = offsets_over_runs(offsets, runs); - // SAFETY: the elements are the runs' elements in order, and the offsets moved down by the - // elements dropped before them, so each list in a run keeps its elements and each list - // between runs is empty. - Ok(unsafe { - ListArray::new_unchecked(elements, offsets.into_array(), self.list_validity()) - }) + let (elements, mut offsets) = + drop_unreferenced_elements(self.elements(), starts, end, &null_runs, ctx.allocator())?; + offsets.push( + num_traits::cast(elements.len()) + .vortex_expect("the elements kept fit the offsets' type"), + ); + let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(offsets_ptype) + .into_array(); + // SAFETY: the offsets keep their type, start at 0, never decrease and end at the length of + // the elements, which are the valid lists' elements in order. + Ok(unsafe { ListArray::new_unchecked(elements, offsets, self.list_validity()) }) } } -/// The elements of the runs of lists, as a mask over `len` elements. -fn elements_of_runs(offsets: &[O], runs: &[(usize, usize)], len: usize) -> Mask { - let mut ranges: Vec<(usize, usize)> = Vec::with_capacity(runs.len()); - for &(run_start, run_end) in runs { - let (start, stop): (usize, usize) = (offsets[run_start].as_(), offsets[run_end].as_()); - match ranges.last_mut() { - Some(range) if range.1 == start => range.1 = stop, - _ if start < stop => ranges.push((start, stop)), - _ => {} +/// The runs of null lists that hold elements, among lists laid out back to back: list `i` starts +/// at `starts[i]` and ends where the next list starts, or at `end` for the last list. +pub(crate) fn null_runs_holding_elements( + starts: &[O], + end: usize, + valid: &Mask, +) -> Vec<(usize, usize)> { + let start_of = |index: usize| starts.get(index).map_or(end, |start| start.as_()); + let mut runs = Vec::new(); + // Lists never start before the previous one, so a run of null lists holds elements if and + // only if the list after it starts later than the run does. + let mut push_if_holding = |run_start: usize, run_end: usize| { + if run_start < run_end && start_of(run_start) != start_of(run_end) { + runs.push((run_start, run_end)); + } + }; + match valid.bit_buffer() { + AllOr::All => {} + AllOr::None => push_if_holding(0, valid.len()), + AllOr::Some(bits) => { + let mut run_start = 0; + // The empty slice at the end closes a run of null lists at the end. + for (valid_start, valid_end) in bits.set_slices().chain([(bits.len(), bits.len())]) { + push_if_holding(run_start, valid_start); + run_start = valid_end; + } } } - let mut bits = BitBufferMut::new_unset(len); - for &(start, stop) in &ranges { - bits.fill_range(start, stop, true); - } - // With the ranges cached on the mask, a filter slices one range and copies several range by - // range. - Mask::from_buffer_with_slices(bits.freeze(), ranges) + runs } -/// The offsets once only the runs' elements remain: each offset moves down by the elements dropped -/// before it, so that the lists between runs are empty. -fn offsets_over_runs(offsets: &[O], runs: &[(usize, usize)]) -> Buffer { - let mut shifted = BufferMut::copy_from(offsets); - let mut dropped = offsets[0]; - let mut between_from = 0; - for &(run_start, run_end) in runs { - shifted[between_from..run_start].fill(offsets[between_from] - dropped); - dropped += offsets[run_start] - offsets[between_from]; - for offset in &mut shifted[run_start..run_end] { - *offset = *offset - dropped; +/// Drops the elements before the first of the lists laid out back to back, after the last, and +/// under the runs of null lists from [`null_runs_holding_elements`], which come out empty. +/// +/// Returns the elements kept and where each list now starts, with room to push one more offset. +pub(crate) fn drop_unreferenced_elements( + elements: &ArrayRef, + starts: &[O], + end: usize, + null_runs: &[(usize, usize)], + allocator: &BufferAllocatorRef, +) -> VortexResult<(ArrayRef, BufferMut)> { + let start_of = |index: usize| starts.get(index).map_or(end, |start| start.as_()); + let first = start_of(0); + // The ranges of elements kept, counted from the first list's start. + let mut kept = Vec::with_capacity(null_runs.len() + 1); + let mut kept_start = 0; + for &(run_start, run_end) in null_runs { + if kept_start < start_of(run_start) - first { + kept.push((kept_start, start_of(run_start) - first)); } - between_from = run_end; + kept_start = start_of(run_end) - first; } - shifted[between_from..].fill(offsets[between_from] - dropped); - shifted.freeze() -} + if kept_start < end - first { + kept.push((kept_start, end - first)); + } + let referenced = elements.slice(first..end)?; + let elements = if let [(start, stop)] = kept[..] { + referenced.slice(start..stop)? + } else { + // With the ranges cached on the mask, the filter copies the elements range by range. + referenced.filter(Mask::from_slices(end - first, kept))? + }; -/// Whether a null list holds elements, given which lists hold elements. -pub(crate) fn null_lists_hold_elements( - valid: &Mask, - mut holds_elements: impl FnMut(usize) -> bool, -) -> bool { - match valid.bit_buffer() { - AllOr::All => false, - AllOr::None => (0..valid.len()).any(holds_elements), - AllOr::Some(valid) => { - BitBuffer::collect_bool(valid.len(), &mut holds_elements) - .bitand_not(valid) - .true_count() - > 0 + let mut offsets = BufferMut::with_capacity_in(starts.len() + 1, allocator.clone()); + let mut dropped = starts.first().copied().unwrap_or_else(O::zero); + let mut from = 0; + for &(run_start, run_end) in null_runs { + offsets.extend(starts[from..run_start].iter().map(|&start| start - dropped)); + offsets.push_n(starts[run_start] - dropped, run_end - run_start); + // A run at the end leaves no list after it to move down. + if let Some(&next) = starts.get(run_end) { + dropped += next - starts[run_start]; } + from = run_end; } + offsets.extend(starts[from..].iter().map(|&start| start - dropped)); + Ok((elements, offsets)) } diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 9af246e002a..7354bb0b41a 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; @@ -13,6 +14,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use super::*; +use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::FilterArray; @@ -1087,6 +1089,21 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { &[0, 0, 0, 0, 1] )] #[case::every_list_null(&[9, 9], &[0, 1, 2], &[false, false], &[], &[0, 0, 0])] +#[case::empty_null_list_with_leftovers_at_the_ends( + &[9, 1, 2, 9], + &[1, 2, 2, 3], + &[true, false, true], + &[1, 2], + &[0, 1, 1, 2] +)] +#[case::every_list_null_and_empty_with_leftovers( + &[9, 9], + &[1, 1, 1], + &[false, false], + &[], + &[0, 0, 0] +)] +#[case::no_lists_with_a_leftover(&[9], &[1], &[], &[], &[0])] fn rebuild_with_referenced_valid_elements( #[case] elements: &[i32], #[case] offsets: &[u32], @@ -1119,3 +1136,60 @@ fn rebuild_with_referenced_valid_elements( assert_arrays_eq!(rebuilt, list, &mut ctx); Ok(()) } + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +#[case::empty_null_list(Validity::from_iter([true, false, true]))] +fn rebuild_with_referenced_valid_elements_returns_an_exact_list_as_is( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let list = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0u32, 1, 1, 3].into_array(), + validity, + &mut ctx, + )?; + let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + assert!(ArrayRef::ptr_eq(rebuilt.elements(), list.elements())); + assert!(ArrayRef::ptr_eq(rebuilt.offsets(), list.offsets())); + Ok(()) +} + +/// A slice keeps all of its parent's elements, and its validity can start mid-byte. +#[rstest] +#[case::no_null_list(2..4, &[3, 4, 5, 6], &[0, 2, 4])] +#[case::a_null_list_holding_elements_first(1..4, &[3, 4, 5, 6], &[0, 0, 2, 4])] +#[case::a_null_list_holding_elements_last(3..5, &[5, 6], &[0, 2, 2])] +fn rebuild_with_referenced_valid_elements_sliced( + #[case] range: Range, + #[case] expected_elements: &[i32], + #[case] expected_offsets: &[u32], +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // [[0, 1], null over [2], [3, 4], [5, 6], null over [7, 8, 9]] + let list = ListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + buffer![0u32, 2, 3, 5, 7, 10].into_array(), + Validity::from_iter([true, false, true, true, false]), + &mut ctx, + )? + .into_array() + .slice(range)? + .as_::() + .into_owned(); + 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 ae688faca6f..78ff59f5ad2 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::FromPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -16,11 +17,10 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; -use crate::arrays::list::null_lists_hold_elements; +use crate::arrays::list::drop_unreferenced_elements; +use crate::arrays::list::null_runs_holding_elements; 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; @@ -87,6 +87,9 @@ pub enum ListViewRebuildMode { /// Equivalent to `MakeZeroCopyToList` plus `TrimElements`. /// + /// An array that is not yet zero-copy to a [`ListArray`] is copied with only its valid lists' + /// elements, and its null lists come out empty. + /// /// This is useful when concatenating multiple [`ListViewArray`]s together to create a new /// [`ListViewArray`] that is also zero-copy to a [`ListArray`]. /// @@ -127,40 +130,78 @@ 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. + /// The result is zero-copy to a [`ListArray`] and its null lists are empty. A non-empty array + /// that is already zero-copy to a [`ListArray`] keeps its offsets' and sizes' types, and is + /// returned as it is when there is nothing to drop. Any other array is rebuilt with + /// [`MakeExact`](ListViewRebuildMode::MakeExact). /// /// [`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. + if self.is_empty() || !self.is_zero_copy_to_list() { return self.rebuild(ListViewRebuildMode::MakeExact, ctx); } - if self.is_empty() { - return self.rebuild(ListViewRebuildMode::TrimElements, ctx); - } - let valid = self.validity()?.execute_mask(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::(); - !null_lists_hold_elements(&valid, |index| sizes[index] != 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 validity = self.validity()?; + // The lists are back to back, so the last one ends where the referenced elements do. + let last = self.len() - 1; + let end = self.offset_at(last, ctx) + self.size_at(last, ctx); + if validity.definitely_no_nulls() + && self.offset_at(0, ctx) == 0 + && end == self.elements().len() + { + return Ok(self.clone()); } - let (start, end) = self.referenced_element_bounds(ctx)?; - if start == 0 && end == self.elements().len() { + + let offsets = self.offsets().clone().execute::(ctx)?; + let offsets_ptype = offsets.ptype(); + // Offsets are never negative, so they are rebuilt unsigned and reinterpreted back. + let offsets = offsets.reinterpret_cast(offsets_ptype.to_unsigned()); + let valid = validity.execute_mask(self.len(), ctx)?; + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + self.rebuild_back_to_back(offsets.as_slice::(), offsets_ptype, end, &valid, ctx) + }) + } + + /// [`rebuild_with_referenced_valid_elements`](Self::rebuild_with_referenced_valid_elements) + /// for an array that is zero-copy to a list, whose lists start at `starts` and end at `end`. + fn rebuild_back_to_back( + &self, + starts: &[O], + offsets_ptype: PType, + end: usize, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let null_runs = null_runs_holding_elements(starts, end, valid); + if null_runs.is_empty() && starts[0].as_() == 0 && end == self.elements().len() { return Ok(self.clone()); } - // SAFETY: the bounds are the referenced elements' bounds. - unsafe { self.trim_elements(start, end) } + + let (elements, offsets) = + drop_unreferenced_elements(self.elements(), starts, end, &null_runs, ctx.allocator())?; + let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(offsets_ptype) + .into_array(); + let sizes = if null_runs.is_empty() { + self.sizes().clone() + } else { + let sizes = self.sizes().clone().execute::(ctx)?; + let sizes_ptype = sizes.ptype(); + let sizes = sizes.reinterpret_cast(sizes_ptype.to_unsigned()); + let sizes = match_each_unsigned_integer_ptype!(sizes.ptype(), |S| { + empty_runs(sizes.as_slice::(), &null_runs, ctx.allocator()) + }); + sizes.reinterpret_cast(sizes_ptype).into_array() + }; + + // SAFETY: the lists stay back to back and in order over the elements kept: each valid list + // keeps its elements, and each null list that held elements is now empty. + Ok(unsafe { + ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) + .with_zero_copy_to_list(true) + }) } /// Rebuilds a [`ListViewArray`], removing all data overlaps and creating a flattened layout. @@ -431,6 +472,23 @@ impl ListViewArray { } } +/// The sizes with the lists in `runs` emptied. +fn empty_runs( + sizes: &[S], + runs: &[(usize, usize)], + allocator: &BufferAllocatorRef, +) -> PrimitiveArray { + let mut emptied = BufferMut::with_capacity_in(sizes.len(), allocator.clone()); + let mut from = 0; + for &(run_start, run_end) in runs { + emptied.extend_from_slice(&sizes[from..run_start]); + emptied.push_n(S::zero(), run_end - run_start); + from = run_end; + } + emptied.extend_from_slice(&sizes[from..]); + PrimitiveArray::new(emptied.freeze(), Validity::NonNullable) +} + struct RebuildRanges { new_offsets: BufferMut, new_sizes: BufferMut, @@ -496,20 +554,31 @@ where #[cfg(test)] mod tests { + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use super::super::tests::common::SESSION; use super::ListViewRebuildMode; + use crate::ArrayRef; + use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; + use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::validity::Validity; #[rstest] @@ -539,6 +608,46 @@ mod tests { &[3, 1, 2], &[(0, 1), (1, 2)] )] + #[case::null_list_holding_elements_not_zero_copy( + &[9, 1, 9, 2], + &[(1, 1), (2, 1), (3, 1)], + &[true, false, true], + false, + &[1, 2], + &[(0, 1), (1, 0), (1, 1)] + )] + #[case::empty_null_list_with_leftovers_at_the_ends( + &[9, 1, 2, 9], + &[(1, 1), (2, 0), (2, 1)], + &[true, false, true], + true, + &[1, 2], + &[(0, 1), (1, 0), (1, 1)] + )] + #[case::every_list_null( + &[9, 9], + &[(0, 1), (1, 1)], + &[false, false], + true, + &[], + &[(0, 0), (0, 0)] + )] + #[case::empty_lists_with_leftovers( + &[9, 9], + &[(1, 0), (1, 0)], + &[true, true], + true, + &[], + &[(0, 0), (0, 0)] + )] + #[case::overlapping_lists( + &[1, 2, 3, 9], + &[(0, 3), (1, 2)], + &[true, true], + false, + &[1, 2, 3, 2, 3], + &[(0, 3), (3, 2)] + )] fn rebuild_with_referenced_valid_elements( #[case] elements: &[i32], #[case] views: &[(u32, u32)], @@ -578,6 +687,134 @@ mod tests { Ok(()) } + #[rstest] + #[case::non_nullable(Validity::NonNullable)] + #[case::empty_null_list(Validity::from_iter([true, false, true]))] + fn rebuild_with_referenced_valid_elements_returns_an_exact_list_view_as_is( + #[case] validity: Validity, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let view = ListViewArray::try_new( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + PrimitiveArray::from_iter([0u32, 1, 1]).into_array(), + PrimitiveArray::from_iter([1u32, 0, 2]).into_array(), + validity, + &mut ctx, + )?; + // SAFETY: the lists are sorted with nothing between them. + let view = unsafe { view.with_zero_copy_to_list(true) }; + let rebuilt = view.rebuild_with_referenced_valid_elements(&mut ctx)?; + assert!(ArrayRef::ptr_eq(rebuilt.elements(), view.elements())); + assert!(ArrayRef::ptr_eq(rebuilt.offsets(), view.offsets())); + assert!(ArrayRef::ptr_eq(rebuilt.sizes(), view.sizes())); + Ok(()) + } + + /// `values` as a primitive array of `ptype`. + fn primitive_of( + values: impl IntoIterator, + ptype: PType, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(PrimitiveArray::from_iter(values) + .into_array() + .cast(DType::Primitive(ptype, Nullability::NonNullable))? + .execute::(ctx)? + .into_array()) + } + + /// Random lists with elements before, after and under null lists, rebuilt as a `ListArray` + /// and as list views with and without the zero-copy flag, against a reference that gathers + /// each valid list's elements back to back. + #[test] + fn rebuild_with_referenced_valid_elements_matches_a_reference() -> VortexResult<()> { + const PTYPES: [PType; 5] = [PType::U8, PType::U16, PType::U64, PType::I32, PType::I64]; + let u32_dtype = DType::Primitive(PType::U32, Nullability::NonNullable); + let mut ctx = SESSION.create_execution_ctx(); + let mut rng = StdRng::seed_from_u64(0); + for _ in 0..500 { + let len = rng.random_range(0..10usize); + let mut offset = rng.random_range(0..3u32); + let mut offsets = vec![offset]; + for _ in 0..len { + offset += rng.random_range(0..3u32); + offsets.push(offset); + } + let elements_len = offset + rng.random_range(0..3u32); + let valid: Vec = (0..len).map(|_| rng.random_bool(0.6)).collect(); + let validity = match rng.random_range(0..3u8) { + 0 if valid.iter().all(|valid| *valid) => Validity::NonNullable, + // Validity that starts mid-byte. + 1 => Validity::Array( + BoolArray::from_iter([false].into_iter().chain(valid.iter().copied())) + .into_array() + .slice(1..len + 1)?, + ), + _ => Validity::from_iter(valid.iter().copied()), + }; + let ptype = PTYPES[rng.random_range(0..PTYPES.len())]; + + let mut expected_elements = Vec::new(); + let mut expected_offsets = vec![0u32]; + let mut kept = 0u32; + for (bounds, is_valid) in offsets.windows(2).zip(&valid) { + if *is_valid { + expected_elements.extend(bounds[0]..bounds[1]); + kept += bounds[1] - bounds[0]; + } + expected_offsets.push(kept); + } + let expected_elements = PrimitiveArray::from_iter(expected_elements); + let expected_sizes = + PrimitiveArray::from_iter(expected_offsets.windows(2).map(|w| w[1] - w[0])); + + let elements = PrimitiveArray::from_iter(0..elements_len).into_array(); + let list = ListArray::try_new( + elements.clone(), + primitive_of(offsets.iter().copied(), ptype, &mut ctx)?, + validity.clone(), + &mut ctx, + )?; + let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + assert_arrays_eq!(rebuilt.elements().clone(), expected_elements.clone(), &mut ctx); + assert_eq!(rebuilt.offsets().dtype(), list.offsets().dtype()); + assert_arrays_eq!( + rebuilt.offsets().cast(u32_dtype.clone())?, + PrimitiveArray::from_iter(expected_offsets.iter().copied()), + &mut ctx + ); + assert_arrays_eq!(rebuilt, list, &mut ctx); + + for zero_copy_to_list in [true, false] { + let view = ListViewArray::try_new( + elements.clone(), + primitive_of(offsets[..len].iter().copied(), ptype, &mut ctx)?, + primitive_of(offsets.windows(2).map(|w| w[1] - w[0]), ptype, &mut ctx)?, + validity.clone(), + &mut ctx, + )?; + // SAFETY: the views are a `ListArray`'s lists, sorted 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(), expected_elements.clone(), &mut ctx); + // The copy widens the offsets, so compare their values. + assert_arrays_eq!( + rebuilt.offsets().cast(u32_dtype.clone())?, + PrimitiveArray::from_iter(expected_offsets[..len].iter().copied()), + &mut ctx + ); + assert_arrays_eq!( + rebuilt.sizes().cast(u32_dtype.clone())?, + expected_sizes.clone(), + &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(); From e6b63728a3fd5c4a59b1e212cb4a3ac61b570bc6 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 14:06:47 +0000 Subject: [PATCH 06/10] fmt Signed-off-by: Matt Katz --- vortex-array/src/arrays/listview/rebuild.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 78ff59f5ad2..198dba86e36 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -776,7 +776,11 @@ mod tests { &mut ctx, )?; let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; - assert_arrays_eq!(rebuilt.elements().clone(), expected_elements.clone(), &mut ctx); + assert_arrays_eq!( + rebuilt.elements().clone(), + expected_elements.clone(), + &mut ctx + ); assert_eq!(rebuilt.offsets().dtype(), list.offsets().dtype()); assert_arrays_eq!( rebuilt.offsets().cast(u32_dtype.clone())?, @@ -797,7 +801,11 @@ mod tests { 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(), expected_elements.clone(), &mut ctx); + assert_arrays_eq!( + rebuilt.elements().clone(), + expected_elements.clone(), + &mut ctx + ); // The copy widens the offsets, so compare their values. assert_arrays_eq!( rebuilt.offsets().cast(u32_dtype.clone())?, From b9449095a10ac255ca7d2a65d1368584d280fdeb Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 14:30:25 +0000 Subject: [PATCH 07/10] build the filter mask and offsets in place again, check null lists with one bitmap Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/rebuild.rs | 57 ++++++++++----------- vortex-array/src/arrays/listview/rebuild.rs | 5 +- 2 files changed, 31 insertions(+), 31 deletions(-) diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index 0231620944f..1ed10faf3f7 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; @@ -58,7 +60,8 @@ impl ListArray { ctx: &mut ExecutionCtx, ) -> VortexResult { let (starts, end) = (&offsets[..self.len()], offsets[self.len()].as_()); - let null_runs = null_runs_holding_elements(starts, end, valid); + let null_runs = + null_runs_holding_elements(valid, |index| offsets[index] != offsets[index + 1]); if null_runs.is_empty() && offsets[0].as_() == 0 && end == self.elements().len() { return Ok(self.clone()); } @@ -78,35 +81,20 @@ impl ListArray { } } -/// The runs of null lists that hold elements, among lists laid out back to back: list `i` starts -/// at `starts[i]` and ends where the next list starts, or at `end` for the last list. -pub(crate) fn null_runs_holding_elements( - starts: &[O], - end: usize, +/// The runs of null lists that hold elements, given which lists hold elements. +pub(crate) fn null_runs_holding_elements( valid: &Mask, + mut holds_elements: impl FnMut(usize) -> bool, ) -> Vec<(usize, usize)> { - let start_of = |index: usize| starts.get(index).map_or(end, |start| start.as_()); - let mut runs = Vec::new(); - // Lists never start before the previous one, so a run of null lists holds elements if and - // only if the list after it starts later than the run does. - let mut push_if_holding = |run_start: usize, run_end: usize| { - if run_start < run_end && start_of(run_start) != start_of(run_end) { - runs.push((run_start, run_end)); - } - }; match valid.bit_buffer() { - AllOr::All => {} - AllOr::None => push_if_holding(0, valid.len()), - AllOr::Some(bits) => { - let mut run_start = 0; - // The empty slice at the end closes a run of null lists at the end. - for (valid_start, valid_end) in bits.set_slices().chain([(bits.len(), bits.len())]) { - push_if_holding(run_start, valid_start); - run_start = valid_end; - } - } + AllOr::All => Vec::new(), + AllOr::None if (0..valid.len()).any(&mut holds_elements) => vec![(0, valid.len())], + AllOr::None => Vec::new(), + AllOr::Some(valid) => BitBuffer::collect_bool(valid.len(), &mut holds_elements) + .bitand_not(valid) + .set_slices() + .collect(), } - runs } /// Drops the elements before the first of the lists laid out back to back, after the last, and @@ -138,22 +126,31 @@ pub(crate) fn drop_unreferenced_elements( let elements = if let [(start, stop)] = kept[..] { referenced.slice(start..stop)? } else { + let mut bits = BitBufferMut::new_unset(end - first); + for &(start, stop) in &kept { + bits.fill_range(start, stop, true); + } // With the ranges cached on the mask, the filter copies the elements range by range. - referenced.filter(Mask::from_slices(end - first, kept))? + referenced.filter(Mask::from_buffer_with_slices(bits.freeze(), kept))? }; let mut offsets = BufferMut::with_capacity_in(starts.len() + 1, allocator.clone()); + offsets.extend_from_slice(starts); let mut dropped = starts.first().copied().unwrap_or_else(O::zero); let mut from = 0; for &(run_start, run_end) in null_runs { - offsets.extend(starts[from..run_start].iter().map(|&start| start - dropped)); - offsets.push_n(starts[run_start] - dropped, run_end - run_start); + for offset in &mut offsets[from..run_start] { + *offset = *offset - dropped; + } + offsets[run_start..run_end].fill(starts[run_start] - dropped); // A run at the end leaves no list after it to move down. if let Some(&next) = starts.get(run_end) { dropped += next - starts[run_start]; } from = run_end; } - offsets.extend(starts[from..].iter().map(|&start| start - dropped)); + for offset in &mut offsets[from..] { + *offset = *offset - dropped; + } Ok((elements, offsets)) } diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 198dba86e36..7cb7f38a0d4 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -174,7 +174,10 @@ impl ListViewArray { valid: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult { - let null_runs = null_runs_holding_elements(starts, end, valid); + // A list holds elements if and only if the next one starts later. + let null_runs = null_runs_holding_elements(valid, |index| { + starts[index].as_() != starts.get(index + 1).map_or(end, |next| next.as_()) + }); if null_runs.is_empty() && starts[0].as_() == 0 && end == self.elements().len() { return Ok(self.clone()); } From 563d4d351a529030ad2e87c6153622ccd9a718ea Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 16:37:22 +0000 Subject: [PATCH 08/10] make the list view rebuild a rebuild mode, MakeExactWithEmptyNulls Signed-off-by: Matt Katz --- vortex-array/src/arrays/listview/rebuild.rs | 50 ++++++++++++--------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 7cb7f38a0d4..96624b20b74 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -88,7 +88,9 @@ pub enum ListViewRebuildMode { /// Equivalent to `MakeZeroCopyToList` plus `TrimElements`. /// /// An array that is not yet zero-copy to a [`ListArray`] is copied with only its valid lists' - /// elements, and its null lists come out empty. + /// elements, and its null lists come out empty. One that already is zero-copy only loses the + /// elements before its first list and after its last, so its null lists keep theirs; use + /// [`MakeExactWithEmptyNulls`](Self::MakeExactWithEmptyNulls) to drop those too. /// /// This is useful when concatenating multiple [`ListViewArray`]s together to create a new /// [`ListViewArray`] that is also zero-copy to a [`ListArray`]. @@ -96,6 +98,16 @@ pub enum ListViewRebuildMode { /// [`ListArray`]: crate::arrays::ListArray MakeExact, + /// Equivalent to `MakeExact`, but the null lists of an array that is already zero-copy to a + /// [`ListArray`] come out empty as well, so only the valid lists' elements are left. + /// + /// This is useful for a function that must not see elements no valid list shows, such as one + /// that can fail on a value. An array that is already zero-copy to a [`ListArray`] keeps its + /// offsets' and sizes' types, and is returned as it is when there is nothing to drop. + /// + /// [`ListArray`]: crate::arrays::ListArray + MakeExactWithEmptyNulls, + // TODO(connor)[ListView]: Implement some version of this. /// Finds the shortest packing / overlapping of list elements. /// @@ -121,27 +133,23 @@ impl ListViewArray { ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(ctx), ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(ctx), ListViewRebuildMode::MakeExact => self.rebuild_make_exact(ctx), + ListViewRebuildMode::MakeExactWithEmptyNulls => { + self.rebuild_make_exact_with_empty_nulls(ctx) + } ListViewRebuildMode::OverlapCompression => { vortex_bail!("ListViewRebuildMode::OverlapCompression is not implemented") } } } - /// 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. - /// - /// The result is zero-copy to a [`ListArray`] and its null lists are empty. A non-empty array - /// that is already zero-copy to a [`ListArray`] keeps its offsets' and sizes' types, and is - /// returned as it is when there is nothing to drop. Any other array is rebuilt with - /// [`MakeExact`](ListViewRebuildMode::MakeExact). - /// - /// [`ListArray`]: crate::arrays::ListArray - pub fn rebuild_with_referenced_valid_elements( + /// Rebuilds a non-empty [`ListViewArray`] over only the elements that valid lists reference: + /// nothing before the first list, after the last, between lists or under a null list. + fn rebuild_make_exact_with_empty_nulls( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { - if self.is_empty() || !self.is_zero_copy_to_list() { - return self.rebuild(ListViewRebuildMode::MakeExact, ctx); + if !self.is_zero_copy_to_list() { + return self.rebuild_make_exact(ctx); } let validity = self.validity()?; // The lists are back to back, so the last one ends where the referenced elements do. @@ -164,8 +172,8 @@ impl ListViewArray { }) } - /// [`rebuild_with_referenced_valid_elements`](Self::rebuild_with_referenced_valid_elements) - /// for an array that is zero-copy to a list, whose lists start at `starts` and end at `end`. + /// [`MakeExactWithEmptyNulls`](ListViewRebuildMode::MakeExactWithEmptyNulls) for an array that + /// is zero-copy to a list, whose lists start at `starts` and end at `end`. fn rebuild_back_to_back( &self, starts: &[O], @@ -651,7 +659,7 @@ mod tests { &[1, 2, 3, 2, 3], &[(0, 3), (3, 2)] )] - fn rebuild_with_referenced_valid_elements( + fn test_rebuild_make_exact_with_empty_nulls( #[case] elements: &[i32], #[case] views: &[(u32, u32)], #[case] validity: &[bool], @@ -669,7 +677,7 @@ mod tests { )?; // 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)?; + let rebuilt = view.rebuild(ListViewRebuildMode::MakeExactWithEmptyNulls, &mut ctx)?; assert!(rebuilt.is_zero_copy_to_list()); assert_arrays_eq!( rebuilt.elements().clone(), @@ -693,7 +701,7 @@ mod tests { #[rstest] #[case::non_nullable(Validity::NonNullable)] #[case::empty_null_list(Validity::from_iter([true, false, true]))] - fn rebuild_with_referenced_valid_elements_returns_an_exact_list_view_as_is( + fn test_rebuild_make_exact_with_empty_nulls_returns_an_exact_list_view_as_is( #[case] validity: Validity, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -706,7 +714,7 @@ mod tests { )?; // SAFETY: the lists are sorted with nothing between them. let view = unsafe { view.with_zero_copy_to_list(true) }; - let rebuilt = view.rebuild_with_referenced_valid_elements(&mut ctx)?; + let rebuilt = view.rebuild(ListViewRebuildMode::MakeExactWithEmptyNulls, &mut ctx)?; assert!(ArrayRef::ptr_eq(rebuilt.elements(), view.elements())); assert!(ArrayRef::ptr_eq(rebuilt.offsets(), view.offsets())); assert!(ArrayRef::ptr_eq(rebuilt.sizes(), view.sizes())); @@ -802,7 +810,8 @@ mod tests { )?; // SAFETY: the views are a `ListArray`'s lists, sorted 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)?; + let rebuilt = + view.rebuild(ListViewRebuildMode::MakeExactWithEmptyNulls, &mut ctx)?; assert!(rebuilt.is_zero_copy_to_list()); assert_arrays_eq!( rebuilt.elements().clone(), @@ -950,6 +959,7 @@ mod tests { #[case::zero_copy_to_list(ListViewRebuildMode::MakeZeroCopyToList)] #[case::trim_elements(ListViewRebuildMode::TrimElements)] #[case::make_exact(ListViewRebuildMode::MakeExact)] + #[case::make_exact_with_empty_nulls(ListViewRebuildMode::MakeExactWithEmptyNulls)] fn test_rebuild_empty_drops_elements(#[case] mode: ListViewRebuildMode) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); // No lists, but three elements left in the buffer. From 98d8115eac9b20497d078b40c64047b801e49e08 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 16:57:51 +0000 Subject: [PATCH 09/10] rename ListArray::rebuild_with_referenced_valid_elements to compact Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/rebuild.rs | 7 ++----- vortex-array/src/arrays/list/tests.rs | 14 ++++++-------- vortex-array/src/arrays/listview/rebuild.rs | 4 ++-- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index 1ed10faf3f7..a070b7247ff 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -24,16 +24,13 @@ use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; impl ListArray { - /// Rebuilds the [`ListArray`] over only the elements that valid lists reference: nothing before + /// Compacts the [`ListArray`] to only the elements that valid lists reference: nothing before /// the first list, after the last or under a null list. /// /// Null lists come out empty, the offsets start at 0, end at the length of the elements and /// keep their type, and the validity is kept. An array with nothing to drop is returned as it /// is. - pub fn rebuild_with_referenced_valid_elements( - &self, - ctx: &mut ExecutionCtx, - ) -> VortexResult { + pub fn compact(&self, ctx: &mut ExecutionCtx) -> VortexResult { let validity = self.list_validity(); if validity.definitely_no_nulls() && self.offset_at(0, ctx)? == 0 diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 7354bb0b41a..d1108df3840 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -1104,7 +1104,7 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { &[0, 0, 0] )] #[case::no_lists_with_a_leftover(&[9], &[1], &[], &[], &[0])] -fn rebuild_with_referenced_valid_elements( +fn test_compact( #[case] elements: &[i32], #[case] offsets: &[u32], #[case] validity: &[bool], @@ -1122,7 +1122,7 @@ fn rebuild_with_referenced_valid_elements( }, &mut ctx, )?; - let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + let rebuilt = list.compact(&mut ctx)?; assert_arrays_eq!( rebuilt.elements().clone(), PrimitiveArray::from_iter(expected_elements.iter().copied()), @@ -1141,9 +1141,7 @@ fn rebuild_with_referenced_valid_elements( #[case::non_nullable(Validity::NonNullable)] #[case::all_valid(Validity::AllValid)] #[case::empty_null_list(Validity::from_iter([true, false, true]))] -fn rebuild_with_referenced_valid_elements_returns_an_exact_list_as_is( - #[case] validity: Validity, -) -> VortexResult<()> { +fn test_compact_returns_an_exact_list_as_is(#[case] validity: Validity) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let list = ListArray::try_new( buffer![1i32, 2, 3].into_array(), @@ -1151,7 +1149,7 @@ fn rebuild_with_referenced_valid_elements_returns_an_exact_list_as_is( validity, &mut ctx, )?; - let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + let rebuilt = list.compact(&mut ctx)?; assert!(ArrayRef::ptr_eq(rebuilt.elements(), list.elements())); assert!(ArrayRef::ptr_eq(rebuilt.offsets(), list.offsets())); Ok(()) @@ -1162,7 +1160,7 @@ fn rebuild_with_referenced_valid_elements_returns_an_exact_list_as_is( #[case::no_null_list(2..4, &[3, 4, 5, 6], &[0, 2, 4])] #[case::a_null_list_holding_elements_first(1..4, &[3, 4, 5, 6], &[0, 0, 2, 4])] #[case::a_null_list_holding_elements_last(3..5, &[5, 6], &[0, 2, 2])] -fn rebuild_with_referenced_valid_elements_sliced( +fn test_compact_sliced( #[case] range: Range, #[case] expected_elements: &[i32], #[case] expected_offsets: &[u32], @@ -1179,7 +1177,7 @@ fn rebuild_with_referenced_valid_elements_sliced( .slice(range)? .as_::() .into_owned(); - let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + let rebuilt = list.compact(&mut ctx)?; assert_arrays_eq!( rebuilt.elements().clone(), PrimitiveArray::from_iter(expected_elements.iter().copied()), diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 96624b20b74..a3cf6ff2188 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -738,7 +738,7 @@ mod tests { /// and as list views with and without the zero-copy flag, against a reference that gathers /// each valid list's elements back to back. #[test] - fn rebuild_with_referenced_valid_elements_matches_a_reference() -> VortexResult<()> { + fn test_compact_and_make_exact_with_empty_nulls_match_a_reference() -> VortexResult<()> { const PTYPES: [PType; 5] = [PType::U8, PType::U16, PType::U64, PType::I32, PType::I64]; let u32_dtype = DType::Primitive(PType::U32, Nullability::NonNullable); let mut ctx = SESSION.create_execution_ctx(); @@ -786,7 +786,7 @@ mod tests { validity.clone(), &mut ctx, )?; - let rebuilt = list.rebuild_with_referenced_valid_elements(&mut ctx)?; + let rebuilt = list.compact(&mut ctx)?; assert_arrays_eq!( rebuilt.elements().clone(), expected_elements.clone(), From 692502fd5d545d57b3a458e9c99453948cc6b7b6 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 26 Oct 2026 18:02:29 +0000 Subject: [PATCH 10/10] walk the null runs once, skip the filter when no elements are kept Signed-off-by: Matt Katz --- vortex-array/src/arrays/list/rebuild.rs | 48 ++++++++++++++----------- 1 file changed, 27 insertions(+), 21 deletions(-) diff --git a/vortex-array/src/arrays/list/rebuild.rs b/vortex-array/src/arrays/list/rebuild.rs index a070b7247ff..d7c895be089 100644 --- a/vortex-array/src/arrays/list/rebuild.rs +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -97,7 +97,12 @@ pub(crate) fn null_runs_holding_elements( /// Drops the elements before the first of the lists laid out back to back, after the last, and /// under the runs of null lists from [`null_runs_holding_elements`], which come out empty. /// -/// Returns the elements kept and where each list now starts, with room to push one more offset. +/// `end` is where the last list ends. It is a `usize` rather than an `O` because a list view's last +/// list can end past what its offsets' type holds: with `u8` offsets `[1, 5, 105]` and sizes +/// `[4, 100, 200]`, it ends at 305. +/// +/// Returns the elements kept and where each list now starts. A [`ListArray`] needs one more offset, +/// where its last list ends, so the buffer has room to push it without reallocating. pub(crate) fn drop_unreferenced_elements( elements: &ArrayRef, starts: &[O], @@ -110,32 +115,16 @@ pub(crate) fn drop_unreferenced_elements( // The ranges of elements kept, counted from the first list's start. let mut kept = Vec::with_capacity(null_runs.len() + 1); let mut kept_start = 0; + let mut offsets = BufferMut::with_capacity_in(starts.len() + 1, allocator.clone()); + offsets.extend_from_slice(starts); + let mut dropped = starts.first().copied().unwrap_or_else(O::zero); + let mut from = 0; for &(run_start, run_end) in null_runs { if kept_start < start_of(run_start) - first { kept.push((kept_start, start_of(run_start) - first)); } kept_start = start_of(run_end) - first; - } - if kept_start < end - first { - kept.push((kept_start, end - first)); - } - let referenced = elements.slice(first..end)?; - let elements = if let [(start, stop)] = kept[..] { - referenced.slice(start..stop)? - } else { - let mut bits = BitBufferMut::new_unset(end - first); - for &(start, stop) in &kept { - bits.fill_range(start, stop, true); - } - // With the ranges cached on the mask, the filter copies the elements range by range. - referenced.filter(Mask::from_buffer_with_slices(bits.freeze(), kept))? - }; - let mut offsets = BufferMut::with_capacity_in(starts.len() + 1, allocator.clone()); - offsets.extend_from_slice(starts); - let mut dropped = starts.first().copied().unwrap_or_else(O::zero); - let mut from = 0; - for &(run_start, run_end) in null_runs { for offset in &mut offsets[from..run_start] { *offset = *offset - dropped; } @@ -146,8 +135,25 @@ pub(crate) fn drop_unreferenced_elements( } from = run_end; } + if kept_start < end - first { + kept.push((kept_start, end - first)); + } for offset in &mut offsets[from..] { *offset = *offset - dropped; } + + let referenced = elements.slice(first..end)?; + let elements = match kept[..] { + [] => referenced.slice(0..0)?, + [(start, stop)] => referenced.slice(start..stop)?, + _ => { + let mut bits = BitBufferMut::new_unset(end - first); + for &(start, stop) in &kept { + bits.fill_range(start, stop, true); + } + // With the ranges cached on the mask, the filter copies the elements range by range. + referenced.filter(Mask::from_buffer_with_slices(bits.freeze(), kept))? + } + }; Ok((elements, offsets)) }