sample_arrow2/
lib.rs

1use arrow2::{array::Array, bitmap::Bitmap};
2use sample_std::{Random, Sample, SampleAll, Shrunk};
3
4pub mod array;
5pub mod chunk;
6pub mod datatypes;
7pub mod fixed_size_list;
8pub mod list;
9pub mod primitive;
10pub mod struct_;
11
12pub type ArrowSampler = Box<dyn Sample<Output = Box<dyn Array>> + Send + Sync>;
13
14pub(crate) fn generate_validity<V>(
15    null: &mut Option<V>,
16    g: &mut Random,
17    len: usize,
18) -> Option<Bitmap>
19where
20    V: Sample<Output = bool>,
21{
22    null.as_mut().map(|null| {
23        Bitmap::from_trusted_len_iter(std::iter::repeat(()).take(len).map(|_| !null.generate(g)))
24    })
25}
26
27pub trait SetLen {
28    fn set_len(&mut self, len: usize);
29}
30
31#[derive(Debug, Clone)]
32pub struct AlwaysValid;
33
34impl Sample for AlwaysValid {
35    type Output = Option<Bitmap>;
36
37    fn generate(&mut self, _: &mut sample_std::Random) -> Self::Output {
38        None
39    }
40
41    fn shrink(&self, _: Self::Output) -> Shrunk<Self::Output> {
42        Box::new(std::iter::empty())
43    }
44}
45
46impl SetLen for AlwaysValid {
47    fn set_len(&mut self, _: usize) {}
48}
49
50impl<C, I, O> SetLen for sample_std::TryConvert<C, I, O>
51where
52    C: SetLen,
53{
54    fn set_len(&mut self, len: usize) {
55        self.inner.set_len(len)
56    }
57}
58
59impl<S, F, I> SampleLen for sample_std::TryConvert<S, F, I>
60where
61    S: Sample + SetLen,
62    F: Fn(S::Output) -> Box<dyn Array>,
63    I: Fn(Box<dyn Array>) -> Option<S::Output>,
64{
65}
66
67impl<C> SetLen for sample_std::SamplerChoice<C>
68where
69    C: SetLen,
70{
71    fn set_len(&mut self, len: usize) {
72        for choice in &mut self.choices {
73            choice.set_len(len);
74        }
75    }
76}
77
78impl SampleLen for sample_std::SamplerChoice<ArrowLenSampler> {}
79
80impl<S: SetLen> SetLen for SampleAll<S> {
81    fn set_len(&mut self, len: usize) {
82        for sampler in &mut self.samplers {
83            sampler.set_len(len);
84        }
85    }
86}
87
88impl Sample for Box<dyn SampleLen> {
89    type Output = Box<dyn Array>;
90
91    fn generate(&mut self, g: &mut Random) -> Self::Output {
92        self.as_mut().generate(g)
93    }
94
95    fn shrink(&self, v: Self::Output) -> Shrunk<'_, Self::Output> {
96        self.as_ref().shrink(v)
97    }
98}
99
100pub trait SampleLen: Sample<Output = Box<dyn Array>> + SetLen {}
101
102pub type ArrowLenSampler = Box<dyn SampleLen>;
103
104impl SetLen for ArrowLenSampler {
105    fn set_len(&mut self, len: usize) {
106        self.as_mut().set_len(len)
107    }
108}
109
110impl SampleLen for ArrowLenSampler {}
111
112pub struct FixedLenSampler<A> {
113    pub len: usize,
114    pub array: A,
115}
116
117impl<A> Sample for FixedLenSampler<A>
118where
119    A: Sample + SetLen,
120    A::Output: Clone,
121{
122    type Output = A::Output;
123
124    fn generate(&mut self, g: &mut Random) -> Self::Output {
125        self.array.set_len(self.len);
126        self.array.generate(g)
127    }
128
129    fn shrink(&self, array: Self::Output) -> Shrunk<Self::Output> {
130        self.array.shrink(array)
131    }
132}
sample_arrow2/lib.rs

sample_arrow2/
lib.rs
