sample_arrow_rs/
array.rs

1//! Chained samplers for generating arbitrary `Arc<dyn Array>` arrow arrays.
2
3use std::ops::Range;
4
5use arrow_array::{Array, ArrayRef, FixedSizeListArray, ListArray};
6use arrow_schema::DataType;
7use sample_std::{Always, Chained, Sample};
8
9use crate::{
10    datatypes::DataTypeSampler,
11    fixed_size_list::FixedSizeListWithLen,
12    list::{ListSampler, ListWithLen},
13    primitive::{
14        f32_array, f32_sampler, f64_array, f64_sampler, i16_array, i16_sampler, i32_array,
15        i32_sampler, i64_array, i64_sampler, i8_array, i8_sampler, u16_array, u16_sampler,
16        u32_array, u32_sampler, u64_array, u64_sampler, u8_array, u8_sampler,
17    },
18    struct_::StructSampler,
19    AlwaysValid, ArrowLenSampler, SetLen,
20};
21
22pub fn sampler_from_example(array: &dyn Array) -> ArrowLenSampler {
23    match array.data_type() {
24        DataType::Float32 => f32_sampler(AlwaysValid),
25        DataType::Float64 => f64_sampler(AlwaysValid),
26        DataType::Int8 => i8_sampler(AlwaysValid),
27        DataType::Int16 => i16_sampler(AlwaysValid),
28        DataType::Int32 => i32_sampler(AlwaysValid),
29        DataType::Int64 => i64_sampler(AlwaysValid),
30        DataType::UInt8 => u8_sampler(AlwaysValid),
31        DataType::UInt16 => u16_sampler(AlwaysValid),
32        DataType::UInt32 => u32_sampler(AlwaysValid),
33        DataType::UInt64 => u64_sampler(AlwaysValid),
34        DataType::List(_) => {
35            let list = array.as_any().downcast_ref::<ListArray>().unwrap();
36            // In arrow-rs we access offsets differently
37            let _offsets_buffer = list.offsets();
38            // Get the raw values as a slice - arrow-rs uses different methods
39            // than arrow2, we'll try to access offsets directly
40            let offsets = list.value_offsets();
41            let lengths: Vec<_> = offsets.windows(2).map(|w| w[1] - w[0]).collect();
42            let min = lengths.iter().min().copied().unwrap_or(0) as i32;
43            let max = lengths.iter().max().copied().unwrap_or(0) as i32 + 1;
44
45            // In arrow-rs, field information comes from the data_type
46            let field_name = if let DataType::List(field) = array.data_type() {
47                field.name().clone()
48            } else {
49                "item".to_string() // Fallback
50            };
51
52            Box::new(ListWithLen {
53                len: array.len(),
54                validity: AlwaysValid,
55                count: min..max,
56                inner_name: Always(field_name),
57                inner: sampler_from_example(list.values()),
58            })
59        }
60        DataType::FixedSizeList(field, size) => {
61            let list = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
62            Box::new(FixedSizeListWithLen {
63                len: array.len(),
64                validity: AlwaysValid,
65                count: Always(*size as i64),
66                inner_name: Always(field.name().clone()),
67                inner: sampler_from_example(list.values()),
68            })
69        }
70        dt => panic!("not implemented: {:?}", dt),
71    }
72}
73
74pub struct FromDataType<V, B> {
75    pub validity: V,
76    pub branch: B,
77}
78
79impl<V, B> FromDataType<V, B>
80where
81    V: Sample<Output = Option<crate::Bitmap>> + SetLen + Clone + Send + Sync + 'static,
82    B: Sample<Output = i32> + Clone + Send + Sync + 'static,
83{
84    pub fn from_data_type(&self, data_type: &DataType) -> ArrowLenSampler {
85        match data_type {
86            DataType::Float32 => f32_sampler(self.validity.clone()),
87            DataType::Float64 => f64_sampler(self.validity.clone()),
88            DataType::Int8 => i8_sampler(self.validity.clone()),
89            DataType::Int16 => i16_sampler(self.validity.clone()),
90            DataType::Int32 => i32_sampler(self.validity.clone()),
91            DataType::Int64 => i64_sampler(self.validity.clone()),
92            DataType::UInt8 => u8_sampler(self.validity.clone()),
93            DataType::UInt16 => u16_sampler(self.validity.clone()),
94            DataType::UInt32 => u32_sampler(self.validity.clone()),
95            DataType::UInt64 => u64_sampler(self.validity.clone()),
96            DataType::List(field) => Box::new(ListWithLen {
97                len: 0,
98                validity: self.validity.clone(),
99                count: self.branch.clone(),
100                inner_name: Always(field.name().clone()),
101                inner: self.from_data_type(field.data_type()),
102            }),
103            DataType::FixedSizeList(field, size) => Box::new(FixedSizeListWithLen {
104                len: 0,
105                validity: self.validity.clone(),
106                count: Always(*size as i64),
107                inner_name: Always(field.name().clone()),
108                inner: self.from_data_type(field.data_type()),
109            }),
110            dt => panic!("not implemented: {:?}", dt),
111        }
112    }
113}
114
115pub type ArraySampler = Box<dyn Sample<Output = ArrayRef> + Send + Sync>;
116
117pub type ChainedArraySampler = Box<dyn Sample<Output = Chained<DataType, ArrayRef>> + Send + Sync>;
118
119#[derive(Clone, Debug)]
120pub struct ArbitraryArray<N, V> {
121    pub names: N,
122    pub branch: Range<usize>,
123    pub len: Range<usize>,
124    pub null: V,
125    pub is_nullable: bool,
126}
127
128impl<N, V> ArbitraryArray<N, V>
129where
130    N: Sample<Output = String> + Send + Sync + Clone + 'static,
131    V: Sample<Output = bool> + Send + Sync + Clone + 'static,
132{
133    pub fn with_len(&self, len: usize) -> Self {
134        Self {
135            len: len..(len + 1),
136            ..self.clone()
137        }
138    }
139
140    pub fn arbitrary_array(self, data_type_sampler: DataTypeSampler) -> ChainedArraySampler {
141        Box::new(data_type_sampler.chain_resample(
142            move |data_type| self.sampler_from_data_type(&data_type),
143            100,
144        ))
145    }
146
147    pub fn sampler_from_data_type(&self, data_type: &DataType) -> ArraySampler {
148        let current_null = if self.is_nullable {
149            Some(self.null.clone())
150        } else {
151            None
152        };
153        let len = self.len.clone();
154
155        match data_type {
156            DataType::Float32 => f32_array(len.clone(), current_null),
157            DataType::Float64 => f64_array(len.clone(), current_null),
158            DataType::Int8 => i8_array(len.clone(), current_null),
159            DataType::Int16 => i16_array(len.clone(), current_null),
160            DataType::Int32 => i32_array(len.clone(), current_null),
161            DataType::Int64 => i64_array(len.clone(), current_null),
162            DataType::UInt8 => u8_array(len.clone(), current_null),
163            DataType::UInt16 => u16_array(len.clone(), current_null),
164            DataType::UInt32 => u32_array(len.clone(), current_null),
165            DataType::UInt64 => u64_array(len.clone(), current_null),
166            DataType::Struct(fields) => Box::new(StructSampler {
167                data_type: data_type.clone(),
168                null: current_null,
169                values: fields
170                    .iter()
171                    .map(|f| {
172                        ArbitraryArray {
173                            len: (len.end.saturating_sub(1))..len.end,
174                            is_nullable: f.is_nullable(),
175                            ..self.clone()
176                        }
177                        .sampler_from_data_type(f.data_type())
178                    })
179                    .collect(),
180            }),
181            DataType::List(field) => Box::new(ListSampler {
182                data_type: data_type.clone(),
183                len: len.clone(),
184                null: current_null,
185                inner: ArbitraryArray {
186                    branch: (self.branch.start * self.len.start)..(self.branch.end * self.len.end),
187                    is_nullable: field.is_nullable(),
188                    ..self.clone()
189                }
190                .sampler_from_data_type(field.data_type()),
191            }),
192            dt => panic!("not implemented: {:?}", dt),
193        }
194    }
195}
sample_arrow_rs/array.rs

sample_arrow_rs/
array.rs
