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_eq;
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_eq!(
93            value.dtype(),
94            &self.dtype,
95            "MapBuilder received a map scalar with the wrong dtype"
96        );
97
98        if value.is_null() {
99            self.entries_builder.append_null();
100            return Ok(());
101        }
102
103        let entry_dtype = self.map_dtype.entries_dtype();
104        let entries = value
105            .entries()
106            .map(|(key, value)| Scalar::struct_(entry_dtype.clone(), vec![key, value]))
107            .collect();
108        let entries = Scalar::list(Arc::new(entry_dtype), entries, self.dtype.nullability());
109        self.entries_builder.append_value(entries.as_list())
110    }
111
112    /// Finishes the builder directly into a [`MapArray`].
113    pub fn finish_into_map(&mut self) -> MapArray {
114        MapArray::new(
115            self.map_dtype.clone(),
116            self.entries_builder.finish_into_listview(),
117        )
118    }
119
120    /// Appends the values of a [`Map`]-encoded `array` to this builder.
121    pub fn append_map_array(
122        &mut self,
123        array: ArrayView<'_, Map>,
124        ctx: &mut ExecutionCtx,
125    ) -> VortexResult<()> {
126        vortex_ensure_eq!(
127            array.dtype(),
128            self.dtype(),
129            "MapBuilder received a map array with the wrong dtype"
130        );
131        self.entries_builder
132            .append_listview_array(array.entries().as_::<ListView>(), ctx)
133    }
134}
135
136impl<O: OffsetBuilderPType, S: OffsetBuilderPType> ArrayBuilder for MapBuilder<O, S> {
137    fn as_any(&self) -> &dyn Any {
138        self
139    }
140
141    fn as_any_mut(&mut self) -> &mut dyn Any {
142        self
143    }
144
145    fn dtype(&self) -> &DType {
146        &self.dtype
147    }
148
149    fn len(&self) -> usize {
150        self.entries_builder.len()
151    }
152
153    fn append_zeros(&mut self, n: usize) {
154        self.entries_builder.append_zeros(n);
155    }
156
157    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
158        unsafe { self.entries_builder.append_nulls_unchecked(n) };
159    }
160
161    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
162        vortex_ensure_eq!(
163            scalar.dtype(),
164            self.dtype(),
165            "MapBuilder received a scalar with the wrong dtype"
166        );
167        self.append_value(scalar.as_map())
168    }
169
170    fn reserve_exact(&mut self, additional: usize) {
171        self.entries_builder.reserve_exact(additional);
172    }
173
174    fn finish(&mut self) -> ArrayRef {
175        self.finish_into_map().into_array()
176    }
177
178    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
179        Canonical::Map(self.finish_into_map())
180    }
181}
vortex_array/builders/map.rs

vortex_array/builders/
map.rs
