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

vortex_array/arrays/listview/vtable/
mod.rs
