vortex_array/arrays/listview/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::Zero;
5use vortex_dtype::{Nullability, match_each_integer_ptype};
6use vortex_error::VortexResult;
7use vortex_scalar::Scalar;
8
9use crate::arrays::{ListViewArray, ListViewRebuildMode, ListViewVTable};
10use crate::compute::{self, TakeKernel, TakeKernelAdapter};
11use crate::vtable::ValidityHelper;
12use crate::{Array, ArrayRef, IntoArray, register_kernel};
13
14// TODO(connor)[ListView]: Make use of this threshold after we start migrating operators.
15/// The threshold for triggering a rebuild of the [`ListViewArray`].
16///
17/// By default, we will not touch the underlying `elements` array of the [`ListViewArray`] since it
18/// can be potentially expensive to reorganize the array based on what views we have into it.
19///
20/// However, we also do not want to carry around a large amount of garbage data. Below this
21/// threshold of the density of the selection mask, we will rebuild the [`ListViewArray`], removing
22/// any garbage data.
23#[allow(unused)]
24const REBUILD_DENSITY_THRESHOLD: f64 = 0.1;
25
26/// [`ListViewArray`] take implementation.
27///
28/// This implementation is deliberately simple and read-optimized. We just take the `offsets` and
29/// `sizes` at the requested indices and reuse the original `elements` array. This works because
30/// `ListView` (unlike `List`) allows non-contiguous and out-of-order lists.
31///
32/// We don't slice the `elements` array because it would require computing min/max offsets and
33/// adjusting all offsets accordingly, which is not really worth the small potential memory we would
34/// be able to get back.
35///
36/// The trade-off is that we may keep unreferenced elements in memory, but this is acceptable since
37/// we're optimizing for read performance and the data isn't being copied.
38impl TakeKernel for ListViewVTable {
39    fn take(&self, array: &ListViewArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
40        let elements = array.elements();
41        let offsets = array.offsets();
42        let sizes = array.sizes();
43
44        // Compute the new validity by combining the array's validity with the indices' validity.
45        let new_validity = array.validity().take(indices)?;
46
47        // Take the offsets and sizes arrays at the requested indices.
48        // Take can reorder offsets, create gaps, and may introduce overlaps if the `indices`
49        // contain duplicates.
50        let nullable_new_offsets = compute::take(offsets.as_ref(), indices)?;
51        let nullable_new_sizes = compute::take(sizes.as_ref(), indices)?;
52
53        // Since `take` returns nullable arrays, we simply cast it back to non-nullable (filled with
54        // zeros to represent null lists).
55        let new_offsets = match_each_integer_ptype!(nullable_new_offsets.dtype().as_ptype(), |O| {
56            compute::fill_null(
57                &nullable_new_offsets,
58                &Scalar::primitive(O::zero(), Nullability::NonNullable),
59            )?
60        });
61        let new_sizes = match_each_integer_ptype!(nullable_new_sizes.dtype().as_ptype(), |S| {
62            compute::fill_null(
63                &nullable_new_sizes,
64                &Scalar::primitive(S::zero(), Nullability::NonNullable),
65            )?
66        });
67        // SAFETY: Take operation maintains all `ListViewArray` invariants:
68        // - `new_offsets` and `new_sizes` are derived from existing valid child arrays.
69        // - `new_offsets` and `new_sizes` are non-nullable.
70        // - `new_offsets` and `new_sizes` have the same length (both taken with the same
71        //   `indices`).
72        // - Validity correctly reflects the combination of array and indices validity.
73        let new_array = unsafe {
74            ListViewArray::new_unchecked(elements.clone(), new_offsets, new_sizes, new_validity)
75        };
76
77        // TODO(connor)[ListView]: Ideally, we would only rebuild after all `take`s and `filter`
78        // compute functions have run, at the "top" of the operator tree. However, we cannot do this
79        // right now, so we will just rebuild every time (similar to `ListArray`).
80
81        Ok(new_array
82            .rebuild(ListViewRebuildMode::MakeZeroCopyToList)
83            .into_array())
84    }
85}
86
87register_kernel!(TakeKernelAdapter(ListViewVTable).lift());
vortex_array/arrays/listview/compute/take.rs

vortex_array/arrays/listview/compute/
take.rs
