Rebuild lists over only the elements valid lists reference

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;
 

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<ListArray> {
+        let offsets = self.offsets().clone().execute::<PrimitiveArray>(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::<O>(), valid)
+        })
+    }
+
+    fn rebuild_with_valid_runs<O: IntegerPType>(
+        &self,
+        offsets: &[O],
+        valid: &[(usize, usize)],
+    ) -> VortexResult<ListArray> {
+        let kept = referenced_valid_ranges(offsets, valid);
+        if kept.iter().map(|(start, end)| end - start).sum::<usize>() == 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<O: IntegerPType>(
+    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<O: IntegerPType>(offsets: &[O], valid: &[(usize, usize)]) -> BufferMut<O> {
+    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
+}

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(())
+}

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<ListViewArray> {
+        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::<PrimitiveArray>(ctx)?;
+        let null_lists_empty = match_each_integer_ptype!(sizes.ptype(), |S| {
+            let sizes = sizes.as_slice::<S>();
+            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();