diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index bfd43cdb401..bbbba9563ca 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -12,6 +12,10 @@ pub use vtable::ListArray; pub(crate) mod compute; +mod rebuild; +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 new file mode 100644 index 00000000000..d7c895be089 --- /dev/null +++ b/vortex-array/src/arrays/list/rebuild.rs @@ -0,0 +1,159 @@ +// 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; +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::dtype::PType; +use crate::match_each_unsigned_integer_ptype; +use crate::validity::Validity; + +impl ListArray { + /// 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 compact(&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 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_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 (starts, end) = (&offsets[..self.len()], offsets[self.len()].as_()); + 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()); + } + + 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 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)> { + match valid.bit_buffer() { + 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(), + } +} + +/// 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. +/// +/// `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], + 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; + 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; + + 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; + } + 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)) +} diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 08792fc251a..d1108df3840 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; +use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -12,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; @@ -1058,3 +1061,133 @@ 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::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], + &[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::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])] +#[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 test_compact( + #[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(), + if validity.iter().all(|valid| *valid) { + Validity::NonNullable + } else { + Validity::from_iter(validity.iter().copied()) + }, + &mut ctx, + )?; + let rebuilt = list.compact(&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(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +#[case::empty_null_list(Validity::from_iter([true, false, true]))] +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(), + buffer![0u32, 1, 1, 3].into_array(), + validity, + &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(()) +} + +/// 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 test_compact_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.compact(&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..a3cf6ff2188 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,6 +17,8 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; +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::primitive::PrimitiveArrayExt; @@ -84,12 +87,27 @@ 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. 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`]. /// /// [`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. /// @@ -115,12 +133,88 @@ 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 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_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. + 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 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) + }) + } + + /// [`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], + offsets_ptype: PType, + end: usize, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // 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()); + } + + 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. /// /// This is useful when the `elements` child array of the [`ListViewArray`] might have @@ -389,6 +483,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, @@ -454,22 +565,276 @@ 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] + #[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)] + )] + #[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 test_rebuild_make_exact_with_empty_nulls( + #[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(ListViewRebuildMode::MakeExactWithEmptyNulls, &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(()) + } + + #[rstest] + #[case::non_nullable(Validity::NonNullable)] + #[case::empty_null_list(Validity::from_iter([true, false, true]))] + 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(); + 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(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())); + 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 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(); + 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.compact(&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(ListViewRebuildMode::MakeExactWithEmptyNulls, &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(); @@ -594,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.