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

vortex_array/arrays/map/vtable/
mod.rs
