sample_arrow_rs/
struct_.rs

1//! Samplers for generating an arrow [`StructArray`].
2
3use arrow_array::{Array, ArrayRef, StructArray};
4use arrow_schema::DataType;
5use sample_std::{Random, Sample, Shrunk};
6use std::sync::Arc;
7
8use crate::{generate_validity, ArrowSampler};
9
10pub struct StructSampler<V> {
11    pub data_type: DataType,
12    pub null: Option<V>,
13    pub values: Vec<ArrowSampler>,
14}
15
16impl<V> Sample for StructSampler<V>
17where
18    V: Sample<Output = bool>,
19{
20    type Output = ArrayRef;
21
22    fn generate(&mut self, g: &mut Random) -> Self::Output {
23        let arrays: Vec<ArrayRef> = self.values.iter_mut().map(|sa| sa.generate(g)).collect();
24        let validity = generate_validity(&mut self.null, g, arrays[0].len());
25
26        // Extract fields from the data_type
27        let fields = if let DataType::Struct(fields) = &self.data_type {
28            fields.clone()
29        } else {
30            panic!("Expected Struct data type")
31        };
32
33        // Create the struct array as ArrayRef
34        Arc::new(StructArray::new(fields, arrays, validity))
35    }
36
37    fn shrink(&self, _v: Self::Output) -> Shrunk<Self::Output> {
38        Box::new(std::iter::empty())
39    }
40}
sample_arrow_rs/struct_.rs

sample_arrow_rs/
struct_.rs
