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

vortex_array/arrays/listview/vtable/
mod.rs
