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

vortex_array/builders/
map.rs
