sample_arrow_rs/
chunk.rs

1//! Chained samplers for generating arbitrary `RecordBatch` arrow record batches.
2
3use std::ops::Range;
4use std::sync::Arc;
5
6use arrow_array::{ArrayRef, RecordBatch};
7use arrow_schema::{DataType, Field, Schema};
8use sample_std::{sample_all, Chained, Random, Sample, Shrunk, VecSampler};
9
10use crate::{array::ArbitraryArray, datatypes::DataTypeSampler};
11
12// In arrow-rs, we use RecordBatch instead of Chunk
13pub type ChainedChunk = Chained<(Vec<DataType>, usize), RecordBatch>;
14pub type ChainedMultiChunk = Chained<(Vec<DataType>, Vec<usize>), Vec<RecordBatch>>;
15
16pub struct ArbitraryChunk<N, V> {
17    pub chunk_len: Range<usize>,
18    pub array_count: Range<usize>,
19    pub data_type: DataTypeSampler,
20    pub array: ArbitraryArray<N, V>,
21}
22
23struct RecordBatchSampler {
24    arrays_sampler: Box<dyn Sample<Output = Vec<ArrayRef>> + Send + Sync>,
25    schema: Arc<Schema>,
26}
27
28impl Sample for RecordBatchSampler {
29    type Output = RecordBatch;
30
31    fn generate(&mut self, g: &mut Random) -> Self::Output {
32        // Generate the arrays
33        let arrays = self.arrays_sampler.generate(g);
34
35        // Create the record batch
36        RecordBatch::try_new(self.schema.clone(), arrays)
37            .unwrap_or_else(|_| panic!("Failed to create record batch"))
38    }
39
40    fn shrink(&self, _: Self::Output) -> Shrunk<Self::Output> {
41        Box::new(std::iter::empty())
42    }
43}
44
45impl<N, V> ArbitraryChunk<N, V>
46where
47    N: Sample<Output = String> + Send + Sync + Clone + 'static,
48    V: Sample<Output = bool> + Send + Sync + Clone + 'static,
49{
50    pub fn sample_one(self) -> Box<dyn Sample<Output = ChainedChunk> + Send + Sync> {
51        Box::new(
52            VecSampler {
53                length: self.array_count,
54                el: self.data_type,
55            }
56            .zip(self.chunk_len)
57            .chain_resample(move |seed| Self::from_seed(&self.array, seed), 100),
58        )
59    }
60
61    pub fn sample_many(
62        self,
63        chunk_count: Range<usize>,
64    ) -> Box<dyn Sample<Output = ChainedMultiChunk> + Send + Sync> {
65        Box::new(
66            VecSampler {
67                length: self.array_count,
68                el: self.data_type,
69            }
70            .zip(VecSampler {
71                length: chunk_count,
72                el: self.chunk_len,
73            })
74            .chain_resample(
75                move |(dts, lens)| {
76                    sample_all(
77                        lens.into_iter()
78                            .map(|len| Self::from_seed(&self.array, (dts.clone(), len)))
79                            .collect(),
80                    )
81                },
82                100,
83            ),
84        )
85    }
86
87    pub fn from_seed(
88        array: &ArbitraryArray<N, V>,
89        seed: (Vec<DataType>, usize),
90    ) -> Box<dyn Sample<Output = RecordBatch> + Send + Sync> {
91        let (dts, len) = seed;
92
93        // Create field names for the schema
94        let field_names: Vec<String> = dts
95            .iter()
96            .enumerate()
97            .map(|(i, _)| format!("field_{}", i))
98            .collect();
99
100        // Create fields for the schema
101        let fields: Vec<Field> = dts
102            .iter()
103            .zip(field_names.iter())
104            .map(|(dt, name)| Field::new(name, dt.clone(), true))
105            .collect();
106
107        // Create the schema
108        let schema = Arc::new(Schema::new(fields));
109
110        // Generate arrays from data types
111        let arrays_sampler = sample_all(
112            dts.into_iter()
113                .map(|data_type| array.with_len(len).sampler_from_data_type(&data_type))
114                .collect(),
115        );
116
117        // Create the record batch sample generator
118        Box::new(RecordBatchSampler {
119            arrays_sampler: Box::new(arrays_sampler),
120            schema,
121        })
122    }
123}
sample_arrow_rs/chunk.rs

sample_arrow_rs/
chunk.rs
