// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Benchmarks for the `coalesce` kernels in Arrow.

use arrow::util::bench_util::*;
use std::sync::Arc;

use arrow::array::*;
use arrow_array::types::{Float64Type, Int32Type, TimestampNanosecondType};
use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
use arrow_select::coalesce::BatchCoalescer;
use criterion::{Criterion, criterion_group, criterion_main};
use rand::{Rng, SeedableRng, rngs::StdRng};

/// Benchmarks for generating evently sized output RecordBatches
/// from a sequence of filtered source batches
///
fn add_all_filter_benchmarks(c: &mut Criterion) {
    let batch_size = 8192; // 8K rows is a commonly used size for batches

    // Multiple primitive types
    let primitive_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new(
            "timestamp_val",
            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
            true,
        ),
    ]));

    // Single StringViewArray
    let single_schema = SchemaRef::new(Schema::new(vec![Field::new(
        "value",
        DataType::Utf8View,
        true,
    )]));

    // Single BinaryViewArray
    let single_binaryview_schema = SchemaRef::new(Schema::new(vec![Field::new(
        "value",
        DataType::BinaryView,
        true,
    )]));

    // Mixed primitive, StringViewArray
    let mixed_utf8view_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("utf8view_val", DataType::Utf8View, true),
    ]));

    // Mixed primitive, BinaryViewArray
    let mixed_binaryview_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("binaryview_val", DataType::BinaryView, true),
    ]));

    // Mixed primitive, StringArray
    let mixed_utf8_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("utf8", DataType::Utf8, true),
    ]));

    // dictionary types
    //
    let mixed_dict_schema = SchemaRef::new(Schema::new(vec![
        Field::new(
            "string_dict",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        ),
        Field::new("float_val1", DataType::Float64, true),
        Field::new("float_val2", DataType::Float64, true),
        // TODO model other dictionary types here (FixedSizeBinary for example)
    ]));

    // Null density: 0, 10%
    for null_density in [0.0, 0.1] {
        // Selectivity: 0.1%, 1%, 10%, 80%
        for selectivity in [0.001, 0.01, 0.1, 0.8] {
            FilterBenchmarkBuilder {
                c,
                name: "primitive",
                batch_size,
                num_output_batches: 50,
                null_density,
                selectivity,
                max_string_len: 30,
                schema: &primitive_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "single_utf8view",
                batch_size,
                num_output_batches: 50,
                null_density,
                selectivity,
                max_string_len: 30,
                schema: &single_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "single_utf8view (max_string_len=8)",
                batch_size,
                num_output_batches: 50,
                null_density,
                selectivity,
                max_string_len: 8,
                schema: &single_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "single_binaryview",
                batch_size,
                num_output_batches: 50,
                null_density,
                selectivity,
                max_string_len: 30,
                schema: &single_binaryview_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "single_binaryview (max_string_len=8)",
                batch_size,
                num_output_batches: 50,
                null_density,
                selectivity,
                max_string_len: 8,
                schema: &single_binaryview_schema,
            }
            .build();

            // Model mostly short strings, but some longer ones
            FilterBenchmarkBuilder {
                c,
                name: "mixed_utf8view (max_string_len=20)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 20,
                schema: &mixed_utf8view_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_utf8view (max_string_len=8)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 8,
                schema: &mixed_utf8view_schema,
            }
            .build();

            // Model mostly longer strings
            FilterBenchmarkBuilder {
                c,
                name: "mixed_utf8view (max_string_len=128)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 128,
                schema: &mixed_utf8view_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_binaryview (max_string_len=20)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 20,
                schema: &mixed_binaryview_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_binaryview (max_string_len=8)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 8,
                schema: &mixed_binaryview_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_binaryview (max_string_len=128)",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 128,
                schema: &mixed_binaryview_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_utf8",
                batch_size,
                num_output_batches: 20,
                null_density,
                selectivity,
                max_string_len: 30,
                schema: &mixed_utf8_schema,
            }
            .build();

            FilterBenchmarkBuilder {
                c,
                name: "mixed_dict",
                batch_size,
                num_output_batches: 10,
                null_density,
                selectivity,
                max_string_len: 30,
                schema: &mixed_dict_schema,
            }
            .build();
        }
    }
}

fn add_all_take_benchmarks(c: &mut Criterion) {
    let batch_size = 8192;

    let primitive_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new(
            "timestamp_val",
            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
            true,
        ),
    ]));

    let single_schema = SchemaRef::new(Schema::new(vec![Field::new(
        "value",
        DataType::Utf8View,
        true,
    )]));

    let single_binaryview_schema = SchemaRef::new(Schema::new(vec![Field::new(
        "value",
        DataType::BinaryView,
        true,
    )]));

    let mixed_utf8view_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("utf8view_val", DataType::Utf8View, true),
    ]));

    let mixed_binaryview_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("binaryview_val", DataType::BinaryView, true),
    ]));

    let mixed_utf8_schema = SchemaRef::new(Schema::new(vec![
        Field::new("int32_val", DataType::Int32, true),
        Field::new("float_val", DataType::Float64, true),
        Field::new("utf8", DataType::Utf8, true),
    ]));

    let mixed_dict_schema = SchemaRef::new(Schema::new(vec![
        Field::new(
            "string_dict",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        ),
        Field::new("float_val1", DataType::Float64, true),
        Field::new("float_val2", DataType::Float64, true),
    ]));

    for null_density in [0.0, 0.1] {
        for selectivity in [0.001, 0.01, 0.1, 0.8] {
            for scenario in [
                TakeBenchmarkScenario {
                    name: "primitive",
                    num_output_batches: 50,
                    max_string_len: 30,
                    schema: &primitive_schema,
                },
                TakeBenchmarkScenario {
                    name: "single_utf8view",
                    num_output_batches: 50,
                    max_string_len: 30,
                    schema: &single_schema,
                },
                TakeBenchmarkScenario {
                    name: "single_binaryview",
                    num_output_batches: 50,
                    max_string_len: 30,
                    schema: &single_binaryview_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_utf8view (max_string_len=20)",
                    num_output_batches: 20,
                    max_string_len: 20,
                    schema: &mixed_utf8view_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_utf8view (max_string_len=128)",
                    num_output_batches: 20,
                    max_string_len: 128,
                    schema: &mixed_utf8view_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_binaryview (max_string_len=20)",
                    num_output_batches: 20,
                    max_string_len: 20,
                    schema: &mixed_binaryview_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_binaryview (max_string_len=128)",
                    num_output_batches: 20,
                    max_string_len: 128,
                    schema: &mixed_binaryview_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_utf8",
                    num_output_batches: 20,
                    max_string_len: 30,
                    schema: &mixed_utf8_schema,
                },
                TakeBenchmarkScenario {
                    name: "mixed_dict",
                    num_output_batches: 10,
                    max_string_len: 30,
                    schema: &mixed_dict_schema,
                },
            ] {
                TakeBenchmarkBuilder::from_scenario(
                    c,
                    batch_size,
                    null_density,
                    selectivity,
                    scenario,
                )
                .build();
            }
        }
    }

    // Repeated indices make the taken batch much larger than the source batch,
    // which exercises the materialized fallback path for unsupported schemas.
    for (name, schema) in [
        ("primitive extra_large_repeat", &primitive_schema),
        ("mixed_utf8 extra_large_repeat", &mixed_utf8_schema),
    ] {
        TakeBenchmarkBuilder::from_scenario(
            c,
            batch_size,
            0.0,
            1.0,
            TakeBenchmarkScenario {
                name,
                num_output_batches: 64,
                max_string_len: 30,
                schema,
            },
        )
        .with_biggest_coalesce_batch_size(1024)
        .with_index_output_len(131_072)
        .with_drain_all_completed_batches()
        .build();
    }
}

criterion_group!(benches, add_all_filter_benchmarks, add_all_take_benchmarks);
criterion_main!(benches);

/// Run the filters with a batch_size, null_density, selectivity, and schema
struct FilterBenchmarkBuilder<'a> {
    /// Benchmark criterion instance
    c: &'a mut Criterion,
    /// Name of the benchmark
    name: &'a str,
    /// Size of the input and output batches
    batch_size: usize,
    /// Number of output batches to collect (tuned to keep benchmark time reasonable)
    num_output_batches: usize,
    /// between 0.0 .. 1.0, percent of data rows (not filter rows) that should be null
    null_density: f32,
    /// between 0.0 .. 1.0, percent of rows that should be kept by the filter
    selectivity: f32,
    /// The maximum length of strings in the data stream
    ///
    /// For StringViewArray, strings <= 12 bytes are stored inline, longer
    /// strings are stored in a separate buffer so it is important to vary to
    /// mix the relative paths
    max_string_len: usize,
    /// Schema of the data stream
    schema: &'a SchemaRef,
}

impl FilterBenchmarkBuilder<'_> {
    fn build(self) {
        let Self {
            c,
            name,
            batch_size,
            num_output_batches,
            null_density,
            selectivity,
            max_string_len,
            schema,
        } = self;

        let filters = FilterStreamBuilder::new()
            .with_batch_size(batch_size)
            .with_true_density(selectivity)
            .with_null_density(0.0) // no nulls in the filter
            .build();

        let data = DataStreamBuilder::new(Arc::clone(schema))
            .with_batch_size(batch_size)
            .with_null_density(null_density)
            .with_max_string_len(max_string_len)
            .build();

        // Keep feeding the filter stream into the coalescer until we hit a total number of output batches
        let id = format!(
            "filter: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}"
        );
        c.bench_function(&id, |b| {
            b.iter(|| {
                filter_streams(num_output_batches, filters.clone(), data.clone());
            })
        });
    }
}

struct TakeBenchmarkBuilder<'a> {
    c: &'a mut Criterion,
    name: &'a str,
    batch_size: usize,
    num_output_batches: usize,
    null_density: f32,
    selectivity: f32,
    max_string_len: usize,
    schema: &'a SchemaRef,
    biggest_coalesce_batch_size: Option<usize>,
    index_output_len: Option<usize>,
    drain_all_completed_batches: bool,
}

#[derive(Clone, Copy)]
struct TakeBenchmarkScenario<'a> {
    name: &'a str,
    num_output_batches: usize,
    max_string_len: usize,
    schema: &'a SchemaRef,
}

impl<'a> TakeBenchmarkBuilder<'a> {
    fn from_scenario(
        c: &'a mut Criterion,
        batch_size: usize,
        null_density: f32,
        selectivity: f32,
        scenario: TakeBenchmarkScenario<'a>,
    ) -> TakeBenchmarkBuilder<'a> {
        let TakeBenchmarkScenario {
            name,
            num_output_batches,
            max_string_len,
            schema,
        } = scenario;

        TakeBenchmarkBuilder {
            c,
            name,
            batch_size,
            num_output_batches,
            null_density,
            selectivity,
            max_string_len,
            schema,
            biggest_coalesce_batch_size: None,
            index_output_len: None,
            drain_all_completed_batches: false,
        }
    }

    fn with_biggest_coalesce_batch_size(mut self, limit: usize) -> Self {
        self.biggest_coalesce_batch_size = Some(limit);
        self
    }

    fn with_index_output_len(mut self, output_len: usize) -> Self {
        self.index_output_len = Some(output_len);
        self
    }

    fn with_drain_all_completed_batches(mut self) -> Self {
        self.drain_all_completed_batches = true;
        self
    }

    fn build(self) {
        let Self {
            c,
            name,
            batch_size,
            num_output_batches,
            null_density,
            selectivity,
            max_string_len,
            schema,
            biggest_coalesce_batch_size,
            index_output_len,
            drain_all_completed_batches,
        } = self;

        let output_len = index_output_len
            .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize);

        let indices = match index_output_len {
            Some(_) => IndexStreamBuilder::new()
                .with_batch_size(batch_size)
                .with_output_len(output_len)
                .build(),
            None => IndexStreamBuilder::new()
                .with_batch_size(batch_size)
                .with_selectivity(selectivity)
                .build(),
        };

        let data = DataStreamBuilder::new(Arc::clone(schema))
            .with_batch_size(batch_size)
            .with_null_density(null_density)
            .with_max_string_len(max_string_len)
            .build();

        let id = if index_output_len.is_some() || biggest_coalesce_batch_size.is_some() {
            format!(
                "take: {name}, input: {batch_size}, output: {output_len}, nulls: {null_density}, biggest: {biggest_coalesce_batch_size:?}"
            )
        } else {
            format!("take: {name}, {batch_size}, nulls: {null_density}, selectivity: {selectivity}")
        };
        c.bench_function(&id, |b| {
            b.iter(|| {
                take_streams(
                    num_output_batches,
                    indices.clone(),
                    data.clone(),
                    biggest_coalesce_batch_size,
                    drain_all_completed_batches,
                );
            })
        });
    }
}

/// Pull RecordBatches from a data stream and apply a sequence of
/// filters from a filter stream until we have a specified number of output
/// batches.
fn filter_streams(
    mut num_output_batches: usize,
    mut filter_stream: FilterStream,
    mut data_stream: DataStream,
) {
    let schema = data_stream.schema();
    let batch_size = data_stream.batch_size();
    let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size);

    while num_output_batches > 0 {
        let filter = filter_stream.next_filter();
        let batch = data_stream.next_batch();
        coalescer
            .push_batch_with_filter(batch.clone(), filter)
            .unwrap();
        // consume (but discard) the output batch
        if coalescer.next_completed_batch().is_some() {
            num_output_batches -= 1;
        }
    }
}

fn take_streams(
    mut num_output_batches: usize,
    mut index_stream: IndexStream,
    mut data_stream: DataStream,
    biggest_coalesce_batch_size: Option<usize>,
    drain_all_completed_batches: bool,
) {
    let schema = data_stream.schema();
    let batch_size = data_stream.batch_size();
    let mut coalescer = BatchCoalescer::new(Arc::clone(schema), batch_size)
        .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size);

    while num_output_batches > 0 {
        let indices = index_stream.next_indices();
        let batch = data_stream.next_batch();
        coalescer
            .push_batch_with_indices(batch.clone(), indices)
            .unwrap();
        if drain_all_completed_batches {
            while num_output_batches > 0 && coalescer.next_completed_batch().is_some() {
                num_output_batches -= 1;
            }
        } else if coalescer.next_completed_batch().is_some() {
            num_output_batches -= 1;
        }
    }
}

/// Stream of filters to apply to a sequence of input RecordBatches
///
/// This pre-computes a sequence of filters and then repeats it forever.
#[derive(Debug, Clone)]
struct FilterStream {
    index: usize,
    // arc'd so it is cheaply cloned
    batches: Arc<[BooleanArray]>,
}

#[derive(Debug, Clone)]
struct IndexStream {
    index: usize,
    batches: Arc<[UInt32Array]>,
}

impl IndexStream {
    fn next_indices(&mut self) -> &UInt32Array {
        let current_index = self.index;
        self.index += 1;
        if self.index >= self.batches.len() {
            self.index = 0;
        }
        self.batches
            .get(current_index)
            .expect("No more index batches available")
    }
}

impl FilterStream {
    pub fn next_filter(&mut self) -> &BooleanArray {
        let current_index = self.index;
        self.index += 1;
        if self.index >= self.batches.len() {
            self.index = 0; // loop back to the start
        }
        self.batches
            .get(current_index)
            .expect("No more filters available")
    }
}

#[derive(Debug)]
struct FilterStreamBuilder {
    batch_size: usize,
    num_batches: usize, // number of unique batches to create
    null_density: f32,
    true_density: f32,
}

impl FilterStreamBuilder {
    fn new() -> Self {
        FilterStreamBuilder {
            batch_size: 8192,  // default batch size
            num_batches: 11,   // default number of unique batches (different than data stream)
            null_density: 0.0, // default null density
            true_density: 0.5, // default true density
        }
    }
    /// set the batch size for the filter stream
    fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// set the null density for the filter stream
    fn with_null_density(mut self, null_density: f32) -> Self {
        assert!((0.0..=1.0).contains(&null_density));
        self.null_density = null_density;
        self
    }

    /// set the true density for the filter stream
    fn with_true_density(mut self, true_density: f32) -> Self {
        assert!((0.0..=1.0).contains(&true_density));
        self.true_density = true_density;
        self
    }
    fn build(self) -> FilterStream {
        let Self {
            batch_size,
            num_batches,
            null_density,
            true_density,
        } = self;
        let batches = (0..num_batches)
            .map(|_| create_boolean_array(batch_size, null_density, true_density))
            .collect::<Vec<_>>();

        FilterStream {
            index: 0,
            batches: Arc::from(batches),
        }
    }
}

#[derive(Debug)]
struct IndexStreamBuilder {
    batch_size: usize,
    num_batches: usize,
    output_len: Option<usize>,
    selectivity: f32,
}

impl IndexStreamBuilder {
    fn new() -> Self {
        Self {
            batch_size: 8192,
            num_batches: 11,
            output_len: None,
            selectivity: 0.5,
        }
    }

    fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    fn with_selectivity(mut self, selectivity: f32) -> Self {
        assert!((0.0..=1.0).contains(&selectivity));
        self.selectivity = selectivity;
        self
    }

    fn with_output_len(mut self, output_len: usize) -> Self {
        self.output_len = Some(output_len.max(1));
        self
    }

    fn build(self) -> IndexStream {
        let Self {
            batch_size,
            num_batches,
            output_len,
            selectivity,
        } = self;

        let output_len = output_len
            .unwrap_or_else(|| ((batch_size as f32) * selectivity).round().max(1.0) as usize);
        let batches = (0..num_batches)
            .map(|seed| create_index_array(batch_size, output_len, seed as u64))
            .collect::<Vec<_>>();

        IndexStream {
            index: 0,
            batches: Arc::from(batches),
        }
    }
}

fn create_index_array(input_len: usize, output_len: usize, seed: u64) -> UInt32Array {
    let mut rng = StdRng::seed_from_u64(seed);
    UInt32Array::from_iter_values(
        (0..output_len).map(|_| rng.random_range(0..u32::try_from(input_len).unwrap())),
    )
}

#[derive(Debug, Clone)]
struct DataStream {
    schema: SchemaRef,
    index: usize,
    batch_size: usize,
    // arc'd so it is cheaply cloned
    batches: Arc<[RecordBatch]>,
}

impl DataStream {
    /// Returns the schema for this data stream
    pub fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    /// Returns the batch size
    pub fn batch_size(&self) -> usize {
        self.batch_size
    }

    fn next_batch(&mut self) -> &RecordBatch {
        let current_index = self.index;
        self.index += 1;
        if self.index >= self.batches.len() {
            self.index = 0; // loop back to the start
        }
        self.batches
            .get(current_index)
            .expect("No more batches available")
    }
}

#[derive(Debug, Clone)]
struct DataStreamBuilder {
    schema: SchemaRef,
    batch_size: usize,
    null_density: f32,
    num_batches: usize,    // number of unique batches to create
    max_string_len: usize, // maximum length of strings in the data stream
}

impl DataStreamBuilder {
    fn new(schema: SchemaRef) -> Self {
        DataStreamBuilder {
            schema,
            batch_size: 8192,
            null_density: 0.0,
            num_batches: 10,
            max_string_len: 30,
        }
    }

    /// set the batch size for the data stream
    fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// set the null density for the data stream
    fn with_null_density(mut self, null_density: f32) -> Self {
        assert!((0.0..=1.0).contains(&null_density));
        self.null_density = null_density;
        self
    }

    fn with_max_string_len(mut self, max_string_len: usize) -> Self {
        self.max_string_len = max_string_len;
        self
    }

    /// build the data stream (not implemented yet)
    fn build(self) -> DataStream {
        let batches = (0..self.num_batches)
            .map(|seed| {
                let columns = self
                    .schema
                    .fields()
                    .iter()
                    .map(|field| self.create_input_array(field, seed as u64))
                    .collect::<Vec<_>>();
                RecordBatch::try_new(self.schema.clone(), columns).unwrap()
            })
            .collect::<Vec<_>>();

        let Self {
            schema,
            batch_size,
            null_density: _,
            num_batches: _,
            max_string_len: _,
        } = self;

        DataStream {
            schema,
            index: 0,
            batch_size,
            batches: Arc::from(batches),
        }
    }

    fn create_input_array(&self, field: &Field, seed: u64) -> ArrayRef {
        match field.data_type() {
            DataType::Int32 => Arc::new(create_primitive_array_with_seed::<Int32Type>(
                self.batch_size,
                self.null_density,
                seed,
            )),
            DataType::Float64 => Arc::new(create_primitive_array_with_seed::<Float64Type>(
                self.batch_size,
                self.null_density,
                seed,
            )),
            DataType::Timestamp(TimeUnit::Nanosecond, Some(tz)) => Arc::new(
                create_primitive_array_with_seed::<TimestampNanosecondType>(
                    self.batch_size,
                    self.null_density,
                    seed,
                )
                .with_timezone(Arc::clone(tz)),
            ),
            DataType::Utf8 => Arc::new(create_string_array::<i32>(
                self.batch_size,
                self.null_density,
            )), // TODO seed
            DataType::Utf8View => {
                Arc::new(create_string_view_array_with_max_len(
                    self.batch_size,
                    self.null_density,
                    self.max_string_len,
                )) // TODO seed
            }
            DataType::BinaryView => Arc::new(BinaryViewArray::from_iter(
                create_binary_array_with_len_range_and_prefix_and_seed::<i32>(
                    self.batch_size,
                    self.null_density,
                    0,
                    self.max_string_len,
                    b"",
                    seed,
                )
                .iter(),
            )),
            DataType::Dictionary(key_type, value_type)
                if key_type.as_ref() == &DataType::Int32
                    && value_type.as_ref() == &DataType::Utf8 =>
            {
                Arc::new(create_string_dict_array::<Int32Type>(
                    self.batch_size,
                    self.null_density,
                    self.max_string_len,
                )) // TODO seed
            }
            _ => panic!("Unsupported data type: {field:?}"),
        }
    }
}
