1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

//! Samplers for generating an arrow [`FixedSizeListArray`].

use crate::{SampleLen, SetLen};
use arrow2::{
    array::{Array, FixedSizeListArray},
    bitmap::Bitmap,
    datatypes::{DataType, Field},
};
use sample_std::{Random, Sample};

pub struct FixedSizeListWithLen<V, C, A, N> {
    pub len: usize,
    pub validity: V,
    pub count: C,

    pub inner: A,
    pub inner_name: N,
}

impl<V: SetLen, C, A, N> SetLen for FixedSizeListWithLen<V, C, A, N> {
    fn set_len(&mut self, len: usize) {
        self.len = len;
        self.validity.set_len(len);
    }
}

impl<V, C, A, N> Sample for FixedSizeListWithLen<V, C, A, N>
where
    V: Sample<Output = Option<Bitmap>> + SetLen,
    C: Sample<Output = usize>,
    A: Sample<Output = Box<dyn Array>> + SetLen,
    N: Sample<Output = String>,
{
    type Output = Box<dyn Array>;

    fn generate(&mut self, g: &mut Random) -> Self::Output {
        let count = self.count.generate(g);
        self.inner.set_len(count * self.len);
        let values = self.inner.generate(g);
        let is_nullable = values.validity().is_some();
        let inner_name = self.inner_name.generate(g);
        let field = Field::new(inner_name, values.data_type().clone(), is_nullable);
        let data_type = DataType::FixedSizeList(Box::new(field), count);
        let validity = self.validity.generate(g);
        FixedSizeListArray::new(data_type, values, validity).boxed()
    }

    fn shrink(&self, _: Self::Output) -> Box<dyn Iterator<Item = Self::Output>> {
        Box::new(std::iter::empty())
    }
}

impl<V, C, A, N> SampleLen for FixedSizeListWithLen<V, C, A, N>
where
    V: Sample<Output = Option<Bitmap>> + SetLen,
    C: Sample<Output = usize>,
    A: Sample<Output = Box<dyn Array>> + SetLen,
    N: Sample<Output = String>,
{
}