Rebuild lists over only the elements valid lists reference

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;
 

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<ListArray> {
+        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::<PrimitiveArray>(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::<O>(), offsets_ptype, &valid, ctx)
+        })
+    }
+
+    fn rebuild_with_valid<O: IntegerPType>(
+        &self,
+        offsets: &[O],
+        offsets_ptype: PType,
+        valid: &Mask,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ListArray> {
+        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<O: IntegerPType>(
+    elements: &ArrayRef,
+    starts: &[O],
+    end: usize,
+    null_runs: &[(usize, usize)],
+    allocator: &BufferAllocatorRef,
+) -> VortexResult<(ArrayRef, BufferMut<O>)> {
+    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))
+}

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<usize>,
+    #[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_::<List>()
+    .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(())
+}