Rebuild lists over only the elements valid lists reference (#10743)

## Summary

Adds `ListArray::compact` and a list view rebuild mode,
`ListViewRebuildMode::MakeExactWithEmptyNulls`. Both return the same
lists over only the elements that valid lists reference: nothing before
the first list, after the last, between lists or under a null list.

A list column can hold elements that no valid list shows. Slicing keeps
the parent's elements on both sides, and a null list may still cover
elements. A function that only reads values does not care, but a
function that can fail on an element, such as a checked `x + 1`, must
not see them: an overflow in an element under a null list would fail the
whole batch. `list_transform` (#10503, stacked on this PR) is the first
user; it applies its lambda to the rebuilt elements whenever the lambda
can fail.

For example:

```text
elements  [9, 1, 2, 9, 9, 3]
offsets   [1, 3, 5, 6]
validity  [valid, null, valid]
lists     [1, 2], null over [9, 9], [3]

rebuilt
elements  [1, 2, 3]
offsets   [0, 2, 2, 3]
```

## How it works

**`ListArray`**

- Non-nullable or all-valid lists with nothing at the ends are returned
as they are, after reading two offsets.
- Otherwise it looks for null lists that hold elements: one pass marks
the lists whose next offset is larger, and one bit operation clears the
valid ones. Only when bits are left does it collect them into runs.
- With no such run, only the ends go: the elements are sliced to `[first
offset, last offset)` and the offsets are shifted down, without copying
an element.
- With some, the elements between those runs are kept, by a slice if one
range is left and otherwise by one filter whose mask carries the ranges,
so it copies range by range. The offsets are copied once and each one
moves down by the elements dropped before it, so the null lists in those
runs become empty.
- The offsets keep their type and the validity is unchanged.

**`ListViewRebuildMode::MakeExactWithEmptyNulls`**

- `MakeExact` only trims the ends of a view that is already zero-copy to
a list, so its null lists keep their elements. The new mode drops those
too.
- Such a view has its lists back to back, so it goes through the same
two steps on its offsets, with the last list's end as the final
boundary. It keeps its layout and its offsets' and sizes' types. Only
the sizes of the null lists that held elements are set to 0; otherwise
the sizes are kept as they are.
- Any other view is rebuilt with `MakeExact`, which already copies only
the valid lists' elements and leaves null lists empty. `MakeExact`'s
documentation now says both.
- The result is always zero-copy to a list.

## Benchmark

`list_transform` with a checked `x + 1` over 100k lists (~800k `i32`
elements, ~1/8 of the lists null), through #10503, exported to Arrow,
median µs over three alternating runs of a local divan benchmark.
"Before" is #10503 without this PR, which used the existing list view
rebuild instead.

| Lists | Before | This PR | DataFusion 54.1 `array_transform` |
|---|---|---|---|
| Null lists holding elements | 1,272 | 253 | 224 |
| Null lists already empty | 97 | 73 | 60 |
| Sliced, with elements on both sides | 78 | 81 | 55 |

## Testing

- `rstest` cases for both types: ends to drop, null lists holding
elements or not, every list null, no lists, an empty valid list between
null lists, overlapping and out-of-order views, views with and without
the zero-copy flag.
- Exact arrays come back as the same arrays (`ArrayRef::ptr_eq`), and a
sliced `ListArray` whose validity starts mid-byte is rebuilt correctly.
- A seeded random test compares `ListArray` and both kinds of view
against a reference that gathers each valid list's elements, over five
offset types, with non-nullable validity and validity arrays that start
mid-byte.

---------

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

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