narrow/array/
mod.rs

1//! Sequences of values with known length all having the same type.
2
3use crate::{
4    Length, NonNullable, Nullability, Nullable,
5    buffer::{BufferType, VecBuffer},
6    logical::{LogicalArray, LogicalArrayType},
7    offset::{self, Offset},
8};
9use std::{collections::VecDeque, marker::PhantomData};
10
11mod boolean;
12pub use boolean::*;
13
14mod fixed_size_binary;
15pub use fixed_size_binary::*;
16
17mod fixed_size_list;
18pub use fixed_size_list::*;
19
20mod fixed_size_primitive;
21pub use fixed_size_primitive::*;
22
23mod null;
24pub use null::*;
25
26mod string;
27pub use string::*;
28
29mod r#struct;
30pub use r#struct::*;
31
32pub mod union;
33pub use union::*;
34
35mod variable_size_binary;
36pub use variable_size_binary::*;
37
38mod variable_size_list;
39pub use variable_size_list::*;
40
41/// Types that store their data in Arrow arrays.
42pub trait Array {
43    /// The items stored in this array.
44    type Item;
45}
46
47/// Types that can be stored in Arrow arrays.
48// Note: the generic `T` is required to allow impls on foreign wrappers e.g.
49// Option. (https://rust-lang.github.io/rfcs/2451-re-rebalancing-coherence.html)
50pub trait ArrayType<T: ?Sized> {
51    /// The [`Array`] type for these items.
52    ///
53    /// It is generic over:
54    /// - `Buffer`: a [`BufferType`] that is used for the array.
55    /// - `OffsetItem`: an [`Offset`] that is used for arrays with offset
56    ///   buffers.
57    /// - `UnionLayout`: a [`UnionType`] that is used for union arrays.
58    ///
59    /// When using this type constructor for arrays that have no offset buffer
60    /// [`offset::NA`] should be used to indicate that this does not apply.
61    ///
62    /// When using this type constructor to construct arrays that are not union
63    /// arrays [`union::NA`] should be used to indicate that does not apply.
64    ///
65    /// This still ends up setting the default type, but this is needed because
66    /// there are no default types for generic associated types.
67    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType>: Array;
68}
69
70/// A helper type that allows extracting the [`ArrayType::Array`] type for any `ArrayType<T> for T`
71pub type ArrayTypeOf<T, Buffer = VecBuffer, OffsetItem = offset::NA, UnionLayout = union::NA> =
72    <T as ArrayType<T>>::Array<Buffer, OffsetItem, UnionLayout>;
73
74/// A helper type that allows extracting the [`ArrayType::Array`] type for any `ArrayType<Option<T>> for T`
75pub type OptionArrayTypeOf<
76    T,
77    Buffer = VecBuffer,
78    OffsetItem = offset::NA,
79    UnionLayout = union::NA,
80> = <Option<T> as ArrayType<T>>::Array<Buffer, OffsetItem, UnionLayout>;
81
82/// A helper type that allows extracting the [`ArrayType::Array`] type for any `ArrayType<T>::Item for <T as Nullability<NULLABLE>`
83pub type NullableArrayTypeOf<
84    Nullable,
85    T,
86    Buffer,
87    OffsetItem = offset::NA,
88    UnionLayout = union::NA,
89> = <<Nullable as Nullability>::Item<T> as ArrayType<T>>::Array<Buffer, OffsetItem, UnionLayout>;
90
91impl<T: ArrayType<U> + ?Sized, U> ArrayType<U> for &T {
92    type Array<Buffer: BufferType, OfsetItem: Offset, UnionLayout: UnionType> =
93        <T as ArrayType<U>>::Array<Buffer, OfsetItem, UnionLayout>;
94}
95
96/// ?
97pub trait Layout {
98    /// a
99    type Buffer: BufferType;
100    /// b
101    type Offset: Offset;
102    /// c
103    type Union: UnionType;
104}
105
106/// Implement [`ArrayType`] for `ty` using `array`.
107macro_rules! impl_array_type {
108    ($ty:ty, $array:ty) => {
109        impl ArrayType<$ty> for $ty {
110            type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> = $array;
111        }
112    };
113    ($ty:ty, $array:ty, $inner:ty) => {
114        impl ArrayType<$inner> for $ty {
115            type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> = $array;
116        }
117    };
118}
119
120impl_array_type!(bool, BooleanArray<NonNullable, Buffer>);
121impl_array_type!(Option<bool>, BooleanArray<Nullable, Buffer>, bool);
122
123impl_array_type!(u8, FixedSizePrimitiveArray<u8, NonNullable, Buffer>);
124impl_array_type!(Option<u8>, FixedSizePrimitiveArray<u8, Nullable, Buffer>, u8);
125impl_array_type!(i8, FixedSizePrimitiveArray<i8, NonNullable, Buffer>);
126impl_array_type!(Option<i8>, FixedSizePrimitiveArray<i8, Nullable, Buffer>, i8);
127impl_array_type!(u16, FixedSizePrimitiveArray<u16, NonNullable, Buffer>);
128impl_array_type!(Option<u16>, FixedSizePrimitiveArray<u16, Nullable, Buffer>, u16);
129impl_array_type!(i16, FixedSizePrimitiveArray<i16, NonNullable, Buffer>);
130impl_array_type!(Option<i16>, FixedSizePrimitiveArray<i16, Nullable, Buffer>, i16);
131impl_array_type!(u32, FixedSizePrimitiveArray<u32, NonNullable, Buffer>);
132impl_array_type!(Option<u32>, FixedSizePrimitiveArray<u32, Nullable, Buffer>, u32);
133impl_array_type!(i32, FixedSizePrimitiveArray<i32, NonNullable, Buffer>);
134impl_array_type!(Option<i32>, FixedSizePrimitiveArray<i32, Nullable, Buffer>, i32);
135impl_array_type!(u64, FixedSizePrimitiveArray<u64, NonNullable, Buffer>);
136impl_array_type!(Option<u64>, FixedSizePrimitiveArray<u64, Nullable, Buffer>, u64);
137impl_array_type!(i64, FixedSizePrimitiveArray<i64, NonNullable, Buffer>);
138impl_array_type!(Option<i64>, FixedSizePrimitiveArray<i64, Nullable, Buffer>, i64);
139#[cfg(not(feature = "arrow-rs"))]
140impl_array_type!(u128, FixedSizePrimitiveArray<u128, NonNullable, Buffer>);
141#[cfg(not(feature = "arrow-rs"))]
142impl_array_type!(Option<u128>, FixedSizePrimitiveArray<u128, Nullable, Buffer>, u128);
143impl_array_type!(i128, FixedSizePrimitiveArray<i128, NonNullable, Buffer>);
144impl_array_type!(Option<i128>, FixedSizePrimitiveArray<i128, Nullable, Buffer>, i128);
145
146#[cfg(not(feature = "arrow-rs"))]
147impl_array_type!(usize, FixedSizePrimitiveArray<usize, NonNullable, Buffer>);
148#[cfg(not(feature = "arrow-rs"))]
149impl_array_type!(Option<usize>, FixedSizePrimitiveArray<usize, Nullable, Buffer>, usize);
150#[cfg(not(feature = "arrow-rs"))]
151impl_array_type!(isize, FixedSizePrimitiveArray<isize, NonNullable, Buffer>);
152#[cfg(not(feature = "arrow-rs"))]
153impl_array_type!(Option<isize>, FixedSizePrimitiveArray<isize, Nullable, Buffer>, isize);
154
155impl_array_type!(f32, FixedSizePrimitiveArray<f32, NonNullable, Buffer>);
156impl_array_type!(Option<f32>, FixedSizePrimitiveArray<f32, Nullable, Buffer>, f32);
157impl_array_type!(f64, FixedSizePrimitiveArray<f64, NonNullable, Buffer>);
158impl_array_type!(Option<f64>, FixedSizePrimitiveArray<f64, Nullable, Buffer>, f64);
159
160impl_array_type!((), NullArray<(), NonNullable, Buffer>);
161impl_array_type!(Option<()>, NullArray<(), Nullable, Buffer>, ());
162
163/// An byte array wrapper that maps to [`FixedSizeBinaryArray`] via its
164/// [`ArrayType`] implementation. Used for example to map `Uuid` to
165/// a [`FixedSizeBinaryArray`] instead of a [`FixedSizeListArray`].
166#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
167pub struct FixedSizeBinary<const N: usize>([u8; N]);
168
169impl<const N: usize> Default for FixedSizeBinary<N> {
170    fn default() -> Self {
171        Self([u8::default(); N])
172    }
173}
174
175impl<const N: usize> ArrayType<FixedSizeBinary<N>> for FixedSizeBinary<N> {
176    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
177        FixedSizeBinaryArray<N, NonNullable, Buffer>;
178}
179
180impl<const N: usize> ArrayType<FixedSizeBinary<N>> for Option<FixedSizeBinary<N>> {
181    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
182        FixedSizeBinaryArray<N, Nullable, Buffer>;
183}
184
185impl<const N: usize> From<&[u8; N]> for FixedSizeBinary<N> {
186    fn from(value: &[u8; N]) -> Self {
187        Self(*value)
188    }
189}
190
191impl<const N: usize> From<[u8; N]> for FixedSizeBinary<N> {
192    fn from(value: [u8; N]) -> Self {
193        Self(value)
194    }
195}
196
197impl<const N: usize> From<FixedSizeBinary<N>> for [u8; N] {
198    fn from(value: FixedSizeBinary<N>) -> Self {
199        value.0
200    }
201}
202
203/// An byte vector wrapper that maps to [`VariableSizeBinaryArray`] via its
204/// [`ArrayType`] implementation. Used for example to map `Vec<u8>` to
205/// a [`VariableSizeBinaryArray`] instead of a [`VariableSizeListArray`].
206#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
207pub struct VariableSizeBinary(Vec<u8>);
208
209impl ArrayType<VariableSizeBinary> for VariableSizeBinary {
210    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
211        LogicalArray<Self, NonNullable, Buffer, OffsetItem, UnionLayout>;
212}
213
214impl ArrayType<VariableSizeBinary> for Option<VariableSizeBinary> {
215    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
216        LogicalArray<VariableSizeBinary, Nullable, Buffer, OffsetItem, UnionLayout>;
217}
218
219#[cfg(feature = "arrow-rs")]
220impl crate::arrow::LogicalArrayType<VariableSizeBinary> for VariableSizeBinary {
221    type ExtensionType = crate::arrow::NoExtensionType;
222}
223
224impl LogicalArrayType<VariableSizeBinary> for VariableSizeBinary {
225    type ArrayType = Vec<u8>;
226
227    fn from_array_type(item: Self::ArrayType) -> Self {
228        item.into()
229    }
230
231    fn into_array_type(self) -> Self::ArrayType {
232        self.into()
233    }
234}
235
236impl From<Vec<u8>> for VariableSizeBinary {
237    fn from(value: Vec<u8>) -> Self {
238        Self(value)
239    }
240}
241
242impl From<VariableSizeBinary> for Vec<u8> {
243    fn from(value: VariableSizeBinary) -> Self {
244        value.0
245    }
246}
247
248impl Length for VariableSizeBinary {
249    fn len(&self) -> usize {
250        self.0.len()
251    }
252}
253
254impl IntoIterator for VariableSizeBinary {
255    type Item = u8;
256    type IntoIter = std::vec::IntoIter<u8>;
257
258    fn into_iter(self) -> Self::IntoIter {
259        self.0.into_iter()
260    }
261}
262
263impl<T: ArrayType<T>, const N: usize> ArrayType<[T; N]> for [T; N] {
264    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
265        FixedSizeListArray<N, ArrayTypeOf<T, Buffer, OffsetItem, UnionLayout>, NonNullable, Buffer>;
266}
267impl<T: ArrayType<T>, const N: usize> ArrayType<[T; N]> for Option<[T; N]> {
268    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
269        FixedSizeListArray<N, ArrayTypeOf<T, Buffer, OffsetItem, UnionLayout>, Nullable, Buffer>;
270}
271impl<T, const N: usize> ArrayType<[Option<T>; N]> for [Option<T>; N]
272where
273    Option<T>: ArrayType<T>,
274{
275    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> = FixedSizeListArray<
276        N,
277        OptionArrayTypeOf<T, Buffer, OffsetItem, UnionLayout>,
278        NonNullable,
279        Buffer,
280    >;
281}
282impl<T, const N: usize> ArrayType<[Option<T>; N]> for Option<[Option<T>; N]>
283where
284    Option<T>: ArrayType<T>,
285{
286    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> = FixedSizeListArray<
287        N,
288        OptionArrayTypeOf<T, Buffer, OffsetItem, UnionLayout>,
289        Nullable,
290        Buffer,
291    >;
292}
293impl ArrayType<str> for str {
294    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
295        StringArray<NonNullable, OffsetItem, Buffer>;
296}
297impl<'a> ArrayType<&'a str> for &'a str {
298    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
299        StringArray<NonNullable, OffsetItem, Buffer>;
300}
301impl<'a> ArrayType<&'a str> for Option<&'a str> {
302    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
303        StringArray<Nullable, OffsetItem, Buffer>;
304}
305impl ArrayType<String> for String {
306    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
307        StringArray<NonNullable, OffsetItem, Buffer>;
308}
309impl ArrayType<String> for Option<String> {
310    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
311        StringArray<Nullable, OffsetItem, Buffer>;
312}
313
314impl<'a, T: ArrayType<T>> ArrayType<&'a [T]> for &'a [T] {
315    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
316        VariableSizeListArray<ArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
317}
318impl<'a, T: ArrayType<T>> ArrayType<&'a [T]> for Option<&'a [T]> {
319    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
320        VariableSizeListArray<ArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
321}
322impl<'a, T> ArrayType<&'a [Option<T>]> for &'a [Option<T>]
323where
324    Option<T>: ArrayType<T>,
325{
326    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
327        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
328}
329impl<'a, T> ArrayType<&'a [Option<T>]> for Option<&'a [Option<T>]>
330where
331    Option<T>: ArrayType<T>,
332{
333    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
334        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
335}
336impl<T: ArrayType<T>> ArrayType<Vec<T>> for Vec<T> {
337    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
338        VariableSizeListArray<ArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
339}
340impl<T: ArrayType<T>> ArrayType<Vec<T>> for Option<Vec<T>> {
341    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
342        VariableSizeListArray<ArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
343}
344impl<T> ArrayType<Vec<Option<T>>> for Vec<Option<T>>
345where
346    Option<T>: ArrayType<T>,
347{
348    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
349        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
350}
351impl<T> ArrayType<Vec<Option<T>>> for Option<Vec<Option<T>>>
352where
353    Option<T>: ArrayType<T>,
354{
355    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
356        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
357}
358
359impl<T: ArrayType<T>> ArrayType<VecDeque<T>> for VecDeque<T> {
360    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
361        VariableSizeListArray<ArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
362}
363impl<T: ArrayType<T>> ArrayType<VecDeque<T>> for Option<VecDeque<T>> {
364    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
365        VariableSizeListArray<ArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
366}
367impl<T> ArrayType<VecDeque<Option<T>>> for VecDeque<Option<T>>
368where
369    Option<T>: ArrayType<T>,
370{
371    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
372        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, NonNullable, OffsetItem, Buffer>;
373}
374impl<T> ArrayType<VecDeque<Option<T>>> for Option<VecDeque<Option<T>>>
375where
376    Option<T>: ArrayType<T>,
377{
378    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> =
379        VariableSizeListArray<OptionArrayTypeOf<T, Buffer>, Nullable, OffsetItem, Buffer>;
380}
381
382impl<T> ArrayType<PhantomData<T>> for PhantomData<T> {
383    type Array<Buffer: BufferType, OffsetItem: Offset, UnionLayout: UnionType> = NullArray;
384}