sample_arrow_rs/
lib.rs

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

sample_arrow_rs/
lib.rs
