vortex_array/builders/
map.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::sync::Arc;
6
7use vortex_error::VortexResult;
8use vortex_error::vortex_ensure;
9
10use crate::ArrayRef;
11use crate::Canonical;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::array::ArrayView;
15use crate::arrays::ListView;
16use crate::arrays::Map;
17use crate::arrays::MapArray;
18use crate::arrays::map::MapArraySlotsExt;
19use crate::builders::ArrayBuilder;
20use crate::builders::DEFAULT_BUILDER_CAPACITY;
21use crate::builders::ListViewBuilder;
22use crate::dtype::DType;
23use crate::dtype::MapDType;
24use crate::dtype::Nullability;
25use crate::dtype::OffsetBuilderPType;
26use crate::scalar::MapScalar;
27use crate::scalar::Scalar;
28
29/// A builder for canonical [`MapArray`] values.
30///
31/// The builder owns a [`ListViewBuilder`] whose elements are non-nullable `{key, value}` structs.
32/// It preserves the map dtype's `keys_sorted` assertion while delegating offsets, sizes, and outer
33/// validity to that list-view builder.
34pub struct MapBuilder<O: OffsetBuilderPType, S: OffsetBuilderPType> {
35    dtype: DType,
36    map_dtype: MapDType,
37    entries_builder: ListViewBuilder<O, S>,
38}
39
40impl<O: OffsetBuilderPType, S: OffsetBuilderPType> MapBuilder<O, S> {
41    /// Creates a map builder with the default capacity.
42    pub fn new(map_dtype: MapDType, nullability: Nullability) -> Self {
43        Self::with_capacity(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY)
44    }
45
46    /// Creates a map builder with space for `capacity` map rows.
47    pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self {
48        let entries_builder = ListViewBuilder::with_capacity(
49            Arc::new(map_dtype.entries_dtype()),
50            nullability,
51            capacity.saturating_mul(2),
52            capacity,
53        );
54        let dtype = DType::Map(map_dtype.clone(), nullability);
55        Self {
56            dtype,
57            map_dtype,
58            entries_builder,
59        }
60    }
61
62    /// Appends one map scalar.
63    pub fn append_value(&mut self, value: MapScalar<'_>) -> VortexResult<()> {
64        vortex_ensure!(
65            value.dtype() == &self.dtype,
66            "MapBuilder expected map scalar with dtype {}, got {}",
67            self.dtype,
68            value.dtype()
69        );
70
71        if value.is_null() {
72            self.entries_builder.append_null();
73            return Ok(());
74        }
75
76        let entry_dtype = self.map_dtype.entries_dtype();
77        let entries = value
78            .entries()
79            .map(|(key, value)| Scalar::struct_(entry_dtype.clone(), vec![key, value]))
80            .collect();
81        let entries = Scalar::list(Arc::new(entry_dtype), entries, self.dtype.nullability());
82        self.entries_builder.append_value(entries.as_list())
83    }
84
85    /// Finishes the builder directly into a [`MapArray`].
86    pub fn finish_into_map(&mut self) -> MapArray {
87        MapArray::new(
88            self.map_dtype.clone(),
89            self.entries_builder.finish_into_listview(),
90        )
91    }
92
93    /// Appends the values of a [`Map`]-encoded `array` to this builder.
94    pub fn append_map_array(
95        &mut self,
96        array: ArrayView<'_, Map>,
97        ctx: &mut ExecutionCtx,
98    ) -> VortexResult<()> {
99        vortex_ensure!(
100            array.dtype() == self.dtype(),
101            "MapBuilder expected map array with dtype {}, got {}",
102            self.dtype(),
103            array.dtype()
104        );
105        self.entries_builder
106            .append_listview_array(array.entries().as_::<ListView>(), ctx)
107    }
108}
109
110impl<O: OffsetBuilderPType, S: OffsetBuilderPType> ArrayBuilder for MapBuilder<O, S> {
111    fn as_any(&self) -> &dyn Any {
112        self
113    }
114
115    fn as_any_mut(&mut self) -> &mut dyn Any {
116        self
117    }
118
119    fn dtype(&self) -> &DType {
120        &self.dtype
121    }
122
123    fn len(&self) -> usize {
124        self.entries_builder.len()
125    }
126
127    fn append_zeros(&mut self, n: usize) {
128        self.entries_builder.append_zeros(n);
129    }
130
131    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
132        unsafe { self.entries_builder.append_nulls_unchecked(n) };
133    }
134
135    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
136        vortex_ensure!(
137            scalar.dtype() == self.dtype(),
138            "MapBuilder expected scalar with dtype {}, got {}",
139            self.dtype(),
140            scalar.dtype()
141        );
142        self.append_value(scalar.as_map())
143    }
144
145    fn reserve_exact(&mut self, additional: usize) {
146        self.entries_builder.reserve_exact(additional);
147    }
148
149    fn finish(&mut self) -> ArrayRef {
150        self.finish_into_map().into_array()
151    }
152
153    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
154        Canonical::Map(self.finish_into_map())
155    }
156}
vortex_array/builders/map.rs

vortex_array/builders/
map.rs
