lance_datagen/

generator.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use std::{collections::HashMap, iter, marker::PhantomData, sync::Arc, sync::LazyLock};

use arrow::{
    array::{ArrayData, AsArray, Float32Builder, GenericBinaryBuilder, GenericStringBuilder},
    buffer::{BooleanBuffer, Buffer, OffsetBuffer, ScalarBuffer},
    datatypes::{
        ArrowPrimitiveType, Float32Type, Int32Type, Int64Type, IntervalDayTime,
        IntervalMonthDayNano, UInt32Type,
    },
};
use arrow_array::{
    Array, BinaryArray, FixedSizeBinaryArray, FixedSizeListArray, Float32Array, LargeListArray,
    LargeStringArray, ListArray, MapArray, NullArray, OffsetSizeTrait, PrimitiveArray, RecordBatch,
    RecordBatchOptions, RecordBatchReader, StringArray, StructArray, make_array,
    types::{ArrowDictionaryKeyType, BinaryType, ByteArrayType, Utf8Type},
};
use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema, SchemaRef};
use futures::{StreamExt, stream::BoxStream};
use rand::{Rng, RngCore, SeedableRng, distr::Uniform};
use rand_distr::Zipf;

use self::array::rand_with_distribution;

#[derive(Copy, Clone, Debug, Default)]
pub struct RowCount(u64);
#[derive(Copy, Clone, Debug, Default)]
pub struct BatchCount(u32);
#[derive(Copy, Clone, Debug, Default)]
pub struct ByteCount(u64);
#[derive(Copy, Clone, Debug, Default)]
pub struct Dimension(u32);

impl From<u32> for BatchCount {
    fn from(n: u32) -> Self {
        Self(n)
    }
}

impl From<u64> for RowCount {
    fn from(n: u64) -> Self {
        Self(n)
    }
}

impl From<u64> for ByteCount {
    fn from(n: u64) -> Self {
        Self(n)
    }
}

impl From<u32> for Dimension {
    fn from(n: u32) -> Self {
        Self(n)
    }
}

/// A trait for anything that can generate arrays of data
pub trait ArrayGenerator: Send + Sync + std::fmt::Debug {
    /// Generate an array of the given length
    ///
    /// # Arguments
    ///
    /// * `length` - The number of elements to generate
    /// * `rng` - The random number generator to use
    ///
    /// # Returns
    ///
    /// An array of the given length
    ///
    /// Note: Not every generator needs an rng.  However, it is passed here because many do and this
    /// lets us manage RNGs at the batch level instead of the array level.
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError>;

    /// Generate an array of the given length using a new RNG with the default seed
    ///
    /// # Arguments
    ///
    /// * `length` - The number of elements to generate
    ///
    /// # Returns
    ///
    /// An array of the given length
    fn generate_default(
        &mut self,
        length: RowCount,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        Self::generate(self, length, &mut rng)
    }
    /// Get the data type of the array that this generator produces
    ///
    /// # Returns
    ///
    /// The data type of the array that this generator produces
    fn data_type(&self) -> &DataType;
    /// Gets metadata that should be associated with the field generated by this generator
    fn metadata(&self) -> Option<HashMap<String, String>> {
        None
    }
    /// Get the size of each element in bytes
    ///
    /// # Returns
    ///
    /// The size of each element in bytes.  Will be None if the size varies by element.
    fn element_size_bytes(&self) -> Option<ByteCount>;
}

#[derive(Debug)]
pub struct CycleNullGenerator {
    generator: Box<dyn ArrayGenerator>,
    validity: Vec<bool>,
    idx: usize,
}
#[derive(Debug)]
pub struct CycleNanGenerator {
    generator: Box<dyn ArrayGenerator>,
    nan_pattern: Vec<bool>,
    idx: usize,
}

impl ArrayGenerator for CycleNanGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let array = self.generator.generate(length, rng)?;

        // Only apply NaN pattern to float types
        match array.data_type() {
            DataType::Float16 => {
                let float_array = array
                    .as_any()
                    .downcast_ref::<arrow_array::Float16Array>()
                    .unwrap();
                let mut values: Vec<half::f16> = float_array.values().to_vec();

                for (i, &should_be_nan) in self
                    .nan_pattern
                    .iter()
                    .cycle()
                    .skip(self.idx)
                    .take(length.0 as usize)
                    .enumerate()
                {
                    if should_be_nan {
                        values[i] = half::f16::NAN;
                    }
                }

                self.idx = (self.idx + (length.0 as usize)) % self.nan_pattern.len();
                Ok(Arc::new(arrow_array::Float16Array::from(values)))
            }
            DataType::Float32 => {
                let float_array = array
                    .as_any()
                    .downcast_ref::<arrow_array::Float32Array>()
                    .unwrap();
                let mut values: Vec<f32> = float_array.values().to_vec();

                for (i, &should_be_nan) in self
                    .nan_pattern
                    .iter()
                    .cycle()
                    .skip(self.idx)
                    .take(length.0 as usize)
                    .enumerate()
                {
                    if should_be_nan {
                        values[i] = f32::NAN;
                    }
                }

                self.idx = (self.idx + (length.0 as usize)) % self.nan_pattern.len();
                Ok(Arc::new(arrow_array::Float32Array::from(values)))
            }
            DataType::Float64 => {
                let float_array = array
                    .as_any()
                    .downcast_ref::<arrow_array::Float64Array>()
                    .unwrap();
                let mut values: Vec<f64> = float_array.values().to_vec();

                for (i, &should_be_nan) in self
                    .nan_pattern
                    .iter()
                    .cycle()
                    .skip(self.idx)
                    .take(length.0 as usize)
                    .enumerate()
                {
                    if should_be_nan {
                        values[i] = f64::NAN;
                    }
                }

                self.idx = (self.idx + (length.0 as usize)) % self.nan_pattern.len();
                Ok(Arc::new(arrow_array::Float64Array::from(values)))
            }
            _ => {
                // For non-float types, just return the original array unchanged
                Ok(array)
            }
        }
    }

    fn data_type(&self) -> &DataType {
        self.generator.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.generator.element_size_bytes()
    }
}

impl ArrayGenerator for CycleNullGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let array = self.generator.generate(length, rng)?;
        let data = array.to_data();
        let validity_itr = self
            .validity
            .iter()
            .cycle()
            .skip(self.idx)
            .take(length.0 as usize)
            .copied();
        let validity_bitmap = BooleanBuffer::from_iter(validity_itr);

        self.idx = (self.idx + (length.0 as usize)) % self.validity.len();
        unsafe {
            let new_data = ArrayData::new_unchecked(
                data.data_type().clone(),
                data.len(),
                None,
                Some(validity_bitmap.into_inner()),
                data.offset(),
                data.buffers().to_vec(),
                data.child_data().into(),
            );
            Ok(make_array(new_data))
        }
    }

    fn data_type(&self) -> &DataType {
        self.generator.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.generator.element_size_bytes()
    }
}

#[derive(Debug)]
pub struct MetadataGenerator {
    generator: Box<dyn ArrayGenerator>,
    metadata: HashMap<String, String>,
}

impl ArrayGenerator for MetadataGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        self.generator.generate(length, rng)
    }

    fn metadata(&self) -> Option<HashMap<String, String>> {
        Some(self.metadata.clone())
    }

    fn data_type(&self) -> &DataType {
        self.generator.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.generator.element_size_bytes()
    }
}

#[derive(Debug)]
pub struct NullGenerator {
    generator: Box<dyn ArrayGenerator>,
    null_probability: f64,
}

impl ArrayGenerator for NullGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let array = self.generator.generate(length, rng)?;
        let data = array.to_data();

        if self.null_probability < 0.0 || self.null_probability > 1.0 {
            return Err(ArrowError::InvalidArgumentError(format!(
                "null_probability must be between 0 and 1, got {}",
                self.null_probability
            )));
        }

        let (null_count, new_validity) = if self.null_probability == 0.0 {
            if data.null_count() == 0 {
                return Ok(array);
            } else {
                (0_usize, None)
            }
        } else if self.null_probability == 1.0 {
            if data.null_count() == data.len() {
                return Ok(array);
            } else {
                let all_nulls = BooleanBuffer::new_unset(array.len());
                (array.len(), Some(all_nulls.into_inner()))
            }
        } else {
            let array_len = array.len();
            let num_validity_bytes = array_len.div_ceil(8);
            let mut null_count = 0;
            // Sampling the RNG once per bit is kind of slow so we do this to sample once
            // per byte.  We only get 8 bits of RNG resolution but that should be good enough.
            let threshold = (self.null_probability * u8::MAX as f64) as u8;
            let bytes = (0..num_validity_bytes)
                .map(|byte_idx| {
                    let mut sample = rng.random::<u64>();
                    let mut byte: u8 = 0;
                    for bit_idx in 0..8 {
                        // We could probably overshoot and fill in extra bits with random data but
                        // this is cleaner and that would mess up the null count
                        byte <<= 1;
                        let pos = byte_idx * 8 + (7 - bit_idx);
                        if pos < array_len {
                            let sample_piece = sample & 0xFF;
                            let is_null = (sample_piece as u8) < threshold;
                            byte |= (!is_null) as u8;
                            null_count += is_null as usize;
                        }
                        sample >>= 8;
                    }
                    byte
                })
                .collect::<Vec<_>>();
            let new_validity = Buffer::from_iter(bytes);
            (null_count, Some(new_validity))
        };

        unsafe {
            let new_data = ArrayData::new_unchecked(
                data.data_type().clone(),
                data.len(),
                Some(null_count),
                new_validity,
                data.offset(),
                data.buffers().to_vec(),
                data.child_data().into(),
            );
            Ok(make_array(new_data))
        }
    }

    fn metadata(&self) -> Option<HashMap<String, String>> {
        self.generator.metadata()
    }

    fn data_type(&self) -> &DataType {
        self.generator.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.generator.element_size_bytes()
    }
}

pub trait ArrayGeneratorExt {
    /// Replaces the validity bitmap of generated arrays, inserting nulls with a given probability
    fn with_random_nulls(self, null_probability: f64) -> Box<dyn ArrayGenerator>;
    /// Replaces the validity bitmap of generated arrays with the inverse of `nulls`, cycling if needed
    fn with_nulls(self, nulls: &[bool]) -> Box<dyn ArrayGenerator>;
    /// Replaces the values of generated arrays with NaN values, cycling if needed
    ///
    /// Will have no effect if the data type is not a floating point data type
    fn with_nans(self, nans: &[bool]) -> Box<dyn ArrayGenerator>;
    /// Replaces the validity bitmap of generated arrays with `validity`, cycling if needed
    fn with_validity(self, nulls: &[bool]) -> Box<dyn ArrayGenerator>;
    fn with_metadata(self, metadata: HashMap<String, String>) -> Box<dyn ArrayGenerator>;
}

impl ArrayGeneratorExt for Box<dyn ArrayGenerator> {
    fn with_random_nulls(self, null_probability: f64) -> Box<dyn ArrayGenerator> {
        Box::new(NullGenerator {
            generator: self,
            null_probability,
        })
    }

    fn with_nulls(self, nulls: &[bool]) -> Box<dyn ArrayGenerator> {
        Box::new(CycleNullGenerator {
            generator: self,
            validity: nulls.iter().map(|v| !*v).collect(),
            idx: 0,
        })
    }

    fn with_nans(self, nans: &[bool]) -> Box<dyn ArrayGenerator> {
        Box::new(CycleNanGenerator {
            generator: self,
            nan_pattern: nans.to_vec(),
            idx: 0,
        })
    }

    fn with_validity(self, validity: &[bool]) -> Box<dyn ArrayGenerator> {
        Box::new(CycleNullGenerator {
            generator: self,
            validity: validity.to_vec(),
            idx: 0,
        })
    }

    fn with_metadata(self, metadata: HashMap<String, String>) -> Box<dyn ArrayGenerator> {
        Box::new(MetadataGenerator {
            generator: self,
            metadata,
        })
    }
}

pub struct NTimesIter<I: Iterator>
where
    I::Item: Copy,
{
    iter: I,
    n: u32,
    cur: I::Item,
    count: u32,
}

// Note: if this is used then there is a performance hit as the
// inner loop cannot experience vectorization
//
// TODO: maybe faster to build the vec and then repeat it into
// the destination array?
impl<I: Iterator> Iterator for NTimesIter<I>
where
    I::Item: Copy,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.count == 0 {
            self.count = self.n - 1;
            self.cur = self.iter.next()?;
        } else {
            self.count -= 1;
        }
        Some(self.cur)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (lower, upper) = self.iter.size_hint();
        let lower = lower * self.n as usize;
        let upper = upper.map(|u| u * self.n as usize);
        (lower, upper)
    }
}

pub struct FnGen<T, ArrayType, F: FnMut(&mut rand_xoshiro::Xoshiro256PlusPlus) -> T>
where
    T: Copy + Default,
    ArrayType: arrow_array::Array + From<Vec<T>>,
{
    data_type: DataType,
    generator: F,
    array_type: PhantomData<ArrayType>,
    repeat: u32,
    leftover: T,
    leftover_count: u32,
    element_size_bytes: Option<ByteCount>,
}

impl<T, ArrayType, F: FnMut(&mut rand_xoshiro::Xoshiro256PlusPlus) -> T> std::fmt::Debug
    for FnGen<T, ArrayType, F>
where
    T: Copy + Default,
    ArrayType: arrow_array::Array + From<Vec<T>>,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FnGen")
            .field("data_type", &self.data_type)
            .field("array_type", &self.array_type)
            .field("repeat", &self.repeat)
            .field("leftover_count", &self.leftover_count)
            .field("element_size_bytes", &self.element_size_bytes)
            .finish()
    }
}

impl<T, ArrayType, F: FnMut(&mut rand_xoshiro::Xoshiro256PlusPlus) -> T> FnGen<T, ArrayType, F>
where
    T: Copy + Default,
    ArrayType: arrow_array::Array + From<Vec<T>>,
{
    fn new_known_size(
        data_type: DataType,
        generator: F,
        repeat: u32,
        element_size_bytes: ByteCount,
    ) -> Self {
        Self {
            data_type,
            generator,
            array_type: PhantomData,
            repeat,
            leftover: T::default(),
            leftover_count: 0,
            element_size_bytes: Some(element_size_bytes),
        }
    }

    fn new_unknown_size(data_type: DataType, generator: F, repeat: u32) -> Self {
        Self {
            data_type,
            generator,
            array_type: PhantomData,
            repeat,
            leftover: T::default(),
            leftover_count: 0,
            element_size_bytes: None,
        }
    }
}

impl<T, ArrayType, F: FnMut(&mut rand_xoshiro::Xoshiro256PlusPlus) -> T> ArrayGenerator
    for FnGen<T, ArrayType, F>
where
    T: Copy + Default + Send + Sync,
    ArrayType: arrow_array::Array + From<Vec<T>> + 'static,
    F: Send + Sync,
{
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let iter = (0..length.0).map(|_| (self.generator)(rng));
        let values = if self.repeat > 1 {
            Vec::from_iter(
                NTimesIter {
                    iter,
                    n: self.repeat,
                    cur: self.leftover,
                    count: self.leftover_count,
                }
                .take(length.0 as usize),
            )
        } else {
            Vec::from_iter(iter)
        };
        self.leftover_count = ((self.leftover_count as u64 + length.0) % self.repeat as u64) as u32;
        self.leftover = values.last().copied().unwrap_or(T::default());
        Ok(Arc::new(ArrayType::from(values)))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.element_size_bytes
    }
}

#[derive(Copy, Clone, Debug)]
pub struct Seed(pub u64);
pub const DEFAULT_SEED: Seed = Seed(42);

impl From<u64> for Seed {
    fn from(n: u64) -> Self {
        Self(n)
    }
}

#[derive(Debug)]
pub struct CycleVectorGenerator {
    underlying_gen: Box<dyn ArrayGenerator>,
    dimension: Dimension,
    data_type: DataType,
}

impl CycleVectorGenerator {
    pub fn new(underlying_gen: Box<dyn ArrayGenerator>, dimension: Dimension) -> Self {
        let data_type = DataType::FixedSizeList(
            Arc::new(Field::new("item", underlying_gen.data_type().clone(), true)),
            dimension.0 as i32,
        );
        Self {
            underlying_gen,
            dimension,
            data_type,
        }
    }
}

impl ArrayGenerator for CycleVectorGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let values = self
            .underlying_gen
            .generate(RowCount::from(length.0 * self.dimension.0 as u64), rng)?;
        let field = Arc::new(Field::new("item", values.data_type().clone(), true));
        let values = Arc::new(values);

        let array = FixedSizeListArray::try_new(field, self.dimension.0 as i32, values, None)?;

        Ok(Arc::new(array))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.underlying_gen
            .element_size_bytes()
            .map(|byte_count| ByteCount::from(byte_count.0 * self.dimension.0 as u64))
    }
}

#[derive(Debug)]
pub struct CycleListGenerator {
    underlying_gen: Box<dyn ArrayGenerator>,
    lengths_gen: Box<dyn ArrayGenerator>,
    data_type: DataType,
}

impl CycleListGenerator {
    pub fn new(
        underlying_gen: Box<dyn ArrayGenerator>,
        min_list_size: Dimension,
        max_list_size: Dimension,
    ) -> Self {
        let data_type = DataType::List(Arc::new(Field::new(
            "item",
            underlying_gen.data_type().clone(),
            true,
        )));
        let lengths_dist = Uniform::new(min_list_size.0, max_list_size.0).unwrap();
        let lengths_gen = rand_with_distribution::<UInt32Type, Uniform<u32>>(lengths_dist);
        Self {
            underlying_gen,
            lengths_gen,
            data_type,
        }
    }
}

impl ArrayGenerator for CycleListGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let lengths = self.lengths_gen.generate(length, rng)?;
        let lengths = lengths.as_primitive::<UInt32Type>();
        let total_length = lengths.values().iter().map(|i| *i as u64).sum::<u64>();
        let offsets = OffsetBuffer::from_lengths(lengths.values().iter().map(|v| *v as usize));
        let values = self
            .underlying_gen
            .generate(RowCount::from(total_length), rng)?;
        let field = Arc::new(Field::new("item", values.data_type().clone(), true));
        let values = Arc::new(values);

        let array = ListArray::try_new(field, offsets, values, None)?;

        Ok(Arc::new(array))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        None
    }
}

#[derive(Debug, Default)]
pub struct PseudoUuidGenerator {}

impl ArrayGenerator for PseudoUuidGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        Ok(Arc::new(FixedSizeBinaryArray::try_from_iter(
            (0..length.0).map(|_| {
                let mut data = vec![0; 16];
                rng.fill_bytes(&mut data);
                data
            }),
        )?))
    }

    fn data_type(&self) -> &DataType {
        &DataType::FixedSizeBinary(16)
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(16))
    }
}

#[derive(Debug, Default)]
pub struct PseudoUuidHexGenerator {}

impl ArrayGenerator for PseudoUuidHexGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let mut data = vec![0; 16 * length.0 as usize];
        rng.fill_bytes(&mut data);
        let data_hex = hex::encode(data);

        Ok(Arc::new(StringArray::from_iter_values(
            (0..length.0 as usize).map(|i| data_hex.get(i * 32..(i + 1) * 32).unwrap()),
        )))
    }

    fn data_type(&self) -> &DataType {
        &DataType::Utf8
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(16))
    }
}

#[derive(Debug, Default)]
pub struct RandomBooleanGenerator {}

impl ArrayGenerator for RandomBooleanGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let num_bytes = length.0.div_ceil(8);
        let mut bytes = vec![0; num_bytes as usize];
        rng.fill_bytes(&mut bytes);
        let bytes = BooleanBuffer::new(Buffer::from(bytes), 0, length.0 as usize);
        Ok(Arc::new(arrow_array::BooleanArray::new(bytes, None)))
    }

    fn data_type(&self) -> &DataType {
        &DataType::Boolean
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // We can't say 1/8th of a byte and 1 byte would be a pretty extreme over-count so let's leave
        // it at None until someone needs this.  Then we can probably special case this (e.g. make a ByteCount::ONE_BIT)
        None
    }
}

// Instead of using the "standard distribution" and generating values there are some cases (e.g. f16 / decimal)
// where we just generate random bytes because there is no rand support
pub struct RandomBytesGenerator<T: ArrowPrimitiveType + Send + Sync> {
    phantom: PhantomData<T>,
    data_type: DataType,
}

impl<T: ArrowPrimitiveType + Send + Sync> std::fmt::Debug for RandomBytesGenerator<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RandomBytesGenerator")
            .field("data_type", &self.data_type)
            .finish()
    }
}

impl<T: ArrowPrimitiveType + Send + Sync> RandomBytesGenerator<T> {
    fn new(data_type: DataType) -> Self {
        Self {
            phantom: Default::default(),
            data_type,
        }
    }

    fn byte_width() -> Result<u64, ArrowError> {
        T::DATA_TYPE.primitive_width().ok_or_else(|| ArrowError::InvalidArgumentError(format!("Cannot generate the data type {} with the RandomBytesGenerator because it is not a fixed-width bytes type", T::DATA_TYPE))).map(|val| val as u64)
    }
}

impl<T: ArrowPrimitiveType + Send + Sync> ArrayGenerator for RandomBytesGenerator<T> {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let num_bytes = length.0 * Self::byte_width()?;
        let mut bytes = vec![0; num_bytes as usize];
        rng.fill_bytes(&mut bytes);
        let bytes = ScalarBuffer::new(Buffer::from(bytes), 0, length.0 as usize);
        Ok(Arc::new(
            PrimitiveArray::<T>::new(bytes, None).with_data_type(self.data_type.clone()),
        ))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Self::byte_width().map(ByteCount::from).ok()
    }
}

// This is pretty much the same thing as RandomBinaryGenerator but we can't use that
// because there is no ArrowPrimitiveType for FixedSizeBinary
#[derive(Debug)]
pub struct RandomFixedSizeBinaryGenerator {
    data_type: DataType,
    size: i32,
}

impl RandomFixedSizeBinaryGenerator {
    fn new(size: i32) -> Self {
        Self {
            size,
            data_type: DataType::FixedSizeBinary(size),
        }
    }
}

impl ArrayGenerator for RandomFixedSizeBinaryGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let num_bytes = length.0 * self.size as u64;
        let mut bytes = vec![0; num_bytes as usize];
        rng.fill_bytes(&mut bytes);
        Ok(Arc::new(FixedSizeBinaryArray::new(
            self.size,
            Buffer::from(bytes),
            None,
        )))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(self.size as u64))
    }
}

#[derive(Debug)]
pub struct RandomIntervalGenerator {
    unit: IntervalUnit,
    data_type: DataType,
}

impl RandomIntervalGenerator {
    pub fn new(unit: IntervalUnit) -> Self {
        Self {
            unit,
            data_type: DataType::Interval(unit),
        }
    }
}

impl ArrayGenerator for RandomIntervalGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        match self.unit {
            IntervalUnit::YearMonth => {
                let months = (0..length.0)
                    .map(|_| rng.random::<i32>())
                    .collect::<Vec<_>>();
                Ok(Arc::new(arrow_array::IntervalYearMonthArray::from(months)))
            }
            IntervalUnit::MonthDayNano => {
                let day_time_array = (0..length.0)
                    .map(|_| IntervalMonthDayNano::new(rng.random(), rng.random(), rng.random()))
                    .collect::<Vec<_>>();
                Ok(Arc::new(arrow_array::IntervalMonthDayNanoArray::from(
                    day_time_array,
                )))
            }
            IntervalUnit::DayTime => {
                let day_time_array = (0..length.0)
                    .map(|_| IntervalDayTime::new(rng.random(), rng.random()))
                    .collect::<Vec<_>>();
                Ok(Arc::new(arrow_array::IntervalDayTimeArray::from(
                    day_time_array,
                )))
            }
        }
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(12))
    }
}
#[derive(Debug)]
pub struct RandomBinaryGenerator {
    bytes_per_element: ByteCount,
    scale_to_utf8: bool,
    is_large: bool,
    data_type: DataType,
}

impl RandomBinaryGenerator {
    pub fn new(bytes_per_element: ByteCount, scale_to_utf8: bool, is_large: bool) -> Self {
        Self {
            bytes_per_element,
            scale_to_utf8,
            is_large,
            data_type: match (scale_to_utf8, is_large) {
                (false, false) => DataType::Binary,
                (false, true) => DataType::LargeBinary,
                (true, false) => DataType::Utf8,
                (true, true) => DataType::LargeUtf8,
            },
        }
    }
}

impl ArrayGenerator for RandomBinaryGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let mut bytes = vec![0; (self.bytes_per_element.0 * length.0) as usize];
        rng.fill_bytes(&mut bytes);
        if self.scale_to_utf8 {
            // This doesn't give us the full UTF-8 range and it isn't statistically correct but
            // it's fast and probably good enough for most cases
            bytes = bytes.into_iter().map(|val| (val % 95) + 32).collect();
        }
        let bytes = Buffer::from(bytes);
        if self.is_large {
            let offsets = OffsetBuffer::from_lengths(iter::repeat_n(
                self.bytes_per_element.0 as usize,
                length.0 as usize,
            ));
            if self.scale_to_utf8 {
                // This is safe because we are only using printable characters
                unsafe {
                    Ok(Arc::new(arrow_array::LargeStringArray::new_unchecked(
                        offsets, bytes, None,
                    )))
                }
            } else {
                unsafe {
                    Ok(Arc::new(arrow_array::LargeBinaryArray::new_unchecked(
                        offsets, bytes, None,
                    )))
                }
            }
        } else {
            let offsets = OffsetBuffer::from_lengths(iter::repeat_n(
                self.bytes_per_element.0 as usize,
                length.0 as usize,
            ));
            if self.scale_to_utf8 {
                // This is safe because we are only using printable characters
                unsafe {
                    Ok(Arc::new(arrow_array::StringArray::new_unchecked(
                        offsets, bytes, None,
                    )))
                }
            } else {
                unsafe {
                    Ok(Arc::new(arrow_array::BinaryArray::new_unchecked(
                        offsets, bytes, None,
                    )))
                }
            }
        }
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // Not exactly correct since there are N + 1 4-byte offsets and this only counts N
        Some(ByteCount::from(
            self.bytes_per_element.0 + std::mem::size_of::<i32>() as u64,
        ))
    }
}

/// Generate a sequence of strings with a prefix and a counter
///
/// For example, if the prefix is "user_" the strings will be "user_0", "user_1", ...
#[derive(Debug)]
pub struct PrefixPlusCounterGenerator {
    prefix: String,
    is_large: bool,
    data_type: DataType,
    current_counter: u64,
}

impl PrefixPlusCounterGenerator {
    pub fn new(prefix: String, is_large: bool) -> Self {
        Self {
            prefix,
            is_large,
            data_type: if is_large {
                DataType::LargeUtf8
            } else {
                DataType::Utf8
            },
            current_counter: 0,
        }
    }

    fn generate_values<T: OffsetSizeTrait>(
        &self,
        start: u64,
        num_values: u64,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let max_counter = start + num_values;
        let max_digits_per_counter = (max_counter as f64).log10().ceil() as u64;
        let max_bytes_per_str = max_digits_per_counter + self.prefix.len() as u64;
        let max_bytes = max_bytes_per_str * num_values;
        let mut builder =
            GenericStringBuilder::<T>::with_capacity(num_values as usize, max_bytes as usize);
        let mut word = String::with_capacity(max_bytes_per_str as usize);
        word.push_str(&self.prefix);
        for i in 0..num_values {
            let counter = start + i;
            word.truncate(self.prefix.len());
            word.push_str(&counter.to_string());
            builder.append_value(&word);
        }
        Ok(Arc::new(builder.finish()))
    }
}

impl ArrayGenerator for PrefixPlusCounterGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let start = self.current_counter;
        self.current_counter += length.0;
        if self.is_large {
            self.generate_values::<i64>(start, length.0)
        } else {
            self.generate_values::<i32>(start, length.0)
        }
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // It's not consistent
        None
    }
}

/// Generate a sequence of binary strings with a prefix and a counter
///
/// The counter will be encoded (little-endian) as a u8, u16, u32, or u64 and added to the prefix
/// As long as more than 256 values are generated then the resulting array will have
/// variable width
#[derive(Debug)]
pub struct BinaryPrefixPlusCounterGenerator {
    prefix: Arc<[u8]>,
    is_large: bool,
    data_type: DataType,
    current_counter: u64,
}

impl BinaryPrefixPlusCounterGenerator {
    pub fn new(prefix: Arc<[u8]>, is_large: bool) -> Self {
        Self {
            prefix,
            is_large,
            data_type: if is_large {
                DataType::LargeBinary
            } else {
                DataType::Binary
            },
            current_counter: 0,
        }
    }

    fn generate_values<T: OffsetSizeTrait>(
        &self,
        start: u64,
        num_values: u64,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let max_bytes = (self.prefix.len() + std::mem::size_of::<u64>()) * num_values as usize;
        let mut builder = GenericBinaryBuilder::<T>::with_capacity(num_values as usize, max_bytes);
        let mut word = Vec::with_capacity(self.prefix.len() + std::mem::size_of::<u64>());
        word.extend_from_slice(&self.prefix);
        for i in 0..num_values {
            let counter = start + i;
            word.truncate(self.prefix.len());
            if counter < u8::MAX as u64 {
                word.push(counter as u8);
            } else if counter < u16::MAX as u64 {
                word.extend_from_slice(&(counter as u16).to_le_bytes());
            } else if counter < u32::MAX as u64 {
                word.extend_from_slice(&(counter as u32).to_le_bytes());
            } else {
                word.extend_from_slice(&counter.to_le_bytes());
            }
            builder.append_value(&word);
        }
        Ok(Arc::new(builder.finish()))
    }
}

impl ArrayGenerator for BinaryPrefixPlusCounterGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let start = self.current_counter;
        self.current_counter += length.0;
        if self.is_large {
            self.generate_values::<i64>(start, length.0)
        } else {
            self.generate_values::<i32>(start, length.0)
        }
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // It's not consistent
        None
    }
}

// Common English stop words placed at the front to be sampled more frequently.
const STOP_WORDS: &[&str] = &[
    "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", "in", "into", "is", "it",
    "no", "not", "of", "on", "or", "such", "that", "the", "their", "then", "there", "these",
    "they", "this", "to", "was", "will", "with",
];

const ENGLISH_WORDS: &[&str] = &[
    "ability",
    "able",
    "about",
    "above",
    "accept",
    "access",
    "account",
    "across",
    "action",
    "active",
    "activity",
    "actual",
    "address",
    "adjust",
    "admin",
    "advance",
    "agent",
    "align",
    "allow",
    "amount",
    "analysis",
    "answer",
    "application",
    "archive",
    "array",
    "asset",
    "async",
    "attribute",
    "available",
    "balance",
    "batch",
    "binary",
    "bitmap",
    "block",
    "branch",
    "buffer",
    "build",
    "cache",
    "capacity",
    "catalog",
    "change",
    "chunk",
    "client",
    "cluster",
    "column",
    "commit",
    "common",
    "compare",
    "compile",
    "compute",
    "condition",
    "config",
    "connect",
    "content",
    "context",
    "control",
    "convert",
    "copy",
    "core",
    "count",
    "create",
    "current",
    "cursor",
    "data",
    "dataset",
    "decode",
    "default",
    "delete",
    "delta",
    "depend",
    "derive",
    "design",
    "detail",
    "detect",
    "device",
    "direct",
    "display",
    "document",
    "domain",
    "drive",
    "dynamic",
    "encode",
    "engine",
    "error",
    "event",
    "example",
    "execute",
    "expand",
    "expect",
    "export",
    "extend",
    "feature",
    "field",
    "filter",
    "final",
    "finish",
    "format",
    "fragment",
    "future",
    "generate",
    "global",
    "group",
    "handle",
    "header",
    "index",
    "input",
    "insert",
    "inspect",
    "instance",
    "integer",
    "internal",
    "item",
    "join",
    "kernel",
    "large",
    "layer",
    "layout",
    "length",
    "level",
    "limit",
    "linear",
    "local",
    "logical",
    "lookup",
    "manage",
    "manifest",
    "memory",
    "merge",
    "metric",
    "model",
    "module",
    "namespace",
    "native",
    "node",
    "normal",
    "number",
    "object",
    "offset",
    "option",
    "output",
    "package",
    "page",
    "parallel",
    "parse",
    "partition",
    "pattern",
    "physical",
    "plan",
    "policy",
    "prefix",
    "prepare",
    "primary",
    "process",
    "profile",
    "project",
    "property",
    "query",
    "range",
    "reader",
    "record",
    "region",
    "registry",
    "request",
    "resolve",
    "resource",
    "result",
    "return",
    "row",
    "runtime",
    "scalar",
    "scan",
    "schema",
    "search",
    "segment",
    "select",
    "session",
    "setting",
    "source",
    "stable",
    "stage",
    "state",
    "static",
    "storage",
    "stream",
    "string",
    "struct",
    "table",
    "target",
    "task",
    "thread",
    "token",
    "trace",
    "transform",
    "type",
    "update",
    "upload",
    "value",
    "vector",
    "version",
    "view",
    "write",
    "writer",
];

/// Word list with stop words at the front for Zipf sampling, computed once.
static SENTENCE_WORDS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
    let mut words = Vec::with_capacity(STOP_WORDS.len() + ENGLISH_WORDS.len());
    words.extend(STOP_WORDS.iter().copied());
    words.extend(ENGLISH_WORDS.iter().copied());
    words
});

struct RandomSentenceGenerator {
    min_words: usize,
    max_words: usize,
    /// Zipf distribution for word selection (favors lower indices)
    zipf: Zipf<f64>,
    is_large: bool,
}

impl std::fmt::Debug for RandomSentenceGenerator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RandomSentenceGenerator")
            .field("min_words", &self.min_words)
            .field("max_words", &self.max_words)
            .field("num_words", &SENTENCE_WORDS.len())
            .field("is_large", &self.is_large)
            .finish()
    }
}

impl RandomSentenceGenerator {
    pub fn new(min_words: usize, max_words: usize, is_large: bool) -> Self {
        // Zipf distribution with exponent ~1.0 approximates natural language
        let zipf = Zipf::new(SENTENCE_WORDS.len() as f64, 1.0).unwrap();

        Self {
            min_words,
            max_words,
            zipf,
            is_large,
        }
    }
}

impl ArrayGenerator for RandomSentenceGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let mut values = Vec::with_capacity(length.0 as usize);

        for _ in 0..length.0 {
            let num_words = rng.random_range(self.min_words..=self.max_words);
            let sentence: String = (0..num_words)
                .map(|_| {
                    // Zipf returns 1-indexed values, subtract 1 for 0-indexed array
                    let idx = rng.sample(self.zipf) as usize - 1;
                    SENTENCE_WORDS[idx]
                })
                .collect::<Vec<_>>()
                .join(" ");
            values.push(sentence);
        }

        if self.is_large {
            Ok(Arc::new(LargeStringArray::from(values)))
        } else {
            Ok(Arc::new(StringArray::from(values)))
        }
    }

    fn data_type(&self) -> &DataType {
        if self.is_large {
            &DataType::LargeUtf8
        } else {
            &DataType::Utf8
        }
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // Estimate average word length as 5, plus space
        // See https://arxiv.org/pdf/1208.6109
        let avg_word_length = 6;
        let avg_words = (self.min_words + self.max_words) / 2;
        Some(ByteCount::from((avg_word_length * avg_words) as u64))
    }
}

#[derive(Debug)]
struct RandomWordGenerator {
    words: &'static [&'static str],
    is_large: bool,
}

impl RandomWordGenerator {
    pub fn new(is_large: bool) -> Self {
        let words = ENGLISH_WORDS;
        Self { words, is_large }
    }
}

impl ArrayGenerator for RandomWordGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let mut values = Vec::with_capacity(length.0 as usize);

        for _ in 0..length.0 {
            let word = self.words[rng.random_range(0..self.words.len())];
            values.push(word.to_string());
        }

        if self.is_large {
            Ok(Arc::new(LargeStringArray::from(values)))
        } else {
            Ok(Arc::new(StringArray::from(values)))
        }
    }

    fn data_type(&self) -> &DataType {
        if self.is_large {
            &DataType::LargeUtf8
        } else {
            &DataType::Utf8
        }
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // Average English word length is ~5 characters
        Some(ByteCount::from(5))
    }
}

#[derive(Debug)]
pub struct VariableRandomBinaryGenerator {
    lengths_gen: Box<dyn ArrayGenerator>,
    data_type: DataType,
}

impl VariableRandomBinaryGenerator {
    pub fn new(min_bytes_per_element: ByteCount, max_bytes_per_element: ByteCount) -> Self {
        let lengths_dist = Uniform::new_inclusive(
            min_bytes_per_element.0 as i32,
            max_bytes_per_element.0 as i32,
        )
        .unwrap();
        let lengths_gen = rand_with_distribution::<Int32Type, Uniform<i32>>(lengths_dist);

        Self {
            lengths_gen,
            data_type: DataType::Binary,
        }
    }
}

impl ArrayGenerator for VariableRandomBinaryGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let lengths = self.lengths_gen.generate(length, rng)?;
        let lengths = lengths.as_primitive::<Int32Type>();
        let total_length = lengths.values().iter().map(|i| *i as usize).sum::<usize>();
        let offsets = OffsetBuffer::from_lengths(lengths.values().iter().map(|v| *v as usize));
        let mut bytes = vec![0; total_length];
        rng.fill_bytes(&mut bytes);
        let bytes = Buffer::from(bytes);
        Ok(Arc::new(BinaryArray::try_new(offsets, bytes, None)?))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        None
    }
}

pub struct CycleBinaryGenerator<T: ByteArrayType> {
    values: Vec<u8>,
    lengths: Vec<usize>,
    data_type: DataType,
    array_type: PhantomData<T>,
    width: Option<ByteCount>,
    idx: usize,
}

impl<T: ByteArrayType> std::fmt::Debug for CycleBinaryGenerator<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CycleBinaryGenerator")
            .field("values", &self.values)
            .field("lengths", &self.lengths)
            .field("data_type", &self.data_type)
            .field("width", &self.width)
            .field("idx", &self.idx)
            .finish()
    }
}

impl<T: ByteArrayType> CycleBinaryGenerator<T> {
    pub fn from_strings(values: &[&str]) -> Self {
        if values.is_empty() {
            panic!("Attempt to create a cycle generator with no values");
        }
        let lengths = values.iter().map(|s| s.len()).collect::<Vec<_>>();
        let typical_length = lengths[0];
        let width = if lengths.iter().all(|item| *item == typical_length) {
            Some(ByteCount::from(
                typical_length as u64 + std::mem::size_of::<i32>() as u64,
            ))
        } else {
            None
        };
        let values = values
            .iter()
            .flat_map(|s| s.as_bytes().iter().copied())
            .collect::<Vec<_>>();
        Self {
            values,
            lengths,
            data_type: T::DATA_TYPE,
            array_type: PhantomData,
            width,
            idx: 0,
        }
    }
}

impl<T: ByteArrayType> ArrayGenerator for CycleBinaryGenerator<T> {
    fn generate(
        &mut self,
        length: RowCount,
        _: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let lengths = self
            .lengths
            .iter()
            .copied()
            .cycle()
            .skip(self.idx)
            .take(length.0 as usize);
        let num_bytes = lengths.clone().sum();
        let byte_offset = self.lengths[0..self.idx].iter().sum();
        let bytes = self
            .values
            .iter()
            .cycle()
            .skip(byte_offset)
            .copied()
            .take(num_bytes)
            .collect::<Vec<_>>();
        let bytes = Buffer::from(bytes);
        let offsets = OffsetBuffer::from_lengths(lengths);
        self.idx = (self.idx + length.0 as usize) % self.lengths.len();
        Ok(Arc::new(arrow_array::GenericByteArray::<T>::new(
            offsets, bytes, None,
        )))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.width
    }
}

pub struct FixedBinaryGenerator<T: ByteArrayType> {
    value: Vec<u8>,
    data_type: DataType,
    array_type: PhantomData<T>,
}

impl<T: ByteArrayType> std::fmt::Debug for FixedBinaryGenerator<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FixedBinaryGenerator")
            .field("value", &self.value)
            .field("data_type", &self.data_type)
            .finish()
    }
}

impl<T: ByteArrayType> FixedBinaryGenerator<T> {
    pub fn new(value: Vec<u8>) -> Self {
        Self {
            value,
            data_type: T::DATA_TYPE,
            array_type: PhantomData,
        }
    }
}

impl<T: ByteArrayType> ArrayGenerator for FixedBinaryGenerator<T> {
    fn generate(
        &mut self,
        length: RowCount,
        _: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let bytes = Buffer::from(Vec::from_iter(
            self.value
                .iter()
                .cycle()
                .take((length.0 * self.value.len() as u64) as usize)
                .copied(),
        ));
        let offsets =
            OffsetBuffer::from_lengths(iter::repeat_n(self.value.len(), length.0 as usize));
        Ok(Arc::new(arrow_array::GenericByteArray::<T>::new(
            offsets, bytes, None,
        )))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        // Not exactly correct since there are N + 1 4-byte offsets and this only counts N
        Some(ByteCount::from(
            self.value.len() as u64 + std::mem::size_of::<i32>() as u64,
        ))
    }
}

pub struct DictionaryGenerator<K: ArrowDictionaryKeyType> {
    generator: Box<dyn ArrayGenerator>,
    data_type: DataType,
    key_type: PhantomData<K>,
    key_width: u64,
}

impl<K: ArrowDictionaryKeyType> std::fmt::Debug for DictionaryGenerator<K> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DictionaryGenerator")
            .field("generator", &self.generator)
            .field("data_type", &self.data_type)
            .field("key_width", &self.key_width)
            .finish()
    }
}

impl<K: ArrowDictionaryKeyType> DictionaryGenerator<K> {
    fn new(generator: Box<dyn ArrayGenerator>) -> Self {
        let key_type = Box::new(K::DATA_TYPE);
        let key_width = key_type
            .primitive_width()
            .expect("dictionary key types should have a known width")
            as u64;
        let val_type = Box::new(generator.data_type().clone());
        let dict_type = DataType::Dictionary(key_type, val_type);
        Self {
            generator,
            data_type: dict_type,
            key_type: PhantomData,
            key_width,
        }
    }
}

impl<K: ArrowDictionaryKeyType + Send + Sync> ArrayGenerator for DictionaryGenerator<K> {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let underlying = self.generator.generate(length, rng)?;
        arrow_cast::cast::cast(&underlying, &self.data_type)
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.generator
            .element_size_bytes()
            .map(|size_bytes| ByteCount::from(size_bytes.0 + self.key_width))
    }
}

/// Generator that produces low-cardinality data by generating a fixed set of
/// unique values and then randomly selecting from them.
struct LowCardinalityGenerator {
    inner: Box<dyn ArrayGenerator>,
    cardinality: usize,
    /// Cached unique values, generated on first call
    unique_values: Option<Arc<dyn Array>>,
}

impl std::fmt::Debug for LowCardinalityGenerator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LowCardinalityGenerator")
            .field("inner", &self.inner)
            .field("cardinality", &self.cardinality)
            .field("initialized", &self.unique_values.is_some())
            .finish()
    }
}

impl LowCardinalityGenerator {
    fn new(inner: Box<dyn ArrayGenerator>, cardinality: usize) -> Self {
        Self {
            inner,
            cardinality,
            unique_values: None,
        }
    }
}

impl ArrayGenerator for LowCardinalityGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        // Generate unique values on first call
        if self.unique_values.is_none() {
            self.unique_values = Some(
                self.inner
                    .generate(RowCount::from(self.cardinality as u64), rng)?,
            );
        }

        let unique_values = self.unique_values.as_ref().unwrap();

        // Generate random indices into the unique values
        let indices: Vec<usize> = (0..length.0)
            .map(|_| rng.random_range(0..self.cardinality))
            .collect();

        // Use arrow's take to select values
        let indices_array =
            arrow_array::UInt32Array::from(indices.iter().map(|&i| i as u32).collect::<Vec<_>>());
        arrow::compute::take(unique_values.as_ref(), &indices_array, None)
            .map(|arr| arr as Arc<dyn Array>)
    }

    fn data_type(&self) -> &DataType {
        self.inner.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        self.inner.element_size_bytes()
    }
}

#[derive(Debug)]
struct RandomListGenerator {
    field: Arc<Field>,
    child_field: Arc<Field>,
    items_gen: Box<dyn ArrayGenerator>,
    lengths_gen: Box<dyn ArrayGenerator>,
    is_large: bool,
}

impl RandomListGenerator {
    // Creates a list generator that generates random lists with lengths between 0 and 10 (inclusive)
    fn new(items_gen: Box<dyn ArrayGenerator>, is_large: bool) -> Self {
        let child_field = Arc::new(Field::new("item", items_gen.data_type().clone(), true));
        let list_type = if is_large {
            DataType::LargeList(child_field.clone())
        } else {
            DataType::List(child_field.clone())
        };
        let field = Field::new("", list_type, true);
        let lengths_gen = if is_large {
            let lengths_dist = Uniform::new_inclusive(0, 10).unwrap();
            rand_with_distribution::<Int64Type, Uniform<i64>>(lengths_dist)
        } else {
            let lengths_dist = Uniform::new_inclusive(0, 10).unwrap();
            rand_with_distribution::<Int32Type, Uniform<i32>>(lengths_dist)
        };
        Self {
            field: Arc::new(field),
            child_field,
            items_gen,
            lengths_gen,
            is_large,
        }
    }
}

impl ArrayGenerator for RandomListGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let lengths = self.lengths_gen.generate(length, rng)?;
        if self.is_large {
            let lengths = lengths.as_primitive::<Int64Type>();
            let total_length = lengths.values().iter().sum::<i64>() as u64;
            let offsets = OffsetBuffer::from_lengths(lengths.values().iter().map(|v| *v as usize));
            let items = self.items_gen.generate(RowCount::from(total_length), rng)?;
            Ok(Arc::new(LargeListArray::try_new(
                self.child_field.clone(),
                offsets,
                items,
                None,
            )?))
        } else {
            let lengths = lengths.as_primitive::<Int32Type>();
            let total_length = lengths.values().iter().sum::<i32>() as u64;
            let offsets = OffsetBuffer::from_lengths(lengths.values().iter().map(|v| *v as usize));
            let items = self.items_gen.generate(RowCount::from(total_length), rng)?;
            Ok(Arc::new(ListArray::try_new(
                self.child_field.clone(),
                offsets,
                items,
                None,
            )?))
        }
    }

    fn data_type(&self) -> &DataType {
        self.field.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        None
    }
}

/// Generates random map arrays where each map has 0-4 entries.
#[derive(Debug)]
struct RandomMapGenerator {
    field: Arc<Field>,
    entries_field: Arc<Field>,
    keys_gen: Box<dyn ArrayGenerator>,
    values_gen: Box<dyn ArrayGenerator>,
    lengths_gen: Box<dyn ArrayGenerator>,
}

impl RandomMapGenerator {
    fn new(keys_gen: Box<dyn ArrayGenerator>, values_gen: Box<dyn ArrayGenerator>) -> Self {
        let entries_fields = Fields::from(vec![
            Field::new("keys", keys_gen.data_type().clone(), false),
            Field::new("values", values_gen.data_type().clone(), true),
        ]);
        let entries_field = Arc::new(Field::new(
            "entries",
            DataType::Struct(entries_fields),
            false,
        ));
        let map_type = DataType::Map(entries_field.clone(), false);
        let field = Arc::new(Field::new("", map_type, true));
        let lengths_dist = Uniform::new_inclusive(0_i32, 4).unwrap();
        let lengths_gen = rand_with_distribution::<Int32Type, Uniform<i32>>(lengths_dist);

        Self {
            field,
            entries_field,
            keys_gen,
            values_gen,
            lengths_gen,
        }
    }
}

impl ArrayGenerator for RandomMapGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let lengths = self.lengths_gen.generate(length, rng)?;
        let lengths = lengths.as_primitive::<Int32Type>();
        let total_entries = lengths.values().iter().sum::<i32>() as u64;
        let offsets = OffsetBuffer::from_lengths(lengths.values().iter().map(|v| *v as usize));

        let keys = self.keys_gen.generate(RowCount::from(total_entries), rng)?;
        let values = self
            .values_gen
            .generate(RowCount::from(total_entries), rng)?;

        let entries = StructArray::new(
            Fields::from(vec![
                Field::new("keys", keys.data_type().clone(), false),
                Field::new("values", values.data_type().clone(), true),
            ]),
            vec![keys, values],
            None,
        );

        Ok(Arc::new(MapArray::try_new(
            self.entries_field.clone(),
            offsets,
            entries,
            None,
            false,
        )?))
    }

    fn data_type(&self) -> &DataType {
        self.field.data_type()
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        None
    }
}

#[derive(Debug)]
struct NullArrayGenerator {}

impl ArrayGenerator for NullArrayGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        _: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        Ok(Arc::new(NullArray::new(length.0 as usize)))
    }

    fn data_type(&self) -> &DataType {
        &DataType::Null
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        None
    }
}

/// Generates 2 dimensional vectors along the unit circle, with a configurable number of steps per circle.
#[derive(Debug)]
struct RadialStepGenerator {
    num_steps_per_circle: u32,
    data_field: Arc<Field>,
    data_type: DataType,
    current_step: u32,
}

impl RadialStepGenerator {
    fn new(num_steps_per_circle: u32) -> Self {
        let data_field = Arc::new(Field::new("item", DataType::Float32, false));
        let data_type = DataType::FixedSizeList(data_field.clone(), 2);
        Self {
            num_steps_per_circle,
            data_field,
            data_type,
            current_step: 0,
        }
    }
}

impl ArrayGenerator for RadialStepGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let mut values_builder = Float32Builder::with_capacity(length.0 as usize * 2);
        for _ in 0..length.0 {
            let angle = (self.current_step as f32) / (self.num_steps_per_circle as f32)
                * 2.0
                * std::f32::consts::PI;
            values_builder.append_value(angle.cos());
            values_builder.append_value(angle.sin());
            self.current_step = (self.current_step + 1) % self.num_steps_per_circle;
        }
        let values = values_builder.finish();
        let vectors =
            FixedSizeListArray::try_new(self.data_field.clone(), 2, Arc::new(values), None)?;
        Ok(Arc::new(vectors))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(8))
    }
}

/// Cycles through a set of centroids, adding noise to each point
#[derive(Debug)]
struct JitterCentroidsGenerator {
    centroids: Float32Array,
    dimension: u32,
    noise_level: f32,
    data_type: DataType,
    data_field: Arc<Field>,

    offset: usize,
}

impl JitterCentroidsGenerator {
    fn try_new(centroids: Arc<dyn Array>, noise_level: f32) -> Result<Self, ArrowError> {
        let DataType::FixedSizeList(values_field, dimension) = centroids.data_type() else {
            return Err(ArrowError::InvalidArgumentError(
                "Centroids must be a FixedSizeList".to_string(),
            ));
        };
        if values_field.data_type() != &DataType::Float32 {
            return Err(ArrowError::InvalidArgumentError(
                "Centroids values must be a Float32".to_string(),
            ));
        }
        let data_type = DataType::FixedSizeList(values_field.clone(), *dimension);
        Ok(Self {
            centroids: centroids
                .as_fixed_size_list()
                .values()
                .as_primitive::<Float32Type>()
                .clone(),
            dimension: *dimension as u32,
            noise_level,
            data_type,
            data_field: values_field.clone(),
            offset: 0,
        })
    }
}

impl ArrayGenerator for JitterCentroidsGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn Array>, ArrowError> {
        let mut values_builder =
            Float32Builder::with_capacity(length.0 as usize * self.dimension as usize);
        for _ in 0..length.0 {
            // Generate random N dimensional point
            let mut noise = (0..self.dimension as usize)
                .map(|_| rng.random::<f32>())
                .collect::<Vec<_>>();
            // Scale point to noise_level length
            let scale = self.noise_level / noise.iter().map(|v| v * v).sum::<f32>().sqrt();
            noise.iter_mut().for_each(|v| *v *= scale);

            // Add noise to centroid and store in values
            for (i, noise) in noise.into_iter().enumerate() {
                let centroid_val = self.centroids.value(self.offset + i);
                let jittered_val = centroid_val + noise;
                values_builder.append_value(jittered_val);
            }
            // Advance to next centroid
            self.offset = (self.offset + self.dimension as usize) % self.centroids.len();
        }
        let values = values_builder.finish();
        let vectors = FixedSizeListArray::try_new(
            self.data_field.clone(),
            self.dimension as i32,
            Arc::new(values),
            None,
        )?;
        Ok(Arc::new(vectors))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        Some(ByteCount::from(self.dimension as u64 * 4))
    }
}
#[derive(Debug)]
struct RandomStructGenerator {
    fields: Fields,
    data_type: DataType,
    child_gens: Vec<Box<dyn ArrayGenerator>>,
}

impl RandomStructGenerator {
    fn new(fields: Fields, child_gens: Vec<Box<dyn ArrayGenerator>>) -> Self {
        let data_type = DataType::Struct(fields.clone());
        Self {
            fields,
            data_type,
            child_gens,
        }
    }
}

impl ArrayGenerator for RandomStructGenerator {
    fn generate(
        &mut self,
        length: RowCount,
        rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        if self.child_gens.is_empty() {
            // Have to create empty struct arrays specially to ensure they have the correct
            // row count
            let struct_arr = StructArray::new_empty_fields(length.0 as usize, None);
            return Ok(Arc::new(struct_arr));
        }
        let child_arrays = self
            .child_gens
            .iter_mut()
            .map(|genn| genn.generate(length, rng))
            .collect::<Result<Vec<_>, ArrowError>>()?;
        let struct_arr = StructArray::new(self.fields.clone(), child_arrays, None);
        Ok(Arc::new(struct_arr))
    }

    fn data_type(&self) -> &DataType {
        &self.data_type
    }

    fn element_size_bytes(&self) -> Option<ByteCount> {
        let mut sum = 0;
        for child_gen in &self.child_gens {
            sum += child_gen.element_size_bytes()?.0;
        }
        Some(ByteCount::from(sum))
    }
}

/// A RecordBatchReader that generates batches of the given size from the given array generators
pub struct FixedSizeBatchGenerator {
    rng: rand_xoshiro::Xoshiro256PlusPlus,
    generators: Vec<Box<dyn ArrayGenerator>>,
    batch_size: RowCount,
    num_batches: BatchCount,
    schema: SchemaRef,
}

impl FixedSizeBatchGenerator {
    fn new(
        generators: Vec<(Option<String>, Box<dyn ArrayGenerator>)>,
        batch_size: RowCount,
        num_batches: BatchCount,
        seed: Option<Seed>,
        default_null_probability: Option<f64>,
    ) -> Self {
        let mut fields = Vec::with_capacity(generators.len());
        for (field_index, field_gen) in generators.iter().enumerate() {
            let (name, genn) = field_gen;
            let default_name = format!("field_{}", field_index);
            let name = name.clone().unwrap_or(default_name);
            let mut field = Field::new(name, genn.data_type().clone(), true);
            if let Some(metadata) = genn.metadata() {
                field = field.with_metadata(metadata);
            }
            fields.push(field);
        }
        let mut generators = generators
            .into_iter()
            .map(|(_, genn)| genn)
            .collect::<Vec<_>>();
        if let Some(null_probability) = default_null_probability {
            generators = generators
                .into_iter()
                .map(|genn| genn.with_random_nulls(null_probability))
                .collect();
        }
        let schema = Arc::new(Schema::new(fields));
        Self {
            rng: rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(
                seed.map(|s| s.0).unwrap_or(DEFAULT_SEED.0),
            ),
            generators,
            batch_size,
            num_batches,
            schema,
        }
    }

    fn gen_next(&mut self) -> Result<RecordBatch, ArrowError> {
        let mut arrays = Vec::with_capacity(self.generators.len());
        for genn in self.generators.iter_mut() {
            let arr = genn.generate(self.batch_size, &mut self.rng)?;
            arrays.push(arr);
        }
        self.num_batches.0 -= 1;
        Ok(RecordBatch::try_new_with_options(
            self.schema.clone(),
            arrays,
            &RecordBatchOptions::new().with_row_count(Some(self.batch_size.0 as usize)),
        )
        .unwrap())
    }
}

impl Iterator for FixedSizeBatchGenerator {
    type Item = Result<RecordBatch, ArrowError>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.num_batches.0 == 0 {
            return None;
        }
        Some(self.gen_next())
    }
}

impl RecordBatchReader for FixedSizeBatchGenerator {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

/// A builder to create a record batch reader with generated data
///
/// This type is meant to be used in a fluent builder style to define the schema and generators
/// for a record batch reader.
#[derive(Default)]
pub struct BatchGeneratorBuilder {
    generators: Vec<(Option<String>, Box<dyn ArrayGenerator>)>,
    default_null_probability: Option<f64>,
    seed: Option<Seed>,
}

pub enum RoundingBehavior {
    ExactOrErr,
    RoundUp,
    RoundDown,
}

impl BatchGeneratorBuilder {
    /// Create a new BatchGeneratorBuilder with a default random seed
    pub fn new() -> Self {
        Default::default()
    }

    /// Create a new BatchGeneratorBuilder with the given seed
    pub fn new_with_seed(seed: Seed) -> Self {
        Self {
            seed: Some(seed),
            ..Default::default()
        }
    }

    /// Adds a new column to the generator
    ///
    /// See [`crate::generator::array`] for methods to create generators
    pub fn col(mut self, name: impl Into<String>, genn: Box<dyn ArrayGenerator>) -> Self {
        self.generators.push((Some(name.into()), genn));
        self
    }

    /// Adds a new column to the generator with a generated unique name
    ///
    /// See [`crate::generator::array`] for methods to create generators
    pub fn anon_col(mut self, genn: Box<dyn ArrayGenerator>) -> Self {
        self.generators.push((None, genn));
        self
    }

    pub fn into_batch_rows(self, batch_size: RowCount) -> Result<RecordBatch, ArrowError> {
        let mut reader = self.into_reader_rows(batch_size, BatchCount::from(1));
        reader
            .next()
            .expect("Asked for 1 batch but reader was empty")
    }

    pub fn into_batch_bytes(
        self,
        batch_size: ByteCount,
        rounding: RoundingBehavior,
    ) -> Result<RecordBatch, ArrowError> {
        let mut reader = self.into_reader_bytes(batch_size, BatchCount::from(1), rounding)?;
        reader
            .next()
            .expect("Asked for 1 batch but reader was empty")
    }

    /// Create a RecordBatchReader that generates batches of the given size (in rows)
    pub fn into_reader_rows(
        self,
        batch_size: RowCount,
        num_batches: BatchCount,
    ) -> impl RecordBatchReader {
        FixedSizeBatchGenerator::new(
            self.generators,
            batch_size,
            num_batches,
            self.seed,
            self.default_null_probability,
        )
    }

    pub fn into_reader_stream(
        self,
        batch_size: RowCount,
        num_batches: BatchCount,
    ) -> (
        BoxStream<'static, Result<RecordBatch, ArrowError>>,
        Arc<Schema>,
    ) {
        // TODO: this is pretty lazy and could be optimized
        let reader = self.into_reader_rows(batch_size, num_batches);
        let schema = reader.schema();
        let batches = reader.collect::<Vec<_>>();
        (futures::stream::iter(batches).boxed(), schema)
    }

    /// Create a RecordBatchReader that generates batches of the given size (in bytes)
    pub fn into_reader_bytes(
        self,
        batch_size_bytes: ByteCount,
        num_batches: BatchCount,
        rounding: RoundingBehavior,
    ) -> Result<impl RecordBatchReader, ArrowError> {
        let bytes_per_row = self
            .generators
            .iter()
            .map(|genn| genn.1.element_size_bytes().map(|byte_count| byte_count.0).ok_or(
                        ArrowError::NotYetImplemented("The function into_reader_bytes currently requires each array generator to have a fixed element size".to_string())
                )
            )
            .sum::<Result<u64, ArrowError>>()?;
        let mut num_rows = RowCount::from(batch_size_bytes.0 / bytes_per_row);
        if !batch_size_bytes.0.is_multiple_of(bytes_per_row) {
            match rounding {
                RoundingBehavior::ExactOrErr => {
                    return Err(ArrowError::NotYetImplemented(format!(
                        "Exact rounding requested but not possible.  Batch size requested {}, row size: {}",
                        batch_size_bytes.0, bytes_per_row
                    )));
                }
                RoundingBehavior::RoundUp => {
                    num_rows = RowCount::from(num_rows.0 + 1);
                }
                RoundingBehavior::RoundDown => (),
            }
        }
        Ok(self.into_reader_rows(num_rows, num_batches))
    }

    /// Set the seed for the generator
    pub fn with_seed(mut self, seed: Seed) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Adds nulls (with the given probability) to all columns
    pub fn with_random_nulls(&mut self, default_null_probability: f64) {
        self.default_null_probability = Some(default_null_probability);
    }
}

/// Factory for creating a single random array
pub struct ArrayGeneratorBuilder {
    generator: Box<dyn ArrayGenerator>,
    seed: Option<Seed>,
}

impl ArrayGeneratorBuilder {
    fn new(generator: Box<dyn ArrayGenerator>) -> Self {
        Self {
            generator,
            seed: None,
        }
    }

    /// Use the given seed for the generator
    pub fn with_seed(mut self, seed: Seed) -> Self {
        self.seed = Some(seed);
        self
    }

    /// Generate a single array with the given length
    pub fn into_array_rows(
        mut self,
        length: RowCount,
    ) -> Result<Arc<dyn arrow_array::Array>, ArrowError> {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(
            self.seed.map(|s| s.0).unwrap_or(DEFAULT_SEED.0),
        );
        self.generator.generate(length, &mut rng)
    }
}

const MS_PER_DAY: i64 = 86400000;

pub mod array {

    use arrow::datatypes::{Int8Type, Int16Type, Int64Type};
    use arrow_array::types::{
        Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType,
        DurationNanosecondType, DurationSecondType, Float16Type, Float32Type, Float64Type,
        UInt8Type, UInt16Type, UInt32Type, UInt64Type,
    };
    use arrow_array::{
        ArrowNativeTypeOp, BooleanArray, Date32Array, Date64Array, Time32MillisecondArray,
        Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
        TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
        TimestampSecondArray,
    };
    use arrow_schema::{IntervalUnit, TimeUnit};
    use chrono::Utc;
    use rand::prelude::Distribution;

    use super::*;

    /// Create a generator of vectors by continuously calling the given generator
    ///
    /// For example, given a step generator and a dimension of 3 this will generate vectors like
    /// [0, 1, 2], [3, 4, 5], [6, 7, 8], ...
    pub fn cycle_vec(
        generator: Box<dyn ArrayGenerator>,
        dimension: Dimension,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(CycleVectorGenerator::new(generator, dimension))
    }

    /// Create a generator of list vectors by continuously calling the given generator
    ///
    /// The lists will have lengths uniformly distributed between `min_list_size` (inclusive) and
    /// `max_list_size` (exclusive).
    pub fn cycle_vec_var(
        generator: Box<dyn ArrayGenerator>,
        min_list_size: Dimension,
        max_list_size: Dimension,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(CycleListGenerator::new(
            generator,
            min_list_size,
            max_list_size,
        ))
    }

    /// Create a generator of vectors around unit circle
    ///
    /// Vectors will be equally spaced around the unit circle so that there are num_steps
    /// vectors per circle.
    pub fn cycle_unit_circle(num_steps: u32) -> Box<dyn ArrayGenerator> {
        Box::new(RadialStepGenerator::new(num_steps))
    }

    /// Create a generator of vectors by cycling through a given set of vectors
    ///
    /// Each value will be spaced in slightly away from the previous value on a ball of radius jitter
    pub fn jitter_centroids(centroids: Arc<dyn Array>, jitter: f32) -> Box<dyn ArrayGenerator> {
        Box::new(JitterCentroidsGenerator::try_new(centroids, jitter).unwrap())
    }

    /// Create a generator from a vector of values
    ///
    /// If more rows are requested than the length of values then it will restart
    /// from the beginning of the vector.
    pub fn cycle<DataType>(values: Vec<DataType::Native>) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        DataType: ArrowPrimitiveType,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
    {
        let mut values_idx = 0;
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |_| {
                    let y = values[values_idx];
                    values_idx = (values_idx + 1) % values.len();
                    y
                },
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    /// Create a generator from a vector of booleans
    ///
    /// If more rows are requested than the length of values then it will restart from
    /// the beginning of the vector
    pub fn cycle_bool(values: Vec<bool>) -> Box<dyn ArrayGenerator> {
        let mut values_idx = 0;
        Box::new(FnGen::<bool, BooleanArray, _>::new_unknown_size(
            DataType::Boolean,
            move |_| {
                let val = values[values_idx];
                values_idx = (values_idx + 1) % values.len();
                val
            },
            1,
        ))
    }

    /// Create a generator that starts at 0 and increments by 1 for each element
    pub fn step<DataType>() -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + Default + std::ops::AddAssign<DataType::Native> + 'static,
        DataType: ArrowPrimitiveType,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
    {
        let mut x = DataType::Native::default();
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |_| {
                    let y = x;
                    x += DataType::Native::ONE;
                    y
                },
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    pub fn blob() -> Box<dyn ArrayGenerator> {
        let mut blob_meta = HashMap::new();
        blob_meta.insert("lance-encoding:blob".to_string(), "true".to_string());
        rand_fixedbin(ByteCount::from(4 * 1024 * 1024), true).with_metadata(blob_meta)
    }

    /// Create a generator that starts at a given value and increments by a given step for each element
    pub fn step_custom<DataType>(
        start: DataType::Native,
        step: DataType::Native,
    ) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + Default + std::ops::AddAssign<DataType::Native> + 'static,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
        DataType: ArrowPrimitiveType,
    {
        let mut x = start;
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |_| {
                    let y = x;
                    x += step;
                    y
                },
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    /// Create a generator that fills each element with the given primitive value
    pub fn fill<DataType>(value: DataType::Native) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        DataType: ArrowPrimitiveType,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
    {
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |_| value,
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    /// Create a generator that fills each element with the given binary value
    pub fn fill_varbin(value: Vec<u8>) -> Box<dyn ArrayGenerator> {
        Box::new(FixedBinaryGenerator::<BinaryType>::new(value))
    }

    /// Create a generator that fills each element with the given string value
    pub fn fill_utf8(value: String) -> Box<dyn ArrayGenerator> {
        Box::new(FixedBinaryGenerator::<Utf8Type>::new(value.into_bytes()))
    }

    pub fn cycle_utf8_literals(values: &[&'static str]) -> Box<dyn ArrayGenerator> {
        Box::new(CycleBinaryGenerator::<Utf8Type>::from_strings(values))
    }

    /// Create a generator of primitive values that are randomly sampled from the entire range available for the value
    pub fn rand<DataType>() -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
        DataType: ArrowPrimitiveType,
        rand::distr::StandardUniform: rand::distr::Distribution<DataType::Native>,
    {
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |rng| rng.random(),
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    /// Create a generator of primitive values that are randomly sampled from the entire range available for the value
    pub fn rand_with_distribution<
        DataType,
        Dist: rand::distr::Distribution<DataType::Native> + Clone + Send + Sync + 'static,
    >(
        dist: Dist,
    ) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
        DataType: ArrowPrimitiveType,
    {
        Box::new(
            FnGen::<DataType::Native, PrimitiveArray<DataType>, _>::new_known_size(
                DataType::DATA_TYPE,
                move |rng| rng.sample(dist.clone()),
                1,
                DataType::DATA_TYPE
                    .primitive_width()
                    .map(|width| ByteCount::from(width as u64))
                    .expect("Primitive types should have a fixed width"),
            ),
        )
    }

    /// Create a generator of 1d vectors (of a primitive type) consisting of randomly sampled primitive values
    pub fn rand_vec<DataType>(dimension: Dimension) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
        DataType: ArrowPrimitiveType,
        rand::distr::StandardUniform: rand::distr::Distribution<DataType::Native>,
    {
        let underlying = rand::<DataType>();
        cycle_vec(underlying, dimension)
    }

    /// Create a generator of 1d vectors (of a primitive type) consisting of randomly sampled nullable values
    pub fn rand_vec_nullable<DataType>(
        dimension: Dimension,
        null_probability: f64,
    ) -> Box<dyn ArrayGenerator>
    where
        DataType::Native: Copy + 'static,
        PrimitiveArray<DataType>: From<Vec<DataType::Native>> + 'static,
        DataType: ArrowPrimitiveType,
        rand::distr::StandardUniform: rand::distr::Distribution<DataType::Native>,
    {
        let underlying = rand::<DataType>().with_random_nulls(null_probability);
        cycle_vec(underlying, dimension)
    }

    /// Create a generator of randomly sampled time32 values covering the entire
    /// range of 1 day
    pub fn rand_time32(resolution: &TimeUnit) -> Box<dyn ArrayGenerator> {
        let start = 0;
        let end = match resolution {
            TimeUnit::Second => 86_400,
            TimeUnit::Millisecond => 86_400_000,
            _ => panic!(),
        };

        let data_type = DataType::Time32(*resolution);
        let size = ByteCount::from(data_type.primitive_width().unwrap() as u64);
        let dist = Uniform::new(start, end).unwrap();
        let sample_fn = move |rng: &mut _| dist.sample(rng);

        match resolution {
            TimeUnit::Second => Box::new(FnGen::<i32, Time32SecondArray, _>::new_known_size(
                data_type, sample_fn, 1, size,
            )),
            TimeUnit::Millisecond => {
                Box::new(FnGen::<i32, Time32MillisecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, size,
                ))
            }
            _ => panic!(),
        }
    }

    /// Create a generator of randomly sampled time64 values covering the entire
    /// range of 1 day
    pub fn rand_time64(resolution: &TimeUnit) -> Box<dyn ArrayGenerator> {
        let start = 0_i64;
        let end: i64 = match resolution {
            TimeUnit::Microsecond => 86_400_000,
            TimeUnit::Nanosecond => 86_400_000_000,
            _ => panic!(),
        };

        let data_type = DataType::Time64(*resolution);
        let size = ByteCount::from(data_type.primitive_width().unwrap() as u64);
        let dist = Uniform::new(start, end).unwrap();
        let sample_fn = move |rng: &mut _| dist.sample(rng);

        match resolution {
            TimeUnit::Microsecond => {
                Box::new(FnGen::<i64, Time64MicrosecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, size,
                ))
            }
            TimeUnit::Nanosecond => {
                Box::new(FnGen::<i64, Time64NanosecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, size,
                ))
            }
            _ => panic!(),
        }
    }

    /// Create a generator of random UUIDs, stored as fixed size binary values
    ///
    /// Note, these are "pseudo UUIDs".  They are 16-byte randomish values but they
    /// are not guaranteed to be unique.  We use a simplistic RNG that trades uniqueness
    /// for speed.
    pub fn rand_pseudo_uuid() -> Box<dyn ArrayGenerator> {
        Box::<PseudoUuidGenerator>::default()
    }

    /// Create a generator of random UUIDs, stored as 32-character strings (hex encoding
    /// of the 16-byte binary value)
    ///
    /// Note, these are "pseudo UUIDs".  They are 16-byte randomish values but they
    /// are not guaranteed to be unique.  We use a simplistic RNG that trades uniqueness
    /// for speed.
    pub fn rand_pseudo_uuid_hex() -> Box<dyn ArrayGenerator> {
        Box::<PseudoUuidHexGenerator>::default()
    }

    pub fn rand_primitive<T: ArrowPrimitiveType + Send + Sync>(
        data_type: DataType,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(RandomBytesGenerator::<T>::new(data_type))
    }

    pub fn rand_fsb(size: i32) -> Box<dyn ArrayGenerator> {
        Box::new(RandomFixedSizeBinaryGenerator::new(size))
    }

    pub fn rand_interval(unit: IntervalUnit) -> Box<dyn ArrayGenerator> {
        Box::new(RandomIntervalGenerator::new(unit))
    }

    /// Create a generator of randomly sampled date32 values
    ///
    /// Instead of sampling the entire range, all values will be drawn from the last year as this
    /// is a more common use pattern
    pub fn rand_date32() -> Box<dyn ArrayGenerator> {
        let now = chrono::Utc::now();
        let one_year_ago = now - chrono::TimeDelta::try_days(365).expect("TimeDelta try days");
        rand_date32_in_range(one_year_ago, now)
    }

    /// Create a generator of randomly sampled date32 values in the given range
    pub fn rand_date32_in_range(
        start: chrono::DateTime<Utc>,
        end: chrono::DateTime<Utc>,
    ) -> Box<dyn ArrayGenerator> {
        let data_type = DataType::Date32;
        let end_ms = end.timestamp_millis();
        let end_days = (end_ms / MS_PER_DAY) as i32;
        let start_ms = start.timestamp_millis();
        let start_days = (start_ms / MS_PER_DAY) as i32;
        let dist = Uniform::new(start_days, end_days).unwrap();

        Box::new(FnGen::<i32, Date32Array, _>::new_known_size(
            data_type,
            move |rng| dist.sample(rng),
            1,
            DataType::Date32
                .primitive_width()
                .map(|width| ByteCount::from(width as u64))
                .expect("Date32 should have a fixed width"),
        ))
    }

    /// Create a generator of randomly sampled date64 values
    ///
    /// Instead of sampling the entire range, all values will be drawn from the last year as this
    /// is a more common use pattern
    pub fn rand_date64() -> Box<dyn ArrayGenerator> {
        let now = chrono::Utc::now();
        let one_year_ago = now - chrono::TimeDelta::try_days(365).expect("TimeDelta try_days");
        rand_date64_in_range(one_year_ago, now)
    }

    /// Create a generator of randomly sampled timestamp values in the given range
    ///
    /// Currently just samples the entire range of u64 values and casts to timestamp
    pub fn rand_timestamp_in_range(
        start: chrono::DateTime<Utc>,
        end: chrono::DateTime<Utc>,
        data_type: &DataType,
    ) -> Box<dyn ArrayGenerator> {
        let end_ms = end.timestamp_millis();
        let start_ms = start.timestamp_millis();
        let (start_ticks, end_ticks) = match data_type {
            DataType::Timestamp(TimeUnit::Nanosecond, _) => {
                (start_ms * 1000 * 1000, end_ms * 1000 * 1000)
            }
            DataType::Timestamp(TimeUnit::Microsecond, _) => (start_ms * 1000, end_ms * 1000),
            DataType::Timestamp(TimeUnit::Millisecond, _) => (start_ms, end_ms),
            DataType::Timestamp(TimeUnit::Second, _) => (start.timestamp(), end.timestamp()),
            _ => panic!(),
        };
        let dist = Uniform::new(start_ticks, end_ticks).unwrap();

        let data_type = data_type.clone();
        let sample_fn = move |rng: &mut _| dist.sample(rng);
        let width = data_type
            .primitive_width()
            .map(|width| ByteCount::from(width as u64))
            .unwrap();

        match data_type {
            DataType::Timestamp(TimeUnit::Nanosecond, _) => {
                Box::new(FnGen::<i64, TimestampNanosecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, width,
                ))
            }
            DataType::Timestamp(TimeUnit::Microsecond, _) => {
                Box::new(FnGen::<i64, TimestampMicrosecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, width,
                ))
            }
            DataType::Timestamp(TimeUnit::Millisecond, _) => {
                Box::new(FnGen::<i64, TimestampMillisecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, width,
                ))
            }
            DataType::Timestamp(TimeUnit::Second, _) => {
                Box::new(FnGen::<i64, TimestampSecondArray, _>::new_known_size(
                    data_type, sample_fn, 1, width,
                ))
            }
            _ => panic!(),
        }
    }

    pub fn rand_timestamp(data_type: &DataType) -> Box<dyn ArrayGenerator> {
        let now = chrono::Utc::now();
        let one_year_ago = now - chrono::Duration::try_days(365).unwrap();
        rand_timestamp_in_range(one_year_ago, now, data_type)
    }

    /// Create a generator of randomly sampled date64 values
    ///
    /// Instead of sampling the entire range, all values will be drawn from the last year as this
    /// is a more common use pattern
    pub fn rand_date64_in_range(
        start: chrono::DateTime<Utc>,
        end: chrono::DateTime<Utc>,
    ) -> Box<dyn ArrayGenerator> {
        let data_type = DataType::Date64;
        let end_ms = end.timestamp_millis();
        let end_days = end_ms / MS_PER_DAY;
        let start_ms = start.timestamp_millis();
        let start_days = start_ms / MS_PER_DAY;
        let dist = Uniform::new(start_days, end_days).unwrap();

        Box::new(FnGen::<i64, Date64Array, _>::new_known_size(
            data_type,
            move |rng| (dist.sample(rng)) * MS_PER_DAY,
            1,
            DataType::Date64
                .primitive_width()
                .map(|width| ByteCount::from(width as u64))
                .expect("Date64 should have a fixed width"),
        ))
    }

    /// Create a generator of random binary values where each value has a fixed number of bytes
    pub fn rand_fixedbin(bytes_per_element: ByteCount, is_large: bool) -> Box<dyn ArrayGenerator> {
        Box::new(RandomBinaryGenerator::new(
            bytes_per_element,
            false,
            is_large,
        ))
    }

    /// Create a generator of random binary values where each value has a variable number of bytes
    ///
    /// The number of bytes per element will be randomly sampled from the given (inclusive) range
    pub fn rand_varbin(
        min_bytes_per_element: ByteCount,
        max_bytes_per_element: ByteCount,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(VariableRandomBinaryGenerator::new(
            min_bytes_per_element,
            max_bytes_per_element,
        ))
    }

    /// Create a generator of random strings
    ///
    /// All strings will consist entirely of printable ASCII characters
    pub fn rand_utf8(bytes_per_element: ByteCount, is_large: bool) -> Box<dyn ArrayGenerator> {
        Box::new(RandomBinaryGenerator::new(
            bytes_per_element,
            true,
            is_large,
        ))
    }

    /// Creates a generator of strings with a prefix and a counter
    ///
    /// For example, if the prefix is "user_" the strings will be "user_0", "user_1", ...
    pub fn utf8_prefix_plus_counter(
        prefix: impl Into<String>,
        is_large: bool,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(PrefixPlusCounterGenerator::new(prefix.into(), is_large))
    }

    pub fn binary_prefix_plus_counter(
        prefix: Arc<[u8]>,
        is_large: bool,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(BinaryPrefixPlusCounterGenerator::new(prefix, is_large))
    }

    /// Create a random generator of boolean values
    pub fn rand_boolean() -> Box<dyn ArrayGenerator> {
        Box::<RandomBooleanGenerator>::default()
    }

    /// Create a generator of random sentences
    ///
    /// Generates strings containing between min_words and max_words random English words joined by spaces
    pub fn random_sentence(
        min_words: usize,
        max_words: usize,
        is_large: bool,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(RandomSentenceGenerator::new(min_words, max_words, is_large))
    }

    /// Create a generator of random words (one word per row)
    ///
    /// Generates strings containing a single random English word per row
    pub fn random_word(is_large: bool) -> Box<dyn ArrayGenerator> {
        Box::new(RandomWordGenerator::new(is_large))
    }

    pub fn rand_list(item_type: &DataType, is_large: bool) -> Box<dyn ArrayGenerator> {
        let child_gen = rand_type(item_type);
        Box::new(RandomListGenerator::new(child_gen, is_large))
    }

    pub fn rand_list_any(
        item_gen: Box<dyn ArrayGenerator>,
        is_large: bool,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(RandomListGenerator::new(item_gen, is_large))
    }

    /// Generates random map arrays where each map has 0-4 entries.
    pub fn rand_map(key_type: &DataType, value_type: &DataType) -> Box<dyn ArrayGenerator> {
        let keys_gen = rand_type(key_type);
        let values_gen = rand_type(value_type);
        Box::new(RandomMapGenerator::new(keys_gen, values_gen))
    }

    pub fn rand_struct(fields: Fields) -> Box<dyn ArrayGenerator> {
        let child_gens = fields
            .iter()
            .map(|f| rand_type(f.data_type()))
            .collect::<Vec<_>>();
        Box::new(RandomStructGenerator::new(fields, child_gens))
    }

    pub fn null_type() -> Box<dyn ArrayGenerator> {
        Box::new(NullArrayGenerator {})
    }

    /// Create a generator of random values
    pub fn rand_type(data_type: &DataType) -> Box<dyn ArrayGenerator> {
        match data_type {
            DataType::Boolean => rand_boolean(),
            DataType::Int8 => rand::<Int8Type>(),
            DataType::Int16 => rand::<Int16Type>(),
            DataType::Int32 => rand::<Int32Type>(),
            DataType::Int64 => rand::<Int64Type>(),
            DataType::UInt8 => rand::<UInt8Type>(),
            DataType::UInt16 => rand::<UInt16Type>(),
            DataType::UInt32 => rand::<UInt32Type>(),
            DataType::UInt64 => rand::<UInt64Type>(),
            DataType::Float16 => rand_primitive::<Float16Type>(data_type.clone()),
            DataType::Float32 => rand::<Float32Type>(),
            DataType::Float64 => rand::<Float64Type>(),
            DataType::Decimal128(_, _) => rand_primitive::<Decimal128Type>(data_type.clone()),
            DataType::Decimal256(_, _) => rand_primitive::<Decimal256Type>(data_type.clone()),
            DataType::Utf8 => rand_utf8(ByteCount::from(12), false),
            DataType::LargeUtf8 => rand_utf8(ByteCount::from(12), true),
            DataType::Binary => rand_fixedbin(ByteCount::from(12), false),
            DataType::LargeBinary => rand_fixedbin(ByteCount::from(12), true),
            DataType::Dictionary(key_type, value_type) => {
                dict_type(rand_type(value_type), key_type)
            }
            DataType::FixedSizeList(child, dimension) => cycle_vec(
                rand_type(child.data_type()),
                Dimension::from(*dimension as u32),
            ),
            DataType::FixedSizeBinary(size) => rand_fsb(*size),
            DataType::List(child) => rand_list(child.data_type(), false),
            DataType::LargeList(child) => rand_list(child.data_type(), true),
            DataType::Map(entries_field, _) => {
                let DataType::Struct(fields) = entries_field.data_type() else {
                    panic!("Map entries field must be a struct");
                };
                let key_type = fields[0].data_type();
                let value_type = fields[1].data_type();
                rand_map(key_type, value_type)
            }
            DataType::Duration(unit) => match unit {
                TimeUnit::Second => rand::<DurationSecondType>(),
                TimeUnit::Millisecond => rand::<DurationMillisecondType>(),
                TimeUnit::Microsecond => rand::<DurationMicrosecondType>(),
                TimeUnit::Nanosecond => rand::<DurationNanosecondType>(),
            },
            DataType::Interval(unit) => rand_interval(*unit),
            DataType::Date32 => rand_date32(),
            DataType::Date64 => rand_date64(),
            DataType::Time32(resolution) => rand_time32(resolution),
            DataType::Time64(resolution) => rand_time64(resolution),
            DataType::Timestamp(_, _) => rand_timestamp(data_type),
            DataType::Struct(fields) => rand_struct(fields.clone()),
            DataType::Null => null_type(),
            _ => unimplemented!("random generation of {}", data_type),
        }
    }

    /// Encodes arrays generated by the underlying generator as dictionaries with the given key type
    ///
    /// Note that this may not be very realistic if the underlying generator is something like a random
    /// generator since most of the underlying values will be unique and the common case for dictionary
    /// encoding is when there is a small set of possible values.
    pub fn dict<K: ArrowDictionaryKeyType + Send + Sync>(
        generator: Box<dyn ArrayGenerator>,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(DictionaryGenerator::<K>::new(generator))
    }

    /// Encodes arrays generated by the underlying generator as dictionaries with the given key type
    pub fn dict_type(
        generator: Box<dyn ArrayGenerator>,
        key_type: &DataType,
    ) -> Box<dyn ArrayGenerator> {
        match key_type {
            DataType::Int8 => dict::<Int8Type>(generator),
            DataType::Int16 => dict::<Int16Type>(generator),
            DataType::Int32 => dict::<Int32Type>(generator),
            DataType::Int64 => dict::<Int64Type>(generator),
            DataType::UInt8 => dict::<UInt8Type>(generator),
            DataType::UInt16 => dict::<UInt16Type>(generator),
            DataType::UInt32 => dict::<UInt32Type>(generator),
            DataType::UInt64 => dict::<UInt64Type>(generator),
            _ => unimplemented!(),
        }
    }

    /// Wraps a generator to produce low-cardinality data.
    ///
    /// Generates `cardinality` unique values on first call, then randomly
    /// selects from them for all subsequent rows.
    pub fn low_cardinality(
        generator: Box<dyn ArrayGenerator>,
        cardinality: usize,
    ) -> Box<dyn ArrayGenerator> {
        Box::new(LowCardinalityGenerator::new(generator, cardinality))
    }
}

/// Create a BatchGeneratorBuilder to start generating batch data
pub fn gen_batch() -> BatchGeneratorBuilder {
    BatchGeneratorBuilder::default()
}

/// Create an ArrayGeneratorBuilder to start generating array data
pub fn gen_array(genn: Box<dyn ArrayGenerator>) -> ArrayGeneratorBuilder {
    ArrayGeneratorBuilder::new(genn)
}

/// Metadata key to specify content type for string generation.
/// Set to "sentence" to use the sentence generator with Zipf distribution.
pub const CONTENT_TYPE_KEY: &str = "lance-datagen:content-type";

/// Metadata key to specify cardinality for low-cardinality data generation.
/// Set to a numeric string (e.g., "100") to limit unique values.
pub const CARDINALITY_KEY: &str = "lance-datagen:cardinality";

/// Create a generator for a field, checking metadata for content type hints.
///
/// Supported metadata keys:
/// - `lance-datagen:content-type`: Set to "sentence" for Utf8/LargeUtf8 fields
///   to use the sentence generator with Zipf distribution.
/// - `lance-datagen:cardinality`: Set to a number to limit unique values.
///   The generator will produce only that many unique values and randomly
///   select from them.
pub fn rand_field(field: &Field) -> Box<dyn ArrayGenerator> {
    let mut generator = if let Some(content_type) = field.metadata().get(CONTENT_TYPE_KEY) {
        match (content_type.as_str(), field.data_type()) {
            ("sentence", DataType::Utf8) => array::random_sentence(1, 10, false),
            ("sentence", DataType::LargeUtf8) => array::random_sentence(1, 10, true),
            _ => array::rand_type(field.data_type()),
        }
    } else {
        array::rand_type(field.data_type())
    };

    if let Some(cardinality_str) = field.metadata().get(CARDINALITY_KEY)
        && let Ok(cardinality) = cardinality_str.parse::<usize>()
        && cardinality > 0
    {
        generator = array::low_cardinality(generator, cardinality);
    }

    generator
}

/// Create a BatchGeneratorBuilder with the given schema
///
/// You can add more columns or convert this into a reader immediately.
///
/// Supported field metadata:
/// - `lance-datagen:content-type` = `"sentence"`: Use sentence generator with
///   Zipf distribution for more realistic text (Utf8/LargeUtf8 only).
/// - `lance-datagen:cardinality` = `"<number>"`: Limit to N unique values.
pub fn rand(schema: &Schema) -> BatchGeneratorBuilder {
    let mut builder = BatchGeneratorBuilder::default();
    for field in schema.fields() {
        builder = builder.col(field.name(), rand_field(field));
    }
    builder
}

#[cfg(test)]
mod tests {

    use arrow::datatypes::{Float32Type, Int8Type, Int16Type, UInt32Type};
    use arrow_array::{BooleanArray, Float32Array, Int8Array, Int16Array, Int32Array, UInt32Array};

    use super::*;

    #[test]
    fn test_step() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::step::<Int32Type>();
        assert_eq!(
            *genn.generate(RowCount::from(5), &mut rng).unwrap(),
            Int32Array::from_iter([0, 1, 2, 3, 4])
        );
        assert_eq!(
            *genn.generate(RowCount::from(5), &mut rng).unwrap(),
            Int32Array::from_iter([5, 6, 7, 8, 9])
        );

        let mut genn = array::step::<Int8Type>();
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Int8Array::from_iter([0, 1, 2])
        );

        let mut genn = array::step::<Float32Type>();
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Float32Array::from_iter([0.0, 1.0, 2.0])
        );

        let mut genn = array::step_custom::<Int16Type>(4, 8);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Int16Array::from_iter([4, 12, 20])
        );
        assert_eq!(
            *genn.generate(RowCount::from(2), &mut rng).unwrap(),
            Int16Array::from_iter([28, 36])
        );
    }

    #[test]
    fn test_cycle() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::cycle::<Int32Type>(vec![1, 2, 3]);
        assert_eq!(
            *genn.generate(RowCount::from(5), &mut rng).unwrap(),
            Int32Array::from_iter([1, 2, 3, 1, 2])
        );

        let mut genn = array::cycle_utf8_literals(&["abc", "def", "xyz"]);
        assert_eq!(
            *genn.generate(RowCount::from(5), &mut rng).unwrap(),
            StringArray::from_iter_values(["abc", "def", "xyz", "abc", "def"])
        );
        assert_eq!(
            *genn.generate(RowCount::from(1), &mut rng).unwrap(),
            StringArray::from_iter_values(["xyz"])
        );

        let mut genn = array::cycle_bool(vec![false, false, true]);
        assert_eq!(
            *genn.generate(RowCount::from(5), &mut rng).unwrap(),
            BooleanArray::from_iter(vec![false, false, true, false, false].into_iter().map(Some))
        );
        assert_eq!(
            *genn.generate(RowCount::from(1), &mut rng).unwrap(),
            BooleanArray::from_iter(vec![Some(true)])
        )
    }

    #[test]
    fn test_fill() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::fill::<Int32Type>(42);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Int32Array::from_iter([42, 42, 42])
        );
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Int32Array::from_iter([42, 42, 42])
        );

        let mut genn = array::fill_varbin(vec![0, 1, 2]);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::BinaryArray::from_iter_values([
                "\x00\x01\x02",
                "\x00\x01\x02",
                "\x00\x01\x02"
            ])
        );

        let mut genn = array::fill_utf8("xyz".to_string());
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::StringArray::from_iter_values(["xyz", "xyz", "xyz"])
        );
    }

    #[test]
    fn test_utf8_prefix_plus_counter() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::utf8_prefix_plus_counter("user_", false);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::StringArray::from_iter_values(["user_0", "user_1", "user_2"])
        );

        let mut genn = array::utf8_prefix_plus_counter("user_", true);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::LargeStringArray::from_iter_values(["user_0", "user_1", "user_2"])
        );
    }

    #[test]
    fn test_rng() {
        // Note: these tests are heavily dependent on the default seed.
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::rand::<Int32Type>();
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            Int32Array::from_iter([-797553329, 1369325940, -69174021])
        );

        let mut genn = array::rand_fixedbin(ByteCount::from(3), false);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::BinaryArray::from_iter_values([
                [184, 53, 216],
                [12, 96, 159],
                [125, 179, 56]
            ])
        );

        let mut genn = array::rand_utf8(ByteCount::from(3), false);
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::StringArray::from_iter_values([">@p", "n `", "NWa"])
        );

        let mut genn = array::random_sentence(1, 5, false);
        let words = genn.generate(RowCount::from(10), &mut rng).unwrap();
        assert_eq!(words.data_type(), &DataType::Utf8);
        let words_array = words.as_any().downcast_ref::<StringArray>().unwrap();
        // Verify each string contains 1-5 words
        for i in 0..10 {
            let sentence = words_array.value(i);
            let word_count = sentence.split_whitespace().count();
            assert!((1..=5).contains(&word_count));
        }

        let mut genn = array::rand_date32();
        let days_32 = genn.generate(RowCount::from(3), &mut rng).unwrap();
        assert_eq!(days_32.data_type(), &DataType::Date32);

        let mut genn = array::rand_date64();
        let days_64 = genn.generate(RowCount::from(3), &mut rng).unwrap();
        assert_eq!(days_64.data_type(), &DataType::Date64);

        let mut genn = array::rand_boolean();
        let bools = genn.generate(RowCount::from(1024), &mut rng).unwrap();
        assert_eq!(bools.data_type(), &DataType::Boolean);
        let bools = bools.as_any().downcast_ref::<BooleanArray>().unwrap();
        // Sanity check to ensure we're getting at least some rng
        assert!(bools.false_count() > 100);
        assert!(bools.true_count() > 100);

        let mut genn = array::rand_varbin(ByteCount::from(2), ByteCount::from(4));
        assert_eq!(
            *genn.generate(RowCount::from(3), &mut rng).unwrap(),
            arrow_array::BinaryArray::from_iter_values([
                vec![111, 9, 80],
                vec![86, 118, 13, 209],
                vec![68, 33, 202]
            ])
        );
    }

    #[test]
    fn test_rng_list() {
        // Note: these tests are heavily dependent on the default seed.
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::rand_list(&DataType::Int32, false);
        let arr = genn.generate(RowCount::from(100), &mut rng).unwrap();
        // Make sure we can generate empty lists (note, test is dependent on seed)
        let arr = arr.as_list::<i32>();
        assert!(arr.iter().any(|l| l.unwrap().is_empty()));
        // Shouldn't generate any giant lists (don't kill performance in normal datagen)
        assert!(arr.iter().any(|l| l.unwrap().len() < 11));
    }

    #[test]
    fn test_rng_distribution() {
        // Sanity test to make sure we our RNG is giving us well distributed values
        // We generates some 4-byte integers, histogram them into 8 buckets, and make
        // sure each bucket has a good # of values
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::rand::<UInt32Type>();
        for _ in 0..10 {
            let arr = genn.generate(RowCount::from(10000), &mut rng).unwrap();
            let int_arr = arr.as_any().downcast_ref::<UInt32Array>().unwrap();
            let mut buckets = vec![0_u32; 256];
            for val in int_arr.values() {
                buckets[(*val >> 24) as usize] += 1;
            }
            for bucket in buckets {
                // Perfectly even distribution would have 10000 / 256 values (~40) per bucket
                // We test for 15 which should be "good enough" and statistically unlikely to fail
                assert!(bucket > 15);
            }
        }
    }

    #[test]
    fn test_nulls() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::rand::<Int32Type>().with_random_nulls(0.3);

        let arr = genn.generate(RowCount::from(1000), &mut rng).unwrap();

        // This assert depends on the default seed
        assert_eq!(arr.null_count(), 297);

        for len in 0..100 {
            let arr = genn.generate(RowCount::from(len), &mut rng).unwrap();
            // Make sure the null count we came up with matches the actual # of unset bits
            assert_eq!(
                arr.null_count(),
                arr.nulls()
                    .map(|nulls| (len as usize)
                        - nulls.buffer().count_set_bits_offset(0, len as usize))
                    .unwrap_or(0)
            );
        }

        let mut genn = array::rand::<Int32Type>().with_random_nulls(0.0);
        let arr = genn.generate(RowCount::from(10), &mut rng).unwrap();

        assert_eq!(arr.null_count(), 0);

        let mut genn = array::rand::<Int32Type>().with_random_nulls(1.0);
        let arr = genn.generate(RowCount::from(10), &mut rng).unwrap();

        assert_eq!(arr.null_count(), 10);
        assert!((0..10).all(|idx| arr.is_null(idx)));

        let mut genn = array::rand::<Int32Type>().with_nulls(&[false, false, true]);
        let arr = genn.generate(RowCount::from(7), &mut rng).unwrap();
        assert!((0..2).all(|idx| arr.is_valid(idx)));
        assert!(arr.is_null(2));
        assert!((3..5).all(|idx| arr.is_valid(idx)));
        assert!(arr.is_null(5));
        assert!(arr.is_valid(6));
    }

    #[test]
    fn test_unit_circle() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut genn = array::cycle_unit_circle(4);
        let arr = genn.generate(RowCount::from(6), &mut rng).unwrap();

        let arr_values = arr
            .as_fixed_size_list()
            .values()
            .as_primitive::<Float32Type>()
            .values()
            .to_vec();
        assert_eq!(arr_values.len(), 12);
        let expected_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 1.0, 0.0, 0.0, 1.0];
        for (actual, expected) in arr_values.iter().zip(expected_values.iter()) {
            assert!((actual - expected).abs() < 0.0001);
        }
    }

    #[test]
    fn test_jitter_centroids() {
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(DEFAULT_SEED.0);
        let mut centroids_gen = array::cycle_unit_circle(4);
        let centroids = centroids_gen.generate(RowCount::from(4), &mut rng).unwrap();

        let centroid_values = centroids
            .as_fixed_size_list()
            .values()
            .as_primitive::<Float32Type>()
            .values()
            .to_vec();

        let mut jitter_jen = array::jitter_centroids(centroids, 0.001);
        let jittered = jitter_jen.generate(RowCount::from(100), &mut rng).unwrap();

        let values = jittered
            .as_fixed_size_list()
            .values()
            .as_primitive::<Float32Type>()
            .values()
            .to_vec();

        for i in 0..100 {
            let centroid = i % 4;
            let centroid_x = centroid_values[centroid * 2];
            let centroid_y = centroid_values[centroid * 2 + 1];
            let value_x = values[i * 2];
            let value_y = values[i * 2 + 1];

            let l2_dist = ((value_x - centroid_x).powi(2) + (value_y - centroid_y).powi(2)).sqrt();
            assert!(l2_dist < 0.001001);
            assert!(l2_dist > 0.000999);
        }
    }

    #[test]
    fn test_rand_schema() {
        let schema = Schema::new(vec![
            Field::new("a", DataType::Int32, true),
            Field::new("b", DataType::Utf8, true),
            Field::new("c", DataType::Float32, true),
            Field::new("d", DataType::Int32, true),
            Field::new("e", DataType::Int32, true),
        ]);
        let rbr = rand(&schema)
            .into_reader_bytes(
                ByteCount::from(1024 * 1024),
                BatchCount::from(8),
                RoundingBehavior::ExactOrErr,
            )
            .unwrap();
        assert_eq!(*rbr.schema(), schema);

        let batches = rbr.map(|val| val.unwrap()).collect::<Vec<_>>();
        assert_eq!(batches.len(), 8);

        for batch in batches {
            assert_eq!(batch.num_rows(), 1024 * 1024 / 32);
            assert_eq!(batch.num_columns(), 5);
        }
    }
}