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

vortex_array/arrays/listview/vtable/
mod.rs
