sample_arrow2/
datatypes.rs

1//! Samplers for generating an arrow [`DataType`].
2
3use arrow2::datatypes::{DataType, Field};
4use sample_std::{sampler_choice, Always, Random, Sample, VecSampler};
5
6pub type DataTypeSampler = Box<dyn Sample<Output = DataType> + Send + Sync>;
7
8struct FieldSampler<N, V> {
9    names: N,
10    nullable: V,
11    inner: DataTypeSampler,
12}
13
14impl<N, V> Sample for FieldSampler<N, V>
15where
16    N: Sample<Output = String>,
17    V: Sample<Output = bool>,
18{
19    type Output = Field;
20
21    fn generate(&mut self, g: &mut Random) -> Self::Output {
22        Field::new(
23            self.names.generate(g),
24            self.inner.generate(g),
25            self.nullable.generate(g),
26        )
27    }
28}
29
30struct StructDataTypeSampler<S, F> {
31    size: S,
32    field: F,
33}
34
35impl<S, F> Sample for StructDataTypeSampler<S, F>
36where
37    S: Sample<Output = usize>,
38    F: Sample<Output = Field>,
39{
40    type Output = DataType;
41
42    fn generate(&mut self, g: &mut Random) -> Self::Output {
43        let size = self.size.generate(g);
44        DataType::Struct((0..size).map(|_| self.field.generate(g)).collect())
45    }
46}
47
48pub fn sample_flat() -> DataTypeSampler {
49    Box::new(sampler_choice([
50        Always(DataType::Float32),
51        Always(DataType::Float64),
52        Always(DataType::Int8),
53        Always(DataType::Int16),
54        Always(DataType::Int32),
55        Always(DataType::Int64),
56        Always(DataType::UInt8),
57        Always(DataType::UInt16),
58        Always(DataType::UInt32),
59        Always(DataType::UInt64),
60    ]))
61}
62
63pub struct ArbitraryDataType<N, V, B, F> {
64    pub names: N,
65    pub nullable: V,
66    pub struct_branch: B,
67    pub flat: F,
68}
69
70impl<N, V, B, F> ArbitraryDataType<N, V, B, F>
71where
72    N: Sample<Output = String> + Clone + Send + Sync + 'static,
73    V: Sample<Output = bool> + Clone + Send + Sync + 'static,
74    B: Sample<Output = usize> + Clone + Send + Sync + 'static,
75    F: Fn() -> DataTypeSampler,
76{
77    pub fn sample_nested<IF>(&self, inner: IF) -> DataTypeSampler
78    where
79        IF: Fn() -> DataTypeSampler,
80    {
81        let field = || FieldSampler {
82            names: self.names.clone(),
83            nullable: self.nullable.clone(),
84            inner: inner(),
85        };
86
87        Box::new(sampler_choice([
88            Box::new((self.flat)()) as DataTypeSampler,
89            Box::new(
90                VecSampler {
91                    length: self.struct_branch.clone(),
92                    el: field(),
93                }
94                .try_convert(DataType::Struct, |_| None),
95            ),
96            Box::new(field().try_convert(|f| DataType::List(Box::new(f)), |_| None)),
97        ]))
98    }
99
100    pub fn sample_depth(&self, depth: usize) -> DataTypeSampler {
101        let flats = (self.flat)();
102        if depth == 0 {
103            flats
104        } else {
105            let inner = || self.sample_depth(depth - 1);
106            Box::new(sampler_choice([self.sample_nested(inner), flats]))
107        }
108    }
109}
sample_arrow2/datatypes.rs

sample_arrow2/
datatypes.rs
