sample_arrow2/
fixed_size_list.rs

1//! Samplers for generating an arrow [`FixedSizeListArray`].
2
3use crate::{SampleLen, SetLen};
4use arrow2::{
5    array::{Array, FixedSizeListArray},
6    bitmap::Bitmap,
7    datatypes::{DataType, Field},
8};
9use sample_std::{Random, Sample};
10
11pub struct FixedSizeListWithLen<V, C, A, N> {
12    pub len: usize,
13    pub validity: V,
14    pub count: C,
15
16    pub inner: A,
17    pub inner_name: N,
18}
19
20impl<V: SetLen, C, A, N> SetLen for FixedSizeListWithLen<V, C, A, N> {
21    fn set_len(&mut self, len: usize) {
22        self.len = len;
23        self.validity.set_len(len);
24    }
25}
26
27impl<V, C, A, N> Sample for FixedSizeListWithLen<V, C, A, N>
28where
29    V: Sample<Output = Option<Bitmap>> + SetLen,
30    C: Sample<Output = usize>,
31    A: Sample<Output = Box<dyn Array>> + SetLen,
32    N: Sample<Output = String>,
33{
34    type Output = Box<dyn Array>;
35
36    fn generate(&mut self, g: &mut Random) -> Self::Output {
37        let count = self.count.generate(g);
38        self.inner.set_len(count * self.len);
39        let values = self.inner.generate(g);
40        let is_nullable = values.validity().is_some();
41        let inner_name = self.inner_name.generate(g);
42        let field = Field::new(inner_name, values.data_type().clone(), is_nullable);
43        let data_type = DataType::FixedSizeList(Box::new(field), count);
44        let validity = self.validity.generate(g);
45        FixedSizeListArray::new(data_type, values, validity).boxed()
46    }
47
48    fn shrink(&self, _: Self::Output) -> Box<dyn Iterator<Item = Self::Output>> {
49        Box::new(std::iter::empty())
50    }
51}
52
53impl<V, C, A, N> SampleLen for FixedSizeListWithLen<V, C, A, N>
54where
55    V: Sample<Output = Option<Bitmap>> + SetLen,
56    C: Sample<Output = usize>,
57    A: Sample<Output = Box<dyn Array>> + SetLen,
58    N: Sample<Output = String>,
59{
60}
sample_arrow2/fixed_size_list.rs

sample_arrow2/
fixed_size_list.rs
