sample_arrow_rs/
fixed_size_list.rs

1//! Samplers for generating an arrow [`FixedSizeListArray`].
2
3use std::sync::Arc;
4
5use arrow_array::{Array, ArrayRef, FixedSizeListArray};
6use arrow_schema::Field;
7use sample_std::{Random, Sample, Shrunk};
8
9use crate::{SampleLen, SetLen};
10
11pub struct FixedSizeListWithLen<V, C, A, N> {
12    pub len: usize,
13    pub validity: V,
14    pub count: C,
15    pub inner: A,
16    pub inner_name: N,
17}
18
19impl<V: SetLen, C, A, N> SetLen for FixedSizeListWithLen<V, C, A, N> {
20    fn set_len(&mut self, len: usize) {
21        self.len = len;
22        self.validity.set_len(len);
23    }
24}
25
26impl<V, C, A, N> Sample for FixedSizeListWithLen<V, C, A, N>
27where
28    V: Sample<Output = Option<crate::Bitmap>> + SetLen,
29    C: Sample<Output = i64>, // Using i64 for size in arrow-rs
30    A: Sample<Output = ArrayRef> + SetLen,
31    N: Sample<Output = String>,
32{
33    type Output = ArrayRef;
34
35    fn generate(&mut self, g: &mut Random) -> Self::Output {
36        let count = self.count.generate(g) as i32; // Convert to i32 for arrow-rs
37        self.inner.set_len(count as usize * self.len);
38        let values = self.inner.generate(g);
39        let is_nullable = values.nulls().is_some();
40        let inner_name = self.inner_name.generate(g);
41        let field = Arc::new(Field::new(
42            inner_name,
43            values.data_type().clone(),
44            is_nullable,
45        ));
46
47        // Convert the validity bitmap to a NullBuffer if present
48        let null_buffer = self.validity.generate(g).map(|bitmap| bitmap);
49
50        // Create a FixedSizeListArray
51        Arc::new(FixedSizeListArray::new(field, count, values, null_buffer))
52    }
53
54    fn shrink(&self, _: Self::Output) -> Shrunk<Self::Output> {
55        Box::new(std::iter::empty())
56    }
57}
58
59impl<V, C, A, N> SampleLen for FixedSizeListWithLen<V, C, A, N>
60where
61    V: Sample<Output = Option<crate::Bitmap>> + SetLen,
62    C: Sample<Output = i64>,
63    A: Sample<Output = ArrayRef> + SetLen,
64    N: Sample<Output = String>,
65{
66}
sample_arrow_rs/fixed_size_list.rs

sample_arrow_rs/
fixed_size_list.rs
