proptest.rs

use proptest::prelude::{Just, Strategy};
use proptest::prop_oneof;
use proptest::sample::SizeRange;

use super::StructArray;
use crate::bitmap::proptest::bitmap;
use crate::datatypes::{ArrowDataType, Field};

pub fn struct_array_with_fields(
    size_range: impl Into<SizeRange>,
    fields: Vec<Field>,
) -> impl Strategy<Value = StructArray> {
    let size_range = size_range.into();
    let (min, max) = size_range.start_end_incl();
    (min..=max).prop_flat_map(move |length| {
        let fields = fields.clone();

        let strategies: Vec<_> = fields
            .iter()
            .map(|field| crate::array::proptest::array_with_dtype(field.dtype.clone(), length))
            .collect();

        (
            strategies,
            prop_oneof![Just(None), bitmap(length).prop_map(Some)],
        )
            .prop_map(move |(child_arrays, validity)| {
                StructArray::new(
                    ArrowDataType::Struct(fields.clone()),
                    length,
                    child_arrays,
                    validity,
                )
            })
    })
}