vortex_array/arrays/listview/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hash;
5use std::hash::Hasher;
6use std::sync::Arc;
7
8use prost::Message;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_ensure;
13use vortex_error::vortex_ensure_eq;
14use vortex_error::vortex_panic;
15use vortex_session::VortexSession;
16use vortex_session::registry::CachedId;
17
18use crate::ArrayEq;
19use crate::ArrayHash;
20use crate::ArrayParts;
21use crate::ArrayRef;
22use crate::EqMode;
23use crate::ExecutionCtx;
24use crate::ExecutionResult;
25use crate::array::Array;
26use crate::array::ArrayId;
27use crate::array::ArrayView;
28use crate::array::VTable;
29use crate::array::with_empty_buffers;
30use crate::arrays::listview::ListViewArraySlotsExt;
31use crate::arrays::listview::ListViewData;
32use crate::arrays::listview::ListViewSlots;
33use crate::arrays::listview::compute::rules::PARENT_RULES;
34use crate::buffer::BufferHandle;
35use crate::builders::ArrayBuilder;
36use crate::dtype::DType;
37use crate::dtype::Nullability;
38use crate::dtype::PType;
39use crate::match_each_list_builder;
40use crate::serde::ArrayChildren;
41use crate::validity::Validity;
42mod kernel;
43mod operations;
44mod validity;
45/// A [`ListView`]-encoded Vortex array.
46pub type ListViewArray = Array<ListView>;
47
48pub(crate) fn initialize(session: &VortexSession) {
49    kernel::initialize(session);
50}
51
52#[derive(Clone, Debug)]
53pub struct ListView;
54
55#[derive(Clone, prost::Message)]
56pub struct ListViewMetadata {
57    #[prost(uint64, tag = "1")]
58    elements_len: u64,
59    #[prost(enumeration = "PType", tag = "2")]
60    offset_ptype: i32,
61    #[prost(enumeration = "PType", tag = "3")]
62    size_ptype: i32,
63}
64
65impl ArrayHash for ListViewData {
66    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
67        self.is_zero_copy_to_list().hash(state);
68    }
69}
70
71impl ArrayEq for ListViewData {
72    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
73        self.is_zero_copy_to_list() == other.is_zero_copy_to_list()
74    }
75}
76
77impl VTable for ListView {
78    type TypedArrayData = ListViewData;
79
80    type OperationsVTable = Self;
81    type ValidityVTable = Self;
82    fn id(&self) -> ArrayId {
83        static ID: CachedId = CachedId::new("vortex.listview");
84        *ID
85    }
86
87    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
88        0
89    }
90
91    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
92        vortex_panic!("ListViewArray buffer index {idx} out of bounds")
93    }
94
95    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
96        vortex_panic!("ListViewArray buffer_name index {idx} out of bounds")
97    }
98
99    fn with_buffers(
100        &self,
101        array: ArrayView<'_, Self>,
102        buffers: &[BufferHandle],
103    ) -> VortexResult<ArrayParts<Self>> {
104        with_empty_buffers(self, array, buffers)
105    }
106
107    fn serialize(
108        array: ArrayView<'_, Self>,
109        _session: &VortexSession,
110    ) -> VortexResult<Option<Vec<u8>>> {
111        Ok(Some(
112            ListViewMetadata {
113                elements_len: array.elements().len() as u64,
114                offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
115                size_ptype: PType::try_from(array.sizes().dtype())? as i32,
116            }
117            .encode_to_vec(),
118        ))
119    }
120
121    fn validate(
122        &self,
123        _data: &ListViewData,
124        dtype: &DType,
125        len: usize,
126        slots: &[Option<ArrayRef>],
127    ) -> VortexResult<()> {
128        vortex_ensure_eq!(slots.len(), ListViewSlots::COUNT);
129        let elements = slots[ListViewSlots::ELEMENTS]
130            .as_ref()
131            .vortex_expect("ListViewArray elements slot");
132        let offsets = slots[ListViewSlots::OFFSETS]
133            .as_ref()
134            .vortex_expect("ListViewArray offsets slot");
135        let sizes = slots[ListViewSlots::SIZES]
136            .as_ref()
137            .vortex_expect("ListViewArray sizes slot");
138        vortex_ensure!(
139            offsets.len() == len && sizes.len() == len,
140            "ListViewArray length {} does not match outer length {}",
141            offsets.len(),
142            len
143        );
144
145        let actual_dtype = DType::List(Arc::new(elements.dtype().clone()), dtype.nullability());
146        vortex_ensure_eq!(
147            &actual_dtype,
148            dtype,
149            "ListViewArray dtype does not match outer dtype",
150        );
151
152        Ok(())
153    }
154
155    fn deserialize(
156        &self,
157        dtype: &DType,
158        len: usize,
159        metadata: &[u8],
160
161        buffers: &[BufferHandle],
162        children: &dyn ArrayChildren,
163        _session: &VortexSession,
164    ) -> VortexResult<ArrayParts<Self>> {
165        let metadata = ListViewMetadata::decode(metadata)?;
166        vortex_ensure!(
167            buffers.is_empty(),
168            "`ListViewArray::build` expects no buffers"
169        );
170
171        let DType::List(element_dtype, _) = dtype else {
172            vortex_bail!("Expected List dtype, got {:?}", dtype);
173        };
174
175        let validity = if children.len() == 3 {
176            Validity::from(dtype.nullability())
177        } else if children.len() == 4 {
178            let validity = children.get(3, &Validity::DTYPE, len)?;
179            Validity::Array(validity)
180        } else {
181            vortex_bail!(
182                "`ListViewArray::build` expects 3 or 4 children, got {}",
183                children.len()
184            );
185        };
186
187        // Get elements with the correct length from metadata.
188        let elements = children.get(
189            0,
190            element_dtype.as_ref(),
191            usize::try_from(metadata.elements_len)?,
192        )?;
193
194        // Get offsets with proper type from metadata.
195        let offsets = children.get(
196            1,
197            &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable),
198            len,
199        )?;
200
201        // Get sizes with proper type from metadata.
202        let sizes = children.get(
203            2,
204            &DType::Primitive(metadata.size_ptype(), Nullability::NonNullable),
205            len,
206        )?;
207
208        ListViewData::validate(&elements, &offsets, &sizes, &validity)?;
209        let data = ListViewData::try_new()?;
210        let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
211        Ok(ArrayParts::new(
212            self.clone(),
213            dtype.clone(),
214            len,
215            data,
216            slots,
217        ))
218    }
219
220    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
221        ListViewSlots::NAMES[idx].to_string()
222    }
223
224    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
225        Ok(ExecutionResult::done(array))
226    }
227
228    // The complexity comes from the expansion of `match_each_list_builder!`.
229    #[expect(clippy::cognitive_complexity)]
230    fn append_to_builder(
231        array: ArrayView<'_, Self>,
232        builder: &mut dyn ArrayBuilder,
233        ctx: &mut ExecutionCtx,
234    ) -> VortexResult<()> {
235        match match_each_list_builder!(&mut *builder, |b| b.append_listview_array(array, ctx)) {
236            Some(result) => result,
237            None => vortex_bail!(
238                "cannot append a ListView array of dtype {} to a {} builder",
239                array.dtype(),
240                builder.dtype()
241            ),
242        }
243    }
244
245    fn reduce_parent(
246        array: ArrayView<'_, Self>,
247        parent: &ArrayRef,
248        child_idx: usize,
249    ) -> VortexResult<Option<ArrayRef>> {
250        PARENT_RULES.evaluate(array, parent, child_idx)
251    }
252}
vortex_array/arrays/listview/vtable/mod.rs

vortex_array/arrays/listview/vtable/
mod.rs
