|
20 | 20 | use crate::array::*; |
21 | 21 | use crate::datatypes::*; |
22 | 22 | use crate::util::test_util::seedable_rng; |
23 | | -use arrow_buffer::{Buffer, IntervalMonthDayNano}; |
| 23 | +use arrow_buffer::{Buffer, IntervalMonthDayNano, NullBuffer}; |
| 24 | +use arrow_schema::Field; |
24 | 25 | use half::f16; |
25 | 26 | use rand::Rng; |
26 | 27 | use rand::SeedableRng; |
27 | 28 | use rand::distr::uniform::SampleUniform; |
28 | 29 | use rand::rng; |
29 | 30 | use rand::{ |
30 | | - distr::{Alphanumeric, Distribution, StandardUniform}, |
| 31 | + distr::{Alphanumeric, Distribution, SampleString, StandardUniform}, |
31 | 32 | prelude::StdRng, |
32 | 33 | }; |
33 | 34 | use std::ops::Range; |
| 35 | +use std::sync::Arc; |
34 | 36 |
|
35 | 37 | /// Creates an random (but fixed-seeded) array of a given size and null density |
36 | 38 | pub fn create_primitive_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T> |
@@ -871,3 +873,77 @@ pub fn create_f64_array_with_seed(size: usize, nan_density: f32, seed: u64) -> F |
871 | 873 | }) |
872 | 874 | .collect() |
873 | 875 | } |
| 876 | + |
| 877 | +/// Create a FixedSizeList array of primitive values |
| 878 | +/// |
| 879 | +/// Arguments: |
| 880 | +/// - `size`: number of fixed-size lists in the array |
| 881 | +/// - `null_density`: density of nulls in the fixed-size list array (row-level nulls) |
| 882 | +/// - `value_null_density`: density of nulls in the primitive values inside each list |
| 883 | +/// - `list_size`: fixed size of each list element |
| 884 | +pub fn create_primitive_fixed_size_list_array<T>( |
| 885 | + size: usize, |
| 886 | + null_density: f32, |
| 887 | + value_null_density: f32, |
| 888 | + list_size: i32, |
| 889 | +) -> FixedSizeListArray |
| 890 | +where |
| 891 | + T: ArrowPrimitiveType, |
| 892 | + StandardUniform: Distribution<T::Native>, |
| 893 | +{ |
| 894 | + let mut rng = seedable_rng(); |
| 895 | + let list_size_usize = usize::try_from(list_size).expect("list_size must be non-negative"); |
| 896 | + let values: PrimitiveArray<T> = (0..size * list_size_usize) |
| 897 | + .map(|_| { |
| 898 | + if rng.random::<f32>() < value_null_density { |
| 899 | + None |
| 900 | + } else { |
| 901 | + Some(rng.random()) |
| 902 | + } |
| 903 | + }) |
| 904 | + .collect(); |
| 905 | + let field = Arc::new(Field::new("item", T::DATA_TYPE, value_null_density > 0.0)); |
| 906 | + let nulls = (null_density > 0.0).then(|| { |
| 907 | + NullBuffer::new(arrow_buffer::BooleanBuffer::collect_bool(size, |_| { |
| 908 | + rng.random::<f32>() >= null_density |
| 909 | + })) |
| 910 | + }); |
| 911 | + FixedSizeListArray::new(field, list_size, Arc::new(values), nulls) |
| 912 | +} |
| 913 | + |
| 914 | +/// Create a Map array with string keys and primitive values |
| 915 | +/// |
| 916 | +/// Arguments: |
| 917 | +/// - `size`: number of map entries in the array |
| 918 | +/// - `null_density`: density of nulls in the map array (row-level nulls) |
| 919 | +/// - `max_map_size`: maximum number of key-value pairs per map entry |
| 920 | +/// (actual size is random between 0 and max_map_size) |
| 921 | +/// - `key_len`: length of each random string key |
| 922 | +pub fn create_string_map_array<T>( |
| 923 | + size: usize, |
| 924 | + null_density: f32, |
| 925 | + max_map_size: usize, |
| 926 | + key_len: usize, |
| 927 | +) -> MapArray |
| 928 | +where |
| 929 | + T: ArrowPrimitiveType, |
| 930 | + StandardUniform: Distribution<T::Native>, |
| 931 | +{ |
| 932 | + let mut rng = seedable_rng(); |
| 933 | + let mut builder = MapBuilder::new(None, StringBuilder::new(), PrimitiveBuilder::<T>::new()); |
| 934 | + for _ in 0..size { |
| 935 | + if rng.random::<f32>() < null_density { |
| 936 | + builder.append(false).unwrap(); |
| 937 | + } else { |
| 938 | + let n = rng.random_range(0..=max_map_size); |
| 939 | + for _ in 0..n { |
| 940 | + builder |
| 941 | + .keys() |
| 942 | + .append_value(Alphanumeric.sample_string(&mut rng, key_len)); |
| 943 | + builder.values().append_value(rng.random()); |
| 944 | + } |
| 945 | + builder.append(true).unwrap(); |
| 946 | + } |
| 947 | + } |
| 948 | + builder.finish() |
| 949 | +} |
0 commit comments