vortex_array/arrays/map/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_error::vortex_ensure;
7use vortex_error::vortex_ensure_eq;
8use vortex_error::vortex_panic;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayParts;
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::ExecutionResult;
16use crate::array::Array;
17use crate::array::ArrayId;
18use crate::array::ArrayView;
19use crate::array::VTable;
20use crate::array::ValidityVTableFromChild;
21use crate::array::with_empty_buffers;
22use crate::arrays::ListView;
23use crate::arrays::map::MapData;
24use crate::arrays::map::MapSlots;
25use crate::arrays::map::MapSlotsView;
26use crate::arrays::map::array::validate_entries;
27use crate::arrays::map::compute::rules::PARENT_RULES;
28use crate::buffer::BufferHandle;
29use crate::builders::ArrayBuilder;
30use crate::dtype::DType;
31use crate::match_each_map_builder;
32use crate::serde::ArrayChildren;
33
34mod kernel;
35mod operations;
36mod validity;
37
38/// A [`Map`]-encoded Vortex array.
39pub type MapArray = Array<Map>;
40
41pub(crate) fn initialize(session: &VortexSession) {
42    kernel::initialize(session);
43}
44
45/// The canonical encoding for [`DType::Map`].
46///
47/// A map array has one `ListView<Struct<key, value>>` child. Its outer dtype retains map-specific
48/// metadata such as the `keys_sorted` assertion.
49#[derive(Clone, Debug, Default)]
50pub struct Map;
51
52impl VTable for Map {
53    type TypedArrayData = MapData;
54
55    type OperationsVTable = Self;
56    type ValidityVTable = ValidityVTableFromChild;
57
58    fn id(&self) -> ArrayId {
59        static ID: CachedId = CachedId::new("vortex.map");
60        *ID
61    }
62
63    fn validate(
64        &self,
65        _data: &MapData,
66        dtype: &DType,
67        len: usize,
68        slots: &[Option<ArrayRef>],
69    ) -> VortexResult<()> {
70        vortex_ensure_eq!(slots.len(), MapSlots::COUNT);
71
72        let DType::Map(map_dtype, nullability) = dtype else {
73            vortex_bail!("Expected map dtype, got {dtype}");
74        };
75        let slots = MapSlotsView::from_slots(slots);
76        validate_entries(map_dtype, *nullability, len, slots.entries)
77    }
78
79    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
80        0
81    }
82
83    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
84        vortex_panic!("MapArray buffer index {idx} out of bounds")
85    }
86
87    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
88        None
89    }
90
91    fn with_buffers(
92        &self,
93        array: ArrayView<'_, Self>,
94        buffers: &[BufferHandle],
95    ) -> VortexResult<ArrayParts<Self>> {
96        with_empty_buffers(self, array, buffers)
97    }
98
99    fn serialize(
100        _array: ArrayView<'_, Self>,
101        _session: &VortexSession,
102    ) -> VortexResult<Option<Vec<u8>>> {
103        Ok(Some(vec![]))
104    }
105
106    fn deserialize(
107        &self,
108        dtype: &DType,
109        len: usize,
110        metadata: &[u8],
111        buffers: &[BufferHandle],
112        children: &dyn ArrayChildren,
113        _session: &VortexSession,
114    ) -> VortexResult<ArrayParts<Self>> {
115        if !metadata.is_empty() {
116            vortex_bail!(
117                "MapArray expects empty metadata, got {} bytes",
118                metadata.len()
119            );
120        }
121        vortex_ensure!(buffers.is_empty(), "MapArray expects no buffers");
122
123        let DType::Map(map_dtype, nullability) = dtype else {
124            vortex_bail!("Expected map dtype, got {dtype}");
125        };
126        vortex_ensure_eq!(children.len(), MapSlots::COUNT);
127
128        let expected_entries_dtype =
129            DType::List(std::sync::Arc::new(map_dtype.entries_dtype()), *nullability);
130        let entries = children.get(MapSlots::ENTRIES, &expected_entries_dtype, len)?;
131        vortex_ensure!(
132            entries.is::<ListView>(),
133            "MapArray entries must use vortex.listview encoding, got {}",
134            entries.encoding_id()
135        );
136
137        let slots = MapData::make_slots(entries);
138        Ok(ArrayParts::new(
139            self.clone(),
140            dtype.clone(),
141            len,
142            MapData,
143            slots,
144        ))
145    }
146
147    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
148        MapSlots::NAMES[idx].to_string()
149    }
150
151    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
152        Ok(ExecutionResult::done(array))
153    }
154
155    fn append_to_builder(
156        array: ArrayView<'_, Self>,
157        builder: &mut dyn ArrayBuilder,
158        ctx: &mut ExecutionCtx,
159    ) -> VortexResult<()> {
160        match match_each_map_builder!(&mut *builder, |b| b.append_map_array(array, ctx)) {
161            Some(result) => result,
162            None => vortex_bail!(
163                "cannot append a Map array of dtype {} to a {} builder",
164                array.dtype(),
165                builder.dtype()
166            ),
167        }
168    }
169
170    fn reduce_parent(
171        array: ArrayView<'_, Self>,
172        parent: &ArrayRef,
173        child_idx: usize,
174    ) -> VortexResult<Option<ArrayRef>> {
175        PARENT_RULES.evaluate(array, parent, child_idx)
176    }
177}
vortex_array/arrays/map/vtable/mod.rs

vortex_array/arrays/map/vtable/
mod.rs
