arrow/util/
data_gen.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Utilities to generate random arrays and batches
19
20use std::sync::Arc;
21
22use rand::{
23    Rng,
24    distr::uniform::{SampleRange, SampleUniform},
25};
26
27use crate::array::*;
28use crate::error::{ArrowError, Result};
29use crate::{
30    buffer::{Buffer, MutableBuffer},
31    datatypes::*,
32};
33
34use super::{bench_util::*, bit_util, test_util::seedable_rng};
35
36/// Create a random [RecordBatch] from a schema
37pub fn create_random_batch(
38    schema: SchemaRef,
39    size: usize,
40    null_density: f32,
41    true_density: f32,
42) -> Result<RecordBatch> {
43    let columns = schema
44        .fields()
45        .iter()
46        .map(|field| create_random_array(field, size, null_density, true_density))
47        .collect::<Result<Vec<ArrayRef>>>()?;
48
49    RecordBatch::try_new_with_options(
50        schema,
51        columns,
52        &RecordBatchOptions::new().with_match_field_names(false),
53    )
54}
55
56/// Create a random [ArrayRef] from a [DataType] with a length,
57/// null density and true density (for [BooleanArray]).
58///
59/// # Arguments
60///
61/// * `field` - The field containing the data type for which to create a random array
62/// * `size` - The number of elements in the generated array
63/// * `null_density` - The approximate fraction of null values in the resulting array (0.0 to 1.0)
64/// * `true_density` - The approximate fraction of true values in boolean arrays (0.0 to 1.0)
65///
66pub fn create_random_array(
67    field: &Field,
68    size: usize,
69    mut null_density: f32,
70    true_density: f32,
71) -> Result<ArrayRef> {
72    // Override nullability in case of not nested and not dictionary
73    // For nested we don't want to override as we want to keep the nullability for the children
74    // For dictionary it handle the nullability internally
75    if !field.data_type().is_nested() && !matches!(field.data_type(), Dictionary(_, _)) {
76        // Override null density with 0.0 if the array is non-nullable
77        null_density = match field.is_nullable() {
78            true => null_density,
79            false => 0.0,
80        };
81    }
82
83    use DataType::*;
84    let array = match field.data_type() {
85        Null => Arc::new(NullArray::new(size)) as ArrayRef,
86        Boolean => Arc::new(create_boolean_array(size, null_density, true_density)),
87        Int8 => Arc::new(create_primitive_array::<Int8Type>(size, null_density)),
88        Int16 => Arc::new(create_primitive_array::<Int16Type>(size, null_density)),
89        Int32 => Arc::new(create_primitive_array::<Int32Type>(size, null_density)),
90        Int64 => Arc::new(create_primitive_array::<Int64Type>(size, null_density)),
91        UInt8 => Arc::new(create_primitive_array::<UInt8Type>(size, null_density)),
92        UInt16 => Arc::new(create_primitive_array::<UInt16Type>(size, null_density)),
93        UInt32 => Arc::new(create_primitive_array::<UInt32Type>(size, null_density)),
94        UInt64 => Arc::new(create_primitive_array::<UInt64Type>(size, null_density)),
95        Float16 => Arc::new(create_primitive_array::<Float16Type>(size, null_density)),
96        Float32 => Arc::new(create_primitive_array::<Float32Type>(size, null_density)),
97        Float64 => Arc::new(create_primitive_array::<Float64Type>(size, null_density)),
98        Timestamp(unit, tz) => match unit {
99            TimeUnit::Second => Arc::new(
100                create_random_temporal_array::<TimestampSecondType>(size, null_density)
101                    .with_timezone_opt(tz.clone()),
102            ) as ArrayRef,
103            TimeUnit::Millisecond => Arc::new(
104                create_random_temporal_array::<TimestampMillisecondType>(size, null_density)
105                    .with_timezone_opt(tz.clone()),
106            ),
107            TimeUnit::Microsecond => Arc::new(
108                create_random_temporal_array::<TimestampMicrosecondType>(size, null_density)
109                    .with_timezone_opt(tz.clone()),
110            ),
111            TimeUnit::Nanosecond => Arc::new(
112                create_random_temporal_array::<TimestampNanosecondType>(size, null_density)
113                    .with_timezone_opt(tz.clone()),
114            ),
115        },
116        Date32 => Arc::new(create_random_temporal_array::<Date32Type>(
117            size,
118            null_density,
119        )),
120        Date64 => Arc::new(create_random_temporal_array::<Date64Type>(
121            size,
122            null_density,
123        )),
124        Time32(unit) => match unit {
125            TimeUnit::Second => Arc::new(create_random_temporal_array::<Time32SecondType>(
126                size,
127                null_density,
128            )) as ArrayRef,
129            TimeUnit::Millisecond => Arc::new(
130                create_random_temporal_array::<Time32MillisecondType>(size, null_density),
131            ),
132            _ => {
133                return Err(ArrowError::InvalidArgumentError(format!(
134                    "Unsupported unit {unit:?} for Time32"
135                )));
136            }
137        },
138        Time64(unit) => match unit {
139            TimeUnit::Microsecond => Arc::new(
140                create_random_temporal_array::<Time64MicrosecondType>(size, null_density),
141            ) as ArrayRef,
142            TimeUnit::Nanosecond => Arc::new(create_random_temporal_array::<Time64NanosecondType>(
143                size,
144                null_density,
145            )),
146            _ => {
147                return Err(ArrowError::InvalidArgumentError(format!(
148                    "Unsupported unit {unit:?} for Time64"
149                )));
150            }
151        },
152        Utf8 => Arc::new(create_string_array::<i32>(size, null_density)),
153        LargeUtf8 => Arc::new(create_string_array::<i64>(size, null_density)),
154        Utf8View => Arc::new(create_string_view_array_with_len(
155            size,
156            null_density,
157            4,
158            false,
159        )),
160        Binary => Arc::new(create_binary_array::<i32>(size, null_density)),
161        LargeBinary => Arc::new(create_binary_array::<i64>(size, null_density)),
162        FixedSizeBinary(len) => {
163            let len = TryInto::<usize>::try_into(*len).map_err(|_| {
164                ArrowError::InvalidArgumentError(format!("cannot use FixedSizeBinary({len})"))
165            })?;
166            Arc::new(create_fsb_array(size, null_density, len))
167        }
168        BinaryView => Arc::new(
169            create_string_view_array_with_len(size, null_density, 4, false).to_binary_view(),
170        ),
171        List(_) => create_random_list_array(field, size, null_density, true_density)?,
172        LargeList(_) => create_random_list_array(field, size, null_density, true_density)?,
173        Struct(_) => create_random_struct_array(field, size, null_density, true_density)?,
174        d @ Dictionary(_, value_type) if crate::compute::can_cast_types(value_type, d) => {
175            let f = Field::new(
176                field.name(),
177                value_type.as_ref().clone(),
178                field.is_nullable(),
179            );
180            let v = create_random_array(&f, size, null_density, true_density)?;
181            crate::compute::cast(&v, d)?
182        }
183        Map(_, _) => create_random_map_array(field, size, null_density, true_density)?,
184        Decimal32(_, _) => create_random_decimal_array(field, size, null_density)?,
185        Decimal64(_, _) => create_random_decimal_array(field, size, null_density)?,
186        Decimal128(_, _) => create_random_decimal_array(field, size, null_density)?,
187        Decimal256(_, _) => create_random_decimal_array(field, size, null_density)?,
188        RunEndEncoded(index, value) => {
189            create_random_run_end_encoded_array(index, value, size, null_density, true_density)?
190        }
191        other => {
192            return Err(ArrowError::NotYetImplemented(format!(
193                "Generating random arrays not yet implemented for {other:?}"
194            )));
195        }
196    };
197
198    if !field.is_nullable() {
199        assert_eq!(array.null_count(), 0);
200    }
201
202    Ok(array)
203}
204
205#[inline]
206fn create_random_decimal_array(field: &Field, size: usize, null_density: f32) -> Result<ArrayRef> {
207    let mut rng = seedable_rng();
208
209    match field.data_type() {
210        DataType::Decimal32(precision, scale) => {
211            let values = (0..size)
212                .map(|_| {
213                    if rng.random::<f32>() < null_density {
214                        None
215                    } else {
216                        Some(rng.random::<i32>())
217                    }
218                })
219                .collect::<Vec<_>>();
220            Ok(Arc::new(
221                Decimal32Array::from(values).with_precision_and_scale(*precision, *scale)?,
222            ))
223        }
224        DataType::Decimal64(precision, scale) => {
225            let values = (0..size)
226                .map(|_| {
227                    if rng.random::<f32>() < null_density {
228                        None
229                    } else {
230                        Some(rng.random::<i64>())
231                    }
232                })
233                .collect::<Vec<_>>();
234            Ok(Arc::new(
235                Decimal64Array::from(values).with_precision_and_scale(*precision, *scale)?,
236            ))
237        }
238        DataType::Decimal128(precision, scale) => {
239            let values = (0..size)
240                .map(|_| {
241                    if rng.random::<f32>() < null_density {
242                        None
243                    } else {
244                        Some(rng.random::<i128>())
245                    }
246                })
247                .collect::<Vec<_>>();
248            Ok(Arc::new(
249                Decimal128Array::from(values).with_precision_and_scale(*precision, *scale)?,
250            ))
251        }
252        DataType::Decimal256(precision, scale) => {
253            let values = (0..size)
254                .map(|_| {
255                    if rng.random::<f32>() < null_density {
256                        None
257                    } else {
258                        Some(i256::from_parts(rng.random::<u128>(), rng.random::<i128>()))
259                    }
260                })
261                .collect::<Vec<_>>();
262            Ok(Arc::new(
263                Decimal256Array::from(values).with_precision_and_scale(*precision, *scale)?,
264            ))
265        }
266        _ => Err(ArrowError::InvalidArgumentError(format!(
267            "Cannot create decimal array for field {field}"
268        ))),
269    }
270}
271#[inline]
272fn create_random_run_end_encoded_array(
273    index: &Field,
274    value: &Field,
275    size: usize,
276    null_density: f32,
277    true_density: f32,
278) -> Result<ArrayRef> {
279    const MIN_RUN: usize = 8;
280    const MAX_RUN: usize = 32;
281
282    let mut rng = seedable_rng();
283    let mut run_lengths: Vec<usize> = Vec::new();
284    let mut remaining = size;
285    while remaining > 0 {
286        let len = rng.random_range(MIN_RUN..=MAX_RUN).min(remaining);
287        run_lengths.push(len);
288        remaining -= len;
289    }
290    let num_runs = run_lengths.len();
291
292    let mut cumulative: i64 = 0;
293    let run_ends_i64: Vec<i64> = run_lengths
294        .iter()
295        .map(|&l| {
296            cumulative += l as i64;
297            cumulative
298        })
299        .collect();
300
301    let values = create_random_array(value, num_runs, null_density, true_density)?;
302
303    match index.data_type() {
304        DataType::Int16 => {
305            let run_ends: Int16Array = run_ends_i64.iter().map(|&v| v as i16).collect();
306            Ok(Arc::new(RunArray::<Int16Type>::try_new(
307                &run_ends, &values,
308            )?))
309        }
310        DataType::Int32 => {
311            let run_ends: Int32Array = run_ends_i64.iter().map(|&v| v as i32).collect();
312            Ok(Arc::new(RunArray::<Int32Type>::try_new(
313                &run_ends, &values,
314            )?))
315        }
316        DataType::Int64 => {
317            let run_ends: Int64Array = run_ends_i64.iter().copied().collect();
318            Ok(Arc::new(RunArray::<Int64Type>::try_new(
319                &run_ends, &values,
320            )?))
321        }
322        other => Err(ArrowError::InvalidArgumentError(format!(
323            "Unsupported run-ends type for REE: {other:?}"
324        ))),
325    }
326}
327
328#[inline]
329fn create_random_list_array(
330    field: &Field,
331    size: usize,
332    null_density: f32,
333    true_density: f32,
334) -> Result<ArrayRef> {
335    // Override null density with 0.0 if the array is non-nullable
336    let list_null_density = match field.is_nullable() {
337        true => null_density,
338        false => 0.0,
339    };
340    let list_field;
341    let (offsets, child_len) = match field.data_type() {
342        DataType::List(f) => {
343            let (offsets, child_len) = create_random_offsets::<i32>(size, 0, 5);
344            list_field = f;
345            (Buffer::from(offsets.to_byte_slice()), child_len as usize)
346        }
347        DataType::LargeList(f) => {
348            let (offsets, child_len) = create_random_offsets::<i64>(size, 0, 5);
349            list_field = f;
350            (Buffer::from(offsets.to_byte_slice()), child_len as usize)
351        }
352        _ => {
353            return Err(ArrowError::InvalidArgumentError(format!(
354                "Cannot create list array for field {field}"
355            )));
356        }
357    };
358
359    // Create list's child data
360    let child_array = create_random_array(list_field, child_len, null_density, true_density)?;
361    let child_data = child_array.to_data();
362    // Create list's null buffers, if it is nullable
363    let null_buffer = match field.is_nullable() {
364        true => Some(create_random_null_buffer(size, list_null_density)),
365        false => None,
366    };
367    let list_data = unsafe {
368        ArrayData::new_unchecked(
369            field.data_type().clone(),
370            size,
371            None,
372            null_buffer,
373            0,
374            vec![offsets],
375            vec![child_data],
376        )
377    };
378    Ok(make_array(list_data))
379}
380
381#[inline]
382fn create_random_struct_array(
383    field: &Field,
384    size: usize,
385    null_density: f32,
386    true_density: f32,
387) -> Result<ArrayRef> {
388    let struct_fields = match field.data_type() {
389        DataType::Struct(fields) => fields,
390        _ => {
391            return Err(ArrowError::InvalidArgumentError(format!(
392                "Cannot create struct array for field {field}"
393            )));
394        }
395    };
396
397    let child_arrays = struct_fields
398        .iter()
399        .map(|struct_field| create_random_array(struct_field, size, null_density, true_density))
400        .collect::<Result<Vec<_>>>()?;
401
402    let null_buffer = match field.is_nullable() {
403        true => {
404            let nulls = arrow_buffer::BooleanBuffer::new(
405                create_random_null_buffer(size, null_density),
406                0,
407                size,
408            );
409            Some(nulls.into())
410        }
411        false => None,
412    };
413
414    Ok(Arc::new(StructArray::try_new(
415        struct_fields.clone(),
416        child_arrays,
417        null_buffer,
418    )?))
419}
420
421#[inline]
422fn create_random_map_array(
423    field: &Field,
424    size: usize,
425    null_density: f32,
426    true_density: f32,
427) -> Result<ArrayRef> {
428    // Override null density with 0.0 if the array is non-nullable
429    let map_null_density = match field.is_nullable() {
430        true => null_density,
431        false => 0.0,
432    };
433
434    let entries_field = match field.data_type() {
435        DataType::Map(f, _) => f,
436        _ => {
437            return Err(ArrowError::InvalidArgumentError(format!(
438                "Cannot create map array for field {field:?}"
439            )));
440        }
441    };
442
443    let (offsets, child_len) = create_random_offsets::<i32>(size, 0, 5);
444    let offsets = Buffer::from(offsets.to_byte_slice());
445
446    let entries = create_random_array(
447        entries_field,
448        child_len as usize,
449        null_density,
450        true_density,
451    )?
452    .to_data();
453
454    let null_buffer = match field.is_nullable() {
455        true => Some(create_random_null_buffer(size, map_null_density)),
456        false => None,
457    };
458
459    let map_data = unsafe {
460        ArrayData::new_unchecked(
461            field.data_type().clone(),
462            size,
463            None,
464            null_buffer,
465            0,
466            vec![offsets],
467            vec![entries],
468        )
469    };
470    Ok(make_array(map_data))
471}
472
473/// Generate random offsets for list arrays
474fn create_random_offsets<T: OffsetSizeTrait + SampleUniform>(
475    size: usize,
476    min: T,
477    max: T,
478) -> (Vec<T>, T) {
479    let rng = &mut seedable_rng();
480
481    let mut current_offset = T::zero();
482
483    let mut offsets = Vec::with_capacity(size + 1);
484    offsets.push(current_offset);
485
486    (0..size).for_each(|_| {
487        current_offset += rng.random_range(min..max);
488        offsets.push(current_offset);
489    });
490
491    (offsets, current_offset)
492}
493
494fn create_random_null_buffer(size: usize, null_density: f32) -> Buffer {
495    let mut rng = seedable_rng();
496    let mut mut_buf = MutableBuffer::new_null(size);
497    {
498        let mut_slice = mut_buf.as_slice_mut();
499        (0..size).for_each(|i| {
500            if rng.random::<f32>() >= null_density {
501                bit_util::set_bit(mut_slice, i)
502            }
503        })
504    };
505    mut_buf.into()
506}
507
508/// Useful for testing. The range of values are not likely to be representative of the
509/// actual bounds.
510pub trait RandomTemporalValue: ArrowTemporalType {
511    /// Returns the range of values for `impl`'d type
512    fn value_range() -> impl SampleRange<Self::Native>;
513
514    /// Generate a random value within the range of the type
515    fn gen_range<R: Rng>(rng: &mut R) -> Self::Native
516    where
517        Self::Native: SampleUniform,
518    {
519        rng.random_range(Self::value_range())
520    }
521
522    /// Generate a random value of the type
523    fn random<R: Rng>(rng: &mut R) -> Self::Native
524    where
525        Self::Native: SampleUniform,
526    {
527        Self::gen_range(rng)
528    }
529}
530
531impl RandomTemporalValue for TimestampSecondType {
532    /// Range of values for a timestamp in seconds. The range begins at the start
533    /// of the unix epoch and continues for 100 years.
534    fn value_range() -> impl SampleRange<Self::Native> {
535        0..60 * 60 * 24 * 365 * 100
536    }
537}
538
539impl RandomTemporalValue for TimestampMillisecondType {
540    /// Range of values for a timestamp in milliseconds. The range begins at the start
541    /// of the unix epoch and continues for 100 years.
542    fn value_range() -> impl SampleRange<Self::Native> {
543        0..1_000 * 60 * 60 * 24 * 365 * 100
544    }
545}
546
547impl RandomTemporalValue for TimestampMicrosecondType {
548    /// Range of values for a timestamp in microseconds. The range begins at the start
549    /// of the unix epoch and continues for 100 years.
550    fn value_range() -> impl SampleRange<Self::Native> {
551        0..1_000 * 1_000 * 60 * 60 * 24 * 365 * 100
552    }
553}
554
555impl RandomTemporalValue for TimestampNanosecondType {
556    /// Range of values for a timestamp in nanoseconds. The range begins at the start
557    /// of the unix epoch and continues for 100 years.
558    fn value_range() -> impl SampleRange<Self::Native> {
559        0..1_000 * 1_000 * 1_000 * 60 * 60 * 24 * 365 * 100
560    }
561}
562
563impl RandomTemporalValue for Date32Type {
564    /// Range of values representing the elapsed time since UNIX epoch in days. The
565    /// range begins at the start of the unix epoch and continues for 100 years.
566    fn value_range() -> impl SampleRange<Self::Native> {
567        0..365 * 100
568    }
569}
570
571impl RandomTemporalValue for Date64Type {
572    /// Range of values  representing the elapsed time since UNIX epoch in milliseconds.
573    /// The range begins at the start of the unix epoch and continues for 100 years.
574    fn value_range() -> impl SampleRange<Self::Native> {
575        0..1_000 * 60 * 60 * 24 * 365 * 100
576    }
577}
578
579impl RandomTemporalValue for Time32SecondType {
580    /// Range of values representing the elapsed time since midnight in seconds. The
581    /// range is from 0 to 24 hours.
582    fn value_range() -> impl SampleRange<Self::Native> {
583        0..60 * 60 * 24
584    }
585}
586
587impl RandomTemporalValue for Time32MillisecondType {
588    /// Range of values representing the elapsed time since midnight in milliseconds. The
589    /// range is from 0 to 24 hours.
590    fn value_range() -> impl SampleRange<Self::Native> {
591        0..1_000 * 60 * 60 * 24
592    }
593}
594
595impl RandomTemporalValue for Time64MicrosecondType {
596    /// Range of values representing the elapsed time since midnight in microseconds. The
597    /// range is from 0 to 24 hours.
598    fn value_range() -> impl SampleRange<Self::Native> {
599        0..1_000 * 1_000 * 60 * 60 * 24
600    }
601}
602
603impl RandomTemporalValue for Time64NanosecondType {
604    /// Range of values representing the elapsed time since midnight in nanoseconds. The
605    /// range is from 0 to 24 hours.
606    fn value_range() -> impl SampleRange<Self::Native> {
607        0..1_000 * 1_000 * 1_000 * 60 * 60 * 24
608    }
609}
610
611fn create_random_temporal_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T>
612where
613    T: RandomTemporalValue,
614    <T as ArrowPrimitiveType>::Native: SampleUniform,
615{
616    let mut rng = seedable_rng();
617
618    (0..size)
619        .map(|_| {
620            if rng.random::<f32>() < null_density {
621                None
622            } else {
623                Some(T::random(&mut rng))
624            }
625        })
626        .collect()
627}
628
629#[cfg(test)]
630mod tests {
631    use super::*;
632
633    #[test]
634    fn test_create_batch() {
635        let size = 32;
636        let fields = vec![
637            Field::new("a", DataType::Int32, true),
638            Field::new("f16", DataType::Float16, true),
639            Field::new(
640                "timestamp_without_timezone",
641                DataType::Timestamp(TimeUnit::Nanosecond, None),
642                true,
643            ),
644            Field::new(
645                "timestamp_with_timezone",
646                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
647                true,
648            ),
649        ];
650        let schema = Schema::new(fields);
651        let schema_ref = Arc::new(schema);
652        let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
653
654        assert_eq!(batch.schema(), schema_ref);
655        assert_eq!(batch.num_columns(), schema_ref.fields().len());
656        for array in batch.columns() {
657            assert_eq!(array.len(), size);
658        }
659    }
660
661    #[test]
662    fn test_create_batch_non_null() {
663        let size = 32;
664        let fields = vec![
665            Field::new("a", DataType::Int32, false),
666            Field::new(
667                "b",
668                DataType::List(Arc::new(Field::new_list_field(DataType::LargeUtf8, false))),
669                false,
670            ),
671            Field::new("a", DataType::Int32, false),
672        ];
673        let schema = Schema::new(fields);
674        let schema_ref = Arc::new(schema);
675        let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
676
677        assert_eq!(batch.schema(), schema_ref);
678        assert_eq!(batch.num_columns(), schema_ref.fields().len());
679        for array in batch.columns() {
680            assert_eq!(array.null_count(), 0);
681            assert_eq!(array.logical_null_count(), 0);
682        }
683        // Test that the list's child values are non-null
684        let b_array = batch.column(1);
685        let list_array = b_array.as_list::<i32>();
686        let child_array = list_array.values();
687        assert_eq!(child_array.null_count(), 0);
688        // There should be more values than the list, to show that it's a list
689        assert!(child_array.len() > list_array.len());
690    }
691
692    #[test]
693    fn test_create_struct_array() {
694        let size = 32;
695        let struct_fields = Fields::from(vec![
696            Field::new("b", DataType::Boolean, true),
697            Field::new(
698                "c",
699                DataType::LargeList(Arc::new(Field::new_list_field(
700                    DataType::List(Arc::new(Field::new_list_field(
701                        DataType::FixedSizeBinary(6),
702                        true,
703                    ))),
704                    false,
705                ))),
706                true,
707            ),
708            Field::new(
709                "d",
710                DataType::Struct(Fields::from(vec![
711                    Field::new("d_x", DataType::Int32, true),
712                    Field::new("d_y", DataType::Float32, false),
713                    Field::new("d_z", DataType::Binary, true),
714                ])),
715                true,
716            ),
717        ]);
718        let field = Field::new("struct", DataType::Struct(struct_fields), true);
719        let array = create_random_array(&field, size, 0.2, 0.5).unwrap();
720
721        assert_eq!(array.len(), 32);
722        let struct_array = array.as_any().downcast_ref::<StructArray>().unwrap();
723        assert_eq!(struct_array.columns().len(), 3);
724
725        // Test that the nested list makes sense,
726        // i.e. its children's values are more than the parent, to show repetition
727        let col_c = struct_array.column_by_name("c").unwrap();
728        let col_c = col_c.as_any().downcast_ref::<LargeListArray>().unwrap();
729        assert_eq!(col_c.len(), size);
730        let col_c_list = col_c.values().as_list::<i32>();
731        assert!(col_c_list.len() > size);
732        // Its values should be FixedSizeBinary(6)
733        let fsb = col_c_list.values();
734        assert_eq!(fsb.data_type(), &DataType::FixedSizeBinary(6));
735        assert!(fsb.len() > col_c_list.len());
736
737        // Test nested struct
738        let col_d = struct_array.column_by_name("d").unwrap();
739        let col_d = col_d.as_any().downcast_ref::<StructArray>().unwrap();
740        let col_d_y = col_d.column_by_name("d_y").unwrap();
741        assert_eq!(col_d_y.data_type(), &DataType::Float32);
742        assert_eq!(col_d_y.null_count(), 0);
743    }
744
745    #[test]
746    fn test_create_run_end_encoded_array() {
747        let size = 1000;
748        let ree_field = Field::new(
749            "ree",
750            DataType::RunEndEncoded(
751                Arc::new(Field::new("run_ends", DataType::Int32, false)),
752                Arc::new(Field::new("values", DataType::Utf8, true)),
753            ),
754            false,
755        );
756
757        let array = create_random_array(&ree_field, size, 0.25, 0.0).unwrap();
758        assert_eq!(array.len(), size);
759
760        let ree = array.as_run::<Int32Type>();
761        let run_ends = ree.run_ends().values();
762        let num_runs = run_ends.len();
763
764        assert_eq!(*run_ends.last().unwrap() as usize, size);
765
766        assert_eq!(ree.values().len(), num_runs);
767    }
768
769    #[test]
770    fn test_create_list_array_nested_nullability() {
771        let list_field = Field::new_list(
772            "not_null_list",
773            Field::new_list_field(DataType::Boolean, true),
774            false,
775        );
776
777        let list_array = create_random_array(&list_field, 100, 0.95, 0.5).unwrap();
778
779        assert_eq!(list_array.null_count(), 0);
780        assert!(list_array.as_list::<i32>().values().null_count() > 0);
781    }
782
783    #[test]
784    fn test_create_struct_array_nested_nullability() {
785        let struct_child_fields = vec![
786            Field::new("null_int", DataType::Int32, true),
787            Field::new("int", DataType::Int32, false),
788        ];
789        let struct_field = Field::new_struct("not_null_struct", struct_child_fields, false);
790
791        let struct_array = create_random_array(&struct_field, 100, 0.95, 0.5).unwrap();
792
793        assert_eq!(struct_array.null_count(), 0);
794        assert!(
795            struct_array
796                .as_struct()
797                .column_by_name("null_int")
798                .unwrap()
799                .null_count()
800                > 0
801        );
802        assert_eq!(
803            struct_array
804                .as_struct()
805                .column_by_name("int")
806                .unwrap()
807                .null_count(),
808            0
809        );
810    }
811
812    #[test]
813    fn test_create_list_array_nested_struct_nullability() {
814        let struct_child_fields = vec![
815            Field::new("null_int", DataType::Int32, true),
816            Field::new("int", DataType::Int32, false),
817        ];
818        let list_item_field =
819            Field::new_list_field(DataType::Struct(struct_child_fields.into()), true);
820        let list_field = Field::new_list("not_null_list", list_item_field, false);
821
822        let list_array = create_random_array(&list_field, 100, 0.95, 0.5).unwrap();
823
824        assert_eq!(list_array.null_count(), 0);
825        assert!(list_array.as_list::<i32>().values().null_count() > 0);
826        assert!(
827            list_array
828                .as_list::<i32>()
829                .values()
830                .as_struct()
831                .column_by_name("null_int")
832                .unwrap()
833                .null_count()
834                > 0
835        );
836        assert_eq!(
837            list_array
838                .as_list::<i32>()
839                .values()
840                .as_struct()
841                .column_by_name("int")
842                .unwrap()
843                .null_count(),
844            0
845        );
846    }
847
848    #[test]
849    fn test_create_map_array() {
850        let map_field = Field::new_map(
851            "map",
852            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
853            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
854            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
855            false,
856            false,
857        );
858        let array = create_random_array(&map_field, 100, 0.8, 0.5).unwrap();
859
860        assert_eq!(array.len(), 100);
861        // Map field is not null
862        assert_eq!(array.null_count(), 0);
863        assert_eq!(array.logical_null_count(), 0);
864        // Maps have multiple values like a list, so internal arrays are longer
865        assert!(array.as_map().keys().len() > array.len());
866        assert!(array.as_map().values().len() > array.len());
867        // Keys are not nullable
868        assert_eq!(array.as_map().keys().null_count(), 0);
869        // Values are nullable
870        assert!(array.as_map().values().null_count() > 0);
871
872        assert_eq!(array.as_map().keys().data_type(), &DataType::Utf8);
873        assert_eq!(array.as_map().values().data_type(), &DataType::Utf8);
874    }
875
876    #[test]
877    fn test_create_decimal_array() {
878        let size = 10;
879        let fields = vec![
880            Field::new("a", DataType::Decimal128(10, -2), true),
881            Field::new("b", DataType::Decimal256(10, -2), true),
882        ];
883        let schema = Schema::new(fields);
884        let schema_ref = Arc::new(schema);
885        let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
886
887        assert_eq!(batch.schema(), schema_ref);
888        assert_eq!(batch.num_columns(), schema_ref.fields().len());
889        for array in batch.columns() {
890            assert_eq!(array.len(), size);
891        }
892    }
893
894    #[test]
895    fn create_non_nullable_decimal_array_with_null_density() {
896        let size = 10;
897        let fields = vec![
898            Field::new("a", DataType::Decimal128(10, -2), false),
899            Field::new("b", DataType::Decimal256(10, -2), false),
900        ];
901        let schema = Schema::new(fields);
902        let schema_ref = Arc::new(schema);
903        let batch = create_random_batch(schema_ref.clone(), size, 0.35, 0.7).unwrap();
904
905        assert_eq!(batch.schema(), schema_ref);
906        assert_eq!(batch.num_columns(), schema_ref.fields().len());
907        for array in batch.columns() {
908            assert_eq!(array.len(), size);
909            assert_eq!(array.null_count(), 0);
910        }
911    }
912}
arrow/util/data_gen.rs

arrow/util/
data_gen.rs
