sample-arrow-rs 55.2.0

Samplers for arrow-rs for use with sample-test
Documentation 
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use arrow_array::ArrayRef;
use arrow_buffer::NullBuffer;
use sample_std::{Random, Sample, SampleAll, Shrunk};

// Define Bitmap as an alias for NullBuffer to make the transition smoother
pub type Bitmap = NullBuffer;

pub mod array;
pub mod chunk;
pub mod datatypes;
pub mod fixed_size_list;
pub mod list;
pub mod primitive;
pub mod struct_;

// Use ArrayRef directly (Arc<dyn Array>)
pub type ArrowSampler = Box<dyn Sample<Output = ArrayRef> + Send + Sync>;

pub(crate) fn generate_validity<V>(
    null: &mut Option<V>,
    g: &mut Random,
    len: usize,
) -> Option<Bitmap>
where
    V: Sample<Output = bool>,
{
    null.as_mut()
        .map(|null| Bitmap::from_iter(std::iter::repeat(()).take(len).map(|_| !null.generate(g))))
}

pub trait SetLen {
    fn set_len(&mut self, len: usize);
}

#[derive(Debug, Clone)]
pub struct AlwaysValid;

impl Sample for AlwaysValid {
    type Output = Option<Bitmap>;

    fn generate(&mut self, _: &mut sample_std::Random) -> Self::Output {
        None
    }

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

impl SetLen for AlwaysValid {
    fn set_len(&mut self, _: usize) {}
}

impl<C, I, O> SetLen for sample_std::TryConvert<C, I, O>
where
    C: SetLen,
{
    fn set_len(&mut self, len: usize) {
        self.inner.set_len(len)
    }
}

impl<S, F, I> SampleLen for sample_std::TryConvert<S, F, I>
where
    S: Sample + SetLen,
    F: Fn(S::Output) -> ArrayRef,
    I: Fn(ArrayRef) -> Option<S::Output>,
{
}

impl<C> SetLen for sample_std::SamplerChoice<C>
where
    C: SetLen,
{
    fn set_len(&mut self, len: usize) {
        for choice in &mut self.choices {
            choice.set_len(len);
        }
    }
}

impl SampleLen for sample_std::SamplerChoice<ArrowLenSampler> {}

impl<S: SetLen> SetLen for SampleAll<S> {
    fn set_len(&mut self, len: usize) {
        for sampler in &mut self.samplers {
            sampler.set_len(len);
        }
    }
}

impl Sample for Box<dyn SampleLen> {
    type Output = ArrayRef;

    fn generate(&mut self, g: &mut Random) -> Self::Output {
        self.as_mut().generate(g)
    }

    fn shrink(&self, v: Self::Output) -> Shrunk<'_, Self::Output> {
        self.as_ref().shrink(v)
    }
}

pub trait SampleLen: Sample<Output = ArrayRef> + SetLen {}

pub type ArrowLenSampler = Box<dyn SampleLen>;

impl SetLen for ArrowLenSampler {
    fn set_len(&mut self, len: usize) {
        self.as_mut().set_len(len)
    }
}

impl SampleLen for ArrowLenSampler {}

pub struct FixedLenSampler<A> {
    pub len: usize,
    pub array: A,
}

impl<A> Sample for FixedLenSampler<A>
where
    A: Sample + SetLen,
    A::Output: Clone,
{
    type Output = A::Output;

    fn generate(&mut self, g: &mut Random) -> Self::Output {
        self.array.set_len(self.len);
        self.array.generate(g)
    }

    fn shrink(&self, array: Self::Output) -> Shrunk<Self::Output> {
        self.array.shrink(array)
    }
}