polars_arrow/array/growable/
mod.rs

1//! Contains the trait [`Growable`] and corresponding concreate implementations, one per concrete array,
2//! that offer the ability to create a new [`Array`] out of slices of existing [`Array`]s.
3
4use crate::array::*;
5use crate::datatypes::*;
6
7mod binary;
8pub use binary::GrowableBinary;
9mod boolean;
10pub use boolean::GrowableBoolean;
11mod fixed_binary;
12pub use fixed_binary::GrowableFixedSizeBinary;
13mod null;
14pub use null::GrowableNull;
15mod primitive;
16pub use primitive::GrowablePrimitive;
17mod list;
18pub use list::GrowableList;
19mod structure;
20pub use structure::GrowableStruct;
21mod fixed_size_list;
22pub use fixed_size_list::GrowableFixedSizeList;
23mod utf8;
24pub use utf8::GrowableUtf8;
25mod dictionary;
26pub use dictionary::GrowableDictionary;
27
28mod binview;
29pub use binview::GrowableBinaryViewArray;
30
31mod utils;
32
33/// Describes a struct that can be extended from slices of other pre-existing [`Array`]s.
34/// This is used in operations where a new array is built out of other arrays, such
35/// as filter and concatenation.
36pub trait Growable<'a> {
37    /// Extends this [`Growable`] with elements from the bounded [`Array`] at index `index` from
38    /// a slice starting at `start` and length `len`.
39    ///
40    /// # Safety
41    /// Doesn't do any bound checks.
42    unsafe fn extend(&mut self, index: usize, start: usize, len: usize);
43
44    /// Same as extend, except it repeats the extension `copies` times.
45    ///
46    /// # Safety
47    /// Doesn't do any bound checks.
48    unsafe fn extend_copies(&mut self, index: usize, start: usize, len: usize, copies: usize) {
49        for _ in 0..copies {
50            self.extend(index, start, len)
51        }
52    }
53
54    /// Extends this [`Growable`] with null elements, disregarding the bound arrays
55    ///
56    /// # Safety
57    /// Doesn't do any bound checks
58    fn extend_validity(&mut self, additional: usize);
59
60    /// The current length of the [`Growable`].
61    fn len(&self) -> usize;
62
63    /// Converts this [`Growable`] to an [`Arc<dyn Array>`], thereby finishing the mutation.
64    /// Self will be empty after such operation.
65    fn as_arc(&mut self) -> Arc<dyn Array> {
66        self.as_box().into()
67    }
68
69    /// Converts this [`Growable`] to an [`Box<dyn Array>`], thereby finishing the mutation.
70    /// Self will be empty after such operation
71    fn as_box(&mut self) -> Box<dyn Array>;
72}
73
74macro_rules! dyn_growable {
75    ($ty:ty, $arrays:expr, $use_validity:expr, $capacity:expr) => {{
76        let arrays = $arrays
77            .iter()
78            .map(|array| array.as_any().downcast_ref().unwrap())
79            .collect::<Vec<_>>();
80        Box::new(<$ty>::new(arrays, $use_validity, $capacity))
81    }};
82}
83
84/// Creates a new [`Growable`] from an arbitrary number of [`Array`]s.
85/// # Panics
86/// This function panics iff
87/// * the arrays do not have the same [`ArrowDataType`].
88/// * `arrays.is_empty()`.
89pub fn make_growable<'a>(
90    arrays: &[&'a dyn Array],
91    use_validity: bool,
92    capacity: usize,
93) -> Box<dyn Growable<'a> + 'a> {
94    assert!(!arrays.is_empty());
95    let dtype = arrays[0].dtype();
96
97    use PhysicalType::*;
98    match dtype.to_physical_type() {
99        Null => Box::new(null::GrowableNull::new(dtype.clone())),
100        Boolean => dyn_growable!(boolean::GrowableBoolean, arrays, use_validity, capacity),
101        Primitive(primitive) => with_match_primitive_type_full!(primitive, |$T| {
102            dyn_growable!(primitive::GrowablePrimitive::<$T>, arrays, use_validity, capacity)
103        }),
104        Binary => dyn_growable!(
105            binary::GrowableBinary::<i32>,
106            arrays,
107            use_validity,
108            capacity
109        ),
110        LargeBinary => dyn_growable!(
111            binary::GrowableBinary::<i64>,
112            arrays,
113            use_validity,
114            capacity
115        ),
116        FixedSizeBinary => dyn_growable!(
117            fixed_binary::GrowableFixedSizeBinary,
118            arrays,
119            use_validity,
120            capacity
121        ),
122        LargeList => dyn_growable!(list::GrowableList::<i64>, arrays, use_validity, capacity),
123        Struct => dyn_growable!(structure::GrowableStruct, arrays, use_validity, capacity),
124        FixedSizeList => dyn_growable!(
125            fixed_size_list::GrowableFixedSizeList,
126            arrays,
127            use_validity,
128            capacity
129        ),
130        BinaryView => {
131            dyn_growable!(
132                binview::GrowableBinaryViewArray::<[u8]>,
133                arrays,
134                use_validity,
135                capacity
136            )
137        },
138        Utf8View => {
139            dyn_growable!(
140                binview::GrowableBinaryViewArray::<str>,
141                arrays,
142                use_validity,
143                capacity
144            )
145        },
146        Dictionary(key_type) => {
147            match_integer_type!(key_type, |$T| {
148                let arrays = arrays
149                    .iter()
150                    .map(|array| {
151                        array
152                            .as_any()
153                            .downcast_ref::<DictionaryArray<$T>>()
154                            .unwrap()
155                    })
156                    .collect::<Vec<_>>();
157                Box::new(dictionary::GrowableDictionary::<$T>::new(
158                    &arrays,
159                    use_validity,
160                    capacity,
161                ))
162            })
163        },
164        Union | Map | Utf8 | LargeUtf8 | List => unimplemented!(),
165    }
166}