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;
9use vortex_mask::Mask;
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    pub fn new(map_dtype: MapDType, nullability: Nullability) -> Self {
44        Self::with_capacity(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY)
45    }
46
47    /// Creates a map builder with space for `capacity` map rows.
48    pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self {
49        let entries_builder = ListViewBuilder::with_capacity(
50            Arc::new(map_dtype.entries_dtype()),
51            nullability,
52            capacity.saturating_mul(2),
53            capacity,
54        );
55        let dtype = DType::Map(map_dtype.clone(), nullability);
56        Self {
57            dtype,
58            map_dtype,
59            entries_builder,
60        }
61    }
62
63    /// Appends one map scalar.
64    pub fn append_value(&mut self, value: MapScalar<'_>) -> VortexResult<()> {
65        vortex_ensure!(
66            value.dtype() == &self.dtype,
67            "MapBuilder expected map scalar with dtype {}, got {}",
68            self.dtype,
69            value.dtype()
70        );
71
72        if value.is_null() {
73            self.entries_builder.append_null();
74            return Ok(());
75        }
76
77        let entry_dtype = self.map_dtype.entries_dtype();
78        let entries = value
79            .entries()
80            .map(|(key, value)| Scalar::struct_(entry_dtype.clone(), vec![key, value]))
81            .collect();
82        let entries = Scalar::list(Arc::new(entry_dtype), entries, self.dtype.nullability());
83        self.entries_builder.append_value(entries.as_list())
84    }
85
86    /// Finishes the builder directly into a [`MapArray`].
87    pub fn finish_into_map(&mut self) -> MapArray {
88        MapArray::new(
89            self.map_dtype.clone(),
90            self.entries_builder.finish_into_listview(),
91        )
92    }
93
94    /// Appends the values of a [`Map`]-encoded `array` to this builder.
95    pub fn append_map_array(
96        &mut self,
97        array: ArrayView<'_, Map>,
98        ctx: &mut ExecutionCtx,
99    ) -> VortexResult<()> {
100        vortex_ensure!(
101            array.dtype() == self.dtype(),
102            "MapBuilder expected map array with dtype {}, got {}",
103            self.dtype(),
104            array.dtype()
105        );
106        self.entries_builder
107            .append_listview_array(array.entries().as_::<ListView>(), ctx)
108    }
109}
110
111impl<O: OffsetBuilderPType, S: OffsetBuilderPType> ArrayBuilder for MapBuilder<O, S> {
112    fn as_any(&self) -> &dyn Any {
113        self
114    }
115
116    fn as_any_mut(&mut self) -> &mut dyn Any {
117        self
118    }
119
120    fn dtype(&self) -> &DType {
121        &self.dtype
122    }
123
124    fn len(&self) -> usize {
125        self.entries_builder.len()
126    }
127
128    fn append_zeros(&mut self, n: usize) {
129        self.entries_builder.append_zeros(n);
130    }
131
132    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
133        unsafe { self.entries_builder.append_nulls_unchecked(n) };
134    }
135
136    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
137        vortex_ensure!(
138            scalar.dtype() == self.dtype(),
139            "MapBuilder expected scalar with dtype {}, got {}",
140            self.dtype(),
141            scalar.dtype()
142        );
143        self.append_value(scalar.as_map())
144    }
145
146    fn reserve_exact(&mut self, additional: usize) {
147        self.entries_builder.reserve_exact(additional);
148    }
149
150    unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
151        unsafe { self.entries_builder.set_validity_unchecked(validity) };
152    }
153
154    fn finish(&mut self) -> ArrayRef {
155        self.finish_into_map().into_array()
156    }
157
158    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
159        Canonical::Map(self.finish_into_map())
160    }
161}
vortex_array/builders/map.rs

vortex_array/builders/
map.rs
