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..21845676422 --- /dev/null +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -0,0 +1,104 @@ +// 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 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_with_valid_runs(offsets.as_slice::(), valid) + }) + } + + fn rebuild_with_valid_runs( + &self, + offsets: &[O], + valid: &[(usize, usize)], + ) -> VortexResult { + let kept = referenced_valid_ranges(offsets, valid); + 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_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)), + } + } + 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; + } + 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 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..9d97d36beb7 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,57 @@ 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 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::(); + 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 + }) + }); + 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 +523,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();