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

vortex_array/arrays/listview/vtable/
mod.rs
