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

vortex_array/arrays/listview/vtable/
mod.rs
