arrow_flight/

encode.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

use std::{collections::VecDeque, fmt::Debug, pin::Pin, sync::Arc, task::Poll};

use crate::{FlightData, FlightDescriptor, SchemaAsIpc, error::Result};

use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UnionArray};
use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteContext, IpcWriteOptions};

use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, Schema, SchemaRef, UnionMode};
use bytes::Bytes;
use futures::{Stream, StreamExt, ready, stream::BoxStream};

/// Creates a [`Stream`] of [`FlightData`]s from a
/// `Stream` of [`Result`]<[`RecordBatch`], [`FlightError`]>.
///
/// This can be used to implement [`FlightService::do_get`] in an
/// Arrow Flight implementation;
///
/// This structure encodes a stream of `Result`s rather than `RecordBatch`es  to
/// propagate errors from streaming execution, where the generation of the
/// `RecordBatch`es is incremental, and an error may occur even after
/// several have already been successfully produced.
///
/// # Caveats
///
/// 1. All input record batches must have the same schema. The first batch sets the
///    expected schema unless [`FlightDataEncoderBuilder::with_schema`] supplies it.
///    A batch with a different schema produces an error and ends the stream.
/// 2. When [`DictionaryHandling`] is [`DictionaryHandling::Hydrate`],
///    [`DictionaryArray`]s are converted to their underlying types prior to
///    transport.
///    When [`DictionaryHandling`] is [`DictionaryHandling::Resend`], Dictionary [`FlightData`] is sent with every
///    [`RecordBatch`] that contains a [`DictionaryArray`](arrow_array::array::DictionaryArray).
///    See <https://github.com/apache/arrow-rs/issues/3389>.
///
/// [`DictionaryArray`]: arrow_array::array::DictionaryArray
///
/// # Example
/// ```no_run
/// # use std::sync::Arc;
/// # use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
/// # async fn f() {
/// # let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);
/// # let batch = RecordBatch::try_from_iter(vec![
/// #      ("a", Arc::new(c1) as ArrayRef)
/// #   ])
/// #   .expect("cannot create record batch");
/// use arrow_flight::encode::FlightDataEncoderBuilder;
///
/// // Get an input stream of Result<RecordBatch, FlightError>
/// let input_stream = futures::stream::iter(vec![Ok(batch)]);
///
/// // Build a stream of `Result<FlightData>` (e.g. to return for do_get)
/// let flight_data_stream = FlightDataEncoderBuilder::new()
///  .build(input_stream);
///
/// // Create a tonic `Response` that can be returned from a Flight server
/// let response = tonic::Response::new(flight_data_stream);
/// # }
/// ```
///
/// # Example: Sending `Vec<RecordBatch>`
///
/// You can create a [`Stream`] to pass to [`Self::build`] from an existing
/// `Vec` of `RecordBatch`es like this:
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
/// # async fn f() {
/// # fn make_batches() -> Vec<RecordBatch> {
/// #   let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);
/// #   let batch = RecordBatch::try_from_iter(vec![
/// #      ("a", Arc::new(c1) as ArrayRef)
/// #   ])
/// #   .expect("cannot create record batch");
/// #   vec![batch.clone(), batch.clone()]
/// # }
/// use arrow_flight::encode::FlightDataEncoderBuilder;
///
/// // Get batches that you want to send via Flight
/// let batches: Vec<RecordBatch> = make_batches();
///
/// // Create an input stream of Result<RecordBatch, FlightError>
/// let input_stream = futures::stream::iter(
///   batches.into_iter().map(Ok)
/// );
///
/// // Build a stream of `Result<FlightData>` (e.g. to return for do_get)
/// let flight_data_stream = FlightDataEncoderBuilder::new()
///  .build(input_stream);
/// # }
/// ```
///
/// # Example: Determining schema of encoded data
///
/// Encoding flight data may hydrate dictionaries, see [`DictionaryHandling`] for more information,
/// which changes the schema of the encoded data compared to the input record batches.
/// The fully hydrated schema can be accessed using the [`FlightDataEncoder::known_schema`] method
/// and explicitly informing the builder of the schema using [`FlightDataEncoderBuilder::with_schema`].
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
/// # async fn f() {
/// # let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);
/// # let batch = RecordBatch::try_from_iter(vec![
/// #      ("a", Arc::new(c1) as ArrayRef)
/// #   ])
/// #   .expect("cannot create record batch");
/// use arrow_flight::encode::FlightDataEncoderBuilder;
///
/// // Get the schema of the input stream
/// let schema = batch.schema();
///
/// // Get an input stream of Result<RecordBatch, FlightError>
/// let input_stream = futures::stream::iter(vec![Ok(batch)]);
///
/// // Build a stream of `Result<FlightData>` (e.g. to return for do_get)
/// let flight_data_stream = FlightDataEncoderBuilder::new()
///  // Inform the builder of the input stream schema
///  .with_schema(schema)
///  .build(input_stream);
///
/// // Retrieve the schema of the encoded data
/// let encoded_schema = flight_data_stream.known_schema();
/// # }
/// ```
///
/// [`FlightService::do_get`]: crate::flight_service_server::FlightService::do_get
/// [`FlightError`]: crate::error::FlightError
#[derive(Debug)]
pub struct FlightDataEncoderBuilder {
    /// The maximum approximate target message size in bytes
    /// (see details on [`Self::with_max_flight_data_size`]).
    max_flight_data_size: usize,
    /// Ipc writer options
    options: IpcWriteOptions,
    /// Metadata to add to the schema message
    app_metadata: Bytes,
    /// Optional schema, if known before data.
    schema: Option<SchemaRef>,
    /// Optional flight descriptor, if known before data.
    descriptor: Option<FlightDescriptor>,
    /// Deterimines how `DictionaryArray`s are encoded for transport.
    /// See [`DictionaryHandling`] for more information.
    dictionary_handling: DictionaryHandling,
}

/// Default target size for encoded [`FlightData`].
///
/// Note this value would normally be 4MB, but the size calculation is
/// somewhat inexact, so we set it to 2MB.
pub const GRPC_TARGET_MAX_FLIGHT_SIZE_BYTES: usize = 2097152;

impl Default for FlightDataEncoderBuilder {
    fn default() -> Self {
        Self {
            max_flight_data_size: GRPC_TARGET_MAX_FLIGHT_SIZE_BYTES,
            options: IpcWriteOptions::default(),
            app_metadata: Bytes::new(),
            schema: None,
            descriptor: None,
            dictionary_handling: DictionaryHandling::Hydrate,
        }
    }
}

impl FlightDataEncoderBuilder {
    /// Create a new [`FlightDataEncoderBuilder`].
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the (approximate) maximum size, in bytes, of the
    /// [`FlightData`] produced by this encoder. Defaults to 2MB.
    ///
    /// Since there is often a maximum message size for gRPC messages
    /// (typically around 4MB), this encoder splits up [`RecordBatch`]s
    /// (preserving order) into multiple [`FlightData`] objects to
    /// limit the size individual messages sent via gRPC.
    ///
    /// The size is approximate because of the additional encoding
    /// overhead on top of the underlying data buffers themselves.
    pub fn with_max_flight_data_size(mut self, max_flight_data_size: usize) -> Self {
        self.max_flight_data_size = max_flight_data_size;
        self
    }

    /// Set [`DictionaryHandling`] for encoder
    pub fn with_dictionary_handling(mut self, dictionary_handling: DictionaryHandling) -> Self {
        self.dictionary_handling = dictionary_handling;
        self
    }

    /// Specify application specific metadata included in the
    /// [`FlightData::app_metadata`] field of the the first Schema
    /// message
    pub fn with_metadata(mut self, app_metadata: Bytes) -> Self {
        self.app_metadata = app_metadata;
        self
    }

    /// Set the [`IpcWriteOptions`] used to encode the [`RecordBatch`]es for transport.
    pub fn with_options(mut self, options: IpcWriteOptions) -> Self {
        self.options = options;
        self
    }

    /// Specify a schema for the RecordBatches being sent. If a schema
    /// is not specified, an encoded Schema message will be sent when
    /// the first [`RecordBatch`], if any, is encoded. Some clients
    /// expect a Schema message even if there is no data sent.
    /// Every input batch must match this schema.
    pub fn with_schema(mut self, schema: SchemaRef) -> Self {
        self.schema = Some(schema);
        self
    }

    /// Specify a flight descriptor in the first FlightData message.
    pub fn with_flight_descriptor(mut self, descriptor: Option<FlightDescriptor>) -> Self {
        self.descriptor = descriptor;
        self
    }

    /// Takes a [`Stream`] of [`Result<RecordBatch>`] and returns a [`Stream`]
    /// of [`FlightData`], consuming self.
    ///
    /// See example on [`Self`] and [`FlightDataEncoder`] for more details
    pub fn build<S>(self, input: S) -> FlightDataEncoder
    where
        S: Stream<Item = Result<RecordBatch>> + Send + 'static,
    {
        let Self {
            max_flight_data_size,
            options,
            app_metadata,
            schema,
            descriptor,
            dictionary_handling,
        } = self;

        FlightDataEncoder::new(
            input.boxed(),
            schema,
            max_flight_data_size,
            options,
            app_metadata,
            descriptor,
            dictionary_handling,
        )
    }
}

/// Stream that encodes a stream of record batches to flight data.
///
/// See [`FlightDataEncoderBuilder`] for details and example.
pub struct FlightDataEncoder {
    /// Input stream
    inner: BoxStream<'static, Result<RecordBatch>>,
    /// Schema of input batches before dictionary handling
    input_schema: Option<SchemaRef>,
    /// schema, set after the first batch
    schema: Option<SchemaRef>,
    /// Target maximum size of flight data
    /// (see details on [`FlightDataEncoderBuilder::with_max_flight_data_size`]).
    max_flight_data_size: usize,
    /// do the encoding / tracking of dictionaries
    encoder: FlightIpcEncoder,
    /// optional metadata to add to schema FlightData
    app_metadata: Option<Bytes>,
    /// data queued up to send but not yet sent
    queue: VecDeque<FlightData>,
    /// Is this stream done (inner is empty or errored)
    done: bool,
    /// cleared after the first FlightData message is sent
    descriptor: Option<FlightDescriptor>,
    /// Deterimines how `DictionaryArray`s are encoded for transport.
    /// See [`DictionaryHandling`] for more information.
    dictionary_handling: DictionaryHandling,
}

impl FlightDataEncoder {
    fn new(
        inner: BoxStream<'static, Result<RecordBatch>>,
        schema: Option<SchemaRef>,
        max_flight_data_size: usize,
        options: IpcWriteOptions,
        app_metadata: Bytes,
        descriptor: Option<FlightDescriptor>,
        dictionary_handling: DictionaryHandling,
    ) -> Self {
        let mut encoder = Self {
            inner,
            input_schema: schema.clone(),
            schema: None,
            max_flight_data_size,
            encoder: FlightIpcEncoder::new(
                options,
                dictionary_handling != DictionaryHandling::Resend,
            ),
            app_metadata: Some(app_metadata),
            queue: VecDeque::new(),
            done: false,
            descriptor,
            dictionary_handling,
        };

        // If schema is known up front, enqueue it immediately
        if let Some(schema) = schema {
            encoder.encode_schema(&schema);
        }

        encoder
    }

    /// Report the schema of the encoded data when known.
    /// A schema is known when provided via the [`FlightDataEncoderBuilder::with_schema`] method.
    pub fn known_schema(&self) -> Option<SchemaRef> {
        self.schema.clone()
    }

    /// Place the `FlightData` in the queue to send
    #[inline]
    fn queue_message(&mut self, mut data: FlightData) {
        if let Some(descriptor) = self.descriptor.take() {
            data.flight_descriptor = Some(descriptor);
        }
        self.queue.push_back(data);
    }

    /// Encodes schema as a [`FlightData`] in self.queue.
    /// Updates `self.schema` and returns the new schema
    fn encode_schema(&mut self, schema: &SchemaRef) -> SchemaRef {
        // The first message is the schema message, and all
        // batches have the same schema
        let send_dictionaries = self.dictionary_handling == DictionaryHandling::Resend;
        let schema = Arc::new(prepare_schema_for_flight(
            schema,
            &mut self.encoder.dictionary_tracker,
            send_dictionaries,
        ));
        let mut schema_flight_data = self.encoder.encode_schema(&schema);

        // attach any metadata requested
        if let Some(app_metadata) = self.app_metadata.take() {
            schema_flight_data.app_metadata = app_metadata;
        }
        self.queue_message(schema_flight_data);
        // remember schema
        self.schema = Some(schema.clone());
        schema
    }

    /// Encodes batch into one or more `FlightData` messages in self.queue
    fn encode_batch(&mut self, batch: RecordBatch) -> Result<()> {
        match &self.input_schema {
            Some(expected) if expected.as_ref() != batch.schema_ref().as_ref() => {
                return Err(ArrowError::SchemaError(
                    "Record batch schema does not match input schema".to_string(),
                )
                .into());
            }
            None => self.input_schema = Some(batch.schema()),
            _ => {}
        }

        let schema = match &self.schema {
            Some(schema) => schema.clone(),
            // encode the schema if this is the first time we have seen it
            None => self.encode_schema(batch.schema_ref()),
        };

        let batch = match self.dictionary_handling {
            DictionaryHandling::Resend => batch,
            DictionaryHandling::Hydrate => hydrate_dictionaries(&batch, schema)?,
        };

        let batches = split_batch_for_grpc_response(batch, self.max_flight_data_size);
        let last = batches.len().saturating_sub(1); // handle empty batches  
        for (i, batch) in batches.into_iter().enumerate() {
            self.encoder
                .ipc_write_context
                .set_reserve_scratch(i != last);
            let (flight_dictionaries, flight_batch) = self.encoder.encode_batch(&batch)?;
            for dict in flight_dictionaries {
                self.queue_message(dict);
            }
            self.queue_message(flight_batch);
        }

        Ok(())
    }
}

impl Stream for FlightDataEncoder {
    type Item = Result<FlightData>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        loop {
            if self.done && self.queue.is_empty() {
                return Poll::Ready(None);
            }

            // Any messages queued to send?
            if let Some(data) = self.queue.pop_front() {
                return Poll::Ready(Some(Ok(data)));
            }

            // Get next batch
            let batch = ready!(self.inner.poll_next_unpin(cx));

            match batch {
                None => {
                    // inner is done
                    self.done = true;
                    // queue must also be empty so we are done
                    assert!(self.queue.is_empty());
                    return Poll::Ready(None);
                }
                Some(Err(e)) => {
                    // error from inner
                    self.done = true;
                    self.queue.clear();
                    return Poll::Ready(Some(Err(e)));
                }
                Some(Ok(batch)) => {
                    // had data, encode into the queue
                    if let Err(e) = self.encode_batch(batch) {
                        self.done = true;
                        self.queue.clear();
                        return Poll::Ready(Some(Err(e)));
                    }
                }
            }
        }
    }
}

/// Defines how a [`FlightDataEncoder`] encodes [`DictionaryArray`]s
///
/// [`DictionaryArray`]: arrow_array::DictionaryArray
///
/// In the arrow flight protocol dictionary values and keys are sent as two separate messages.
/// When a sender is encoding a [`RecordBatch`] containing [`DictionaryArray`] columns, it will
/// first send a dictionary batch (a batch with header `MessageHeader::DictionaryBatch`) containing
/// the dictionary values. The receiver is responsible for reading this batch and maintaining state that associates
/// those dictionary values with the corresponding array using the `dict_id` as a key.
///
/// After sending the dictionary batch the sender will send the array data in a batch with header `MessageHeader::RecordBatch`.
/// For any dictionary array batches in this message, the encoded flight message will only contain the dictionary keys. The receiver
/// is then responsible for rebuilding the `DictionaryArray` on the client side using the dictionary values from the DictionaryBatch message
/// and the keys from the RecordBatch message.
///
/// For example, if we have a batch with a `TypedDictionaryArray<'_, UInt32Type, Utf8Type>` (a dictionary array where they keys are `u32` and the
/// values are `String`), then the DictionaryBatch will contain a `StringArray` and the RecordBatch will contain a `UInt32Array`.
///
/// Note that since `dict_id` defined in the `Schema` is used as a key to associate dictionary values to their arrays it is required that each
/// `DictionaryArray` in a `RecordBatch` have a unique `dict_id`.
///
/// The current implementation does not support "delta" dictionaries so a new dictionary batch will be sent each time the encoder sees a
/// dictionary which is not pointer-equal to the previously observed dictionary for a given `dict_id`.
///
/// For clients which may not support `DictionaryEncoding`, the `DictionaryHandling::Hydrate` method will bypass the process defined above
/// and "hydrate" any `DictionaryArray` in the batch to their underlying value type (e.g. `TypedDictionaryArray<'_, UInt32Type, Utf8Type>` will
/// be sent as a `StringArray`). With this method all data will be sent in `MessageHeader::RecordBatch` messages and the batch schema
/// will be adjusted so that all dictionary encoded fields are changed to fields of the dictionary value type.
#[derive(Debug, PartialEq)]
pub enum DictionaryHandling {
    /// Expands to the underlying type (default). This likely sends more data
    /// over the network but requires less memory (dictionaries are not tracked)
    /// and is more compatible with other arrow flight client implementations
    /// that may not support `DictionaryEncoding`
    ///
    /// See also:
    /// * <https://github.com/apache/arrow-rs/issues/1206>
    Hydrate,
    /// Send dictionary FlightData with every RecordBatch that contains a
    /// [`DictionaryArray`]. See [`Self::Hydrate`] for more tradeoffs. No
    /// attempt is made to skip sending the same (logical) dictionary values
    /// twice.
    ///
    /// [`DictionaryArray`]: arrow_array::DictionaryArray
    ///
    /// This requires identifying the different dictionaries in use and assigning
    //  them unique IDs
    Resend,
}

fn prepare_field_for_flight(
    field: &FieldRef,
    dictionary_tracker: &mut DictionaryTracker,
    send_dictionaries: bool,
) -> Field {
    match field.data_type() {
        DataType::List(inner) => Field::new_list(
            field.name(),
            prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries),
            field.is_nullable(),
        )
        .with_metadata(field.metadata().clone()),
        DataType::LargeList(inner) => Field::new_large_list(
            field.name(),
            prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries),
            field.is_nullable(),
        )
        .with_metadata(field.metadata().clone()),
        DataType::Struct(fields) => {
            let new_fields: Vec<Field> = fields
                .iter()
                .map(|f| prepare_field_for_flight(f, dictionary_tracker, send_dictionaries))
                .collect();
            Field::new_struct(field.name(), new_fields, field.is_nullable())
                .with_metadata(field.metadata().clone())
        }
        DataType::Union(fields, mode) => {
            let (type_ids, new_fields): (Vec<i8>, Vec<Field>) = fields
                .iter()
                .map(|(type_id, f)| {
                    (
                        type_id,
                        prepare_field_for_flight(f, dictionary_tracker, send_dictionaries),
                    )
                })
                .unzip();

            Field::new_union(field.name(), type_ids, new_fields, *mode)
        }
        DataType::Dictionary(_, value_type) => {
            if !send_dictionaries {
                // Recurse into value type to handle nested dicts being stripped
                let value_field = Field::new(
                    field.name(),
                    value_type.as_ref().clone(),
                    field.is_nullable(),
                );
                prepare_field_for_flight(
                    &Arc::new(value_field),
                    dictionary_tracker,
                    send_dictionaries,
                )
                .with_metadata(field.metadata().clone())
            } else {
                // Recurse into value type BEFORE registering this dict's id,
                // matching the depth-first order of encode_dictionaries in the
                // IPC writer which processes nested dicts before the parent.
                let value_field = Field::new("values", value_type.as_ref().clone(), true);
                prepare_field_for_flight(
                    &Arc::new(value_field),
                    dictionary_tracker,
                    send_dictionaries,
                );
                dictionary_tracker.next_dict_id();
                #[expect(deprecated)]
                Field::new_dict(
                    field.name(),
                    field.data_type().clone(),
                    field.is_nullable(),
                    0,
                    field.dict_is_ordered().unwrap_or_default(),
                )
                .with_metadata(field.metadata().clone())
            }
        }
        DataType::ListView(inner) | DataType::LargeListView(inner) => {
            let prepared = prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries);
            Field::new(
                field.name(),
                match field.data_type() {
                    DataType::ListView(_) => DataType::ListView(Arc::new(prepared)),
                    _ => DataType::LargeListView(Arc::new(prepared)),
                },
                field.is_nullable(),
            )
            .with_metadata(field.metadata().clone())
        }
        DataType::FixedSizeList(inner, size) => Field::new(
            field.name(),
            DataType::FixedSizeList(
                Arc::new(prepare_field_for_flight(
                    inner,
                    dictionary_tracker,
                    send_dictionaries,
                )),
                *size,
            ),
            field.is_nullable(),
        )
        .with_metadata(field.metadata().clone()),
        DataType::RunEndEncoded(run_ends, values) => Field::new(
            field.name(),
            DataType::RunEndEncoded(
                run_ends.clone(),
                Arc::new(prepare_field_for_flight(
                    values,
                    dictionary_tracker,
                    send_dictionaries,
                )),
            ),
            field.is_nullable(),
        )
        .with_metadata(field.metadata().clone()),
        DataType::Map(inner, sorted) => Field::new(
            field.name(),
            DataType::Map(
                prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries).into(),
                *sorted,
            ),
            field.is_nullable(),
        )
        .with_metadata(field.metadata().clone()),
        DataType::Null
        | DataType::Boolean
        | DataType::Int8
        | DataType::Int16
        | DataType::Int32
        | DataType::Int64
        | DataType::UInt8
        | DataType::UInt16
        | DataType::UInt32
        | DataType::UInt64
        | DataType::Float16
        | DataType::Float32
        | DataType::Float64
        | DataType::Timestamp(_, _)
        | DataType::Date32
        | DataType::Date64
        | DataType::Time32(_)
        | DataType::Time64(_)
        | DataType::Duration(_)
        | DataType::Interval(_)
        | DataType::Binary
        | DataType::FixedSizeBinary(_)
        | DataType::LargeBinary
        | DataType::BinaryView
        | DataType::Utf8
        | DataType::LargeUtf8
        | DataType::Utf8View
        | DataType::Decimal32(_, _)
        | DataType::Decimal64(_, _)
        | DataType::Decimal128(_, _)
        | DataType::Decimal256(_, _) => field.as_ref().clone(),
    }
}

/// Prepare an arrow Schema for transport over the Arrow Flight protocol
///
/// Convert dictionary types to underlying types
///
/// See hydrate_dictionary for more information
fn prepare_schema_for_flight(
    schema: &Schema,
    dictionary_tracker: &mut DictionaryTracker,
    send_dictionaries: bool,
) -> Schema {
    let fields: Fields = schema
        .fields()
        .iter()
        .map(|field| prepare_field_for_flight(field, dictionary_tracker, send_dictionaries))
        .collect();

    Schema::new(fields).with_metadata(schema.metadata().clone())
}

/// Split [`RecordBatch`] so it hopefully fits into a gRPC response.
///
/// Data is zero-copy sliced into batches.
///
/// Note: this method does not take into account already sliced
/// arrays: <https://github.com/apache/arrow-rs/issues/3407>
fn split_batch_for_grpc_response(
    batch: RecordBatch,
    max_flight_data_size: usize,
) -> Vec<RecordBatch> {
    let size = batch
        .columns()
        .iter()
        .map(|col| col.get_buffer_memory_size())
        .sum::<usize>();

    let n_batches =
        (size / max_flight_data_size + usize::from(size % max_flight_data_size != 0)).max(1);
    let num_rows = batch.num_rows();
    let rows_per_batch = (num_rows / n_batches).max(1);
    let mut offset = 0;
    let mut batches = Vec::with_capacity(n_batches);

    while offset < num_rows {
        let length = rows_per_batch.min(num_rows - offset);
        batches.push(batch.slice(offset, length));
        offset += length;
    }

    batches
}

/// The data needed to encode a stream of flight data, holding on to
/// shared Dictionaries.
///
/// TODO: at allow dictionaries to be flushed / avoid building them
///
/// TODO limit on the number of dictionaries???
struct FlightIpcEncoder {
    options: IpcWriteOptions,
    data_gen: IpcDataGenerator,
    dictionary_tracker: DictionaryTracker,
    ipc_write_context: IpcWriteContext,
}

impl FlightIpcEncoder {
    fn new(options: IpcWriteOptions, error_on_replacement: bool) -> Self {
        Self {
            options,
            data_gen: IpcDataGenerator::default(),
            dictionary_tracker: DictionaryTracker::new(error_on_replacement),
            ipc_write_context: IpcWriteContext::default(),
        }
    }

    /// Encode a schema as a FlightData
    fn encode_schema(&self, schema: &Schema) -> FlightData {
        SchemaAsIpc::new(schema, &self.options).into()
    }

    /// Convert a `RecordBatch` to a Vec of `FlightData` representing
    /// dictionaries and a `FlightData` representing the batch
    fn encode_batch(
        &mut self,
        batch: &RecordBatch,
    ) -> Result<(impl Iterator<Item = FlightData> + use<>, FlightData)> {
        let (encoded_dictionaries, encoded_batch) = self.data_gen.encode(
            batch,
            &mut self.dictionary_tracker,
            &self.options,
            &mut self.ipc_write_context,
        )?;

        let flight_dictionaries = encoded_dictionaries.into_iter().map(|e| e.into());
        let flight_batch = encoded_batch.into();

        Ok((flight_dictionaries, flight_batch))
    }
}

/// Hydrates any dictionaries arrays in `batch` to its underlying type. See
/// hydrate_dictionary for more information.
fn hydrate_dictionaries(batch: &RecordBatch, schema: SchemaRef) -> Result<RecordBatch> {
    let columns = schema
        .fields()
        .iter()
        .zip(batch.columns())
        .map(|(field, c)| hydrate_dictionary(c, field.data_type()))
        .collect::<Result<Vec<_>>>()?;

    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));

    Ok(RecordBatch::try_new_with_options(
        schema, columns, &options,
    )?)
}

/// Hydrates a dictionary to its underlying type.
fn hydrate_dictionary(array: &ArrayRef, data_type: &DataType) -> Result<ArrayRef> {
    let arr = match (array.data_type(), data_type) {
        (DataType::Union(_, UnionMode::Sparse), DataType::Union(fields, UnionMode::Sparse)) => {
            let union_arr = array.as_any().downcast_ref::<UnionArray>().unwrap();

            Arc::new(UnionArray::try_new(
                fields.clone(),
                union_arr.type_ids().clone(),
                None,
                fields
                    .iter()
                    .map(|(type_id, field)| {
                        Ok(arrow_cast::cast(
                            union_arr.child(type_id),
                            field.data_type(),
                        )?)
                    })
                    .collect::<Result<Vec<_>>>()?,
            )?)
        }
        (_, data_type) => arrow_cast::cast(array, data_type)?,
    };
    Ok(arr)
}

#[cfg(test)]
mod tests {
    use crate::decode::{DecodedPayload, FlightDataDecoder};
    use arrow_array::builder::{
        FixedSizeListBuilder, GenericByteDictionaryBuilder, GenericListViewBuilder, Int32Builder,
        LargeListBuilder, ListBuilder, StringDictionaryBuilder, StructBuilder,
    };
    use arrow_array::*;
    use arrow_array::{cast::downcast_array, types::*};
    use arrow_buffer::ScalarBuffer;
    use arrow_cast::pretty::pretty_format_batches;
    use arrow_ipc::{CompressionType, MetadataVersion};
    use arrow_schema::{UnionFields, UnionMode};
    use builder::MapBuilder;
    use std::collections::HashMap;

    use super::*;

    #[test]
    /// ensure only the batch's used data (not the allocated data) is sent
    /// <https://github.com/apache/arrow-rs/issues/208>
    fn test_encode_flight_data() {
        // use 8-byte alignment - default alignment is 64 which produces bigger ipc data
        let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
        let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);

        let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c1) as ArrayRef)])
            .expect("cannot create record batch");
        let schema = batch.schema_ref();

        let (_, baseline_flight_batch) = make_flight_data(&batch, &options);

        let big_batch = batch.slice(0, batch.num_rows() - 1);
        let optimized_big_batch =
            hydrate_dictionaries(&big_batch, Arc::clone(schema)).expect("failed to optimize");
        let (_, optimized_big_flight_batch) = make_flight_data(&optimized_big_batch, &options);

        assert_eq!(
            baseline_flight_batch.data_body.len(),
            optimized_big_flight_batch.data_body.len()
        );

        let small_batch = batch.slice(0, 1);
        let optimized_small_batch =
            hydrate_dictionaries(&small_batch, Arc::clone(schema)).expect("failed to optimize");
        let (_, optimized_small_flight_batch) = make_flight_data(&optimized_small_batch, &options);

        assert!(
            baseline_flight_batch.data_body.len() > optimized_small_flight_batch.data_body.len()
        );
    }

    #[tokio::test]
    async fn test_dictionary_hydration() {
        let arr1: DictionaryArray<UInt16Type> = vec!["a", "a", "b"].into_iter().collect();
        let arr2: DictionaryArray<UInt16Type> = vec!["c", "c", "d"].into_iter().collect();

        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
            "dict",
            DataType::UInt16,
            DataType::Utf8,
            false,
        )]));
        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);
        let mut decoder = FlightDataDecoder::new(encoder);
        let expected_schema = Schema::new(vec![Field::new("dict", DataType::Utf8, false)]);
        let expected_schema = Arc::new(expected_schema);
        let mut expected_arrays = vec![
            StringArray::from(vec!["a", "a", "b"]),
            StringArray::from(vec!["c", "c", "d"]),
        ]
        .into_iter();
        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let actual_array = b.column_by_name("dict").unwrap();
                    let actual_array = downcast_array::<StringArray>(actual_array);

                    assert_eq!(actual_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_dictionary_resend() {
        let arr1: DictionaryArray<UInt16Type> = vec!["a", "a", "b"].into_iter().collect();
        let arr2: DictionaryArray<UInt16Type> = vec!["c", "c", "d"].into_iter().collect();

        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
            "dict",
            DataType::UInt16,
            DataType::Utf8,
            false,
        )]));
        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_compression_round_trip() {
        // Round trip a batch through Flight with IPC body compression enabled. This exercises
        // the compressed `IpcDataGenerator::encode` path (per-buffer codec output), which the
        // uncompressed Flight tests and the writer-based compression tests do not cover.
        let ints = Int32Array::from_iter_values((0..1024).map(|i| i % 8));
        let strings = StringArray::from_iter_values((0..1024).map(|i| format!("value-{}", i % 8)));
        let batch = RecordBatch::try_from_iter(vec![
            ("ints", Arc::new(ints) as ArrayRef),
            ("strings", Arc::new(strings) as ArrayRef),
        ])
        .unwrap();

        for compression in [CompressionType::LZ4_FRAME, CompressionType::ZSTD] {
            let options = IpcWriteOptions::default()
                .try_with_compression(Some(compression))
                .unwrap();
            verify_flight_round_trip_with_options(vec![batch.clone()], options).await;
        }
    }

    #[tokio::test]
    async fn test_dictionary_hydration_known_schema() {
        let arr1: DictionaryArray<UInt16Type> = vec!["a", "a", "b"].into_iter().collect();
        let arr2: DictionaryArray<UInt16Type> = vec!["c", "c", "d"].into_iter().collect();

        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
            "dict",
            DataType::UInt16,
            DataType::Utf8,
            false,
        )]));
        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default()
            .with_schema(schema)
            .build(stream);
        let expected_schema =
            Arc::new(Schema::new(vec![Field::new("dict", DataType::Utf8, false)]));
        assert_eq!(Some(expected_schema), encoder.known_schema())
    }

    #[tokio::test]
    async fn test_dictionary_resend_known_schema() {
        let arr1: DictionaryArray<UInt16Type> = vec!["a", "a", "b"].into_iter().collect();
        let arr2: DictionaryArray<UInt16Type> = vec!["c", "c", "d"].into_iter().collect();

        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
            "dict",
            DataType::UInt16,
            DataType::Utf8,
            false,
        )]));
        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default()
            .with_dictionary_handling(DictionaryHandling::Resend)
            .with_schema(schema.clone())
            .build(stream);
        assert_eq!(Some(schema), encoder.known_schema())
    }

    #[tokio::test]
    async fn test_multiple_dictionaries_resend() {
        // Create a schema with two dictionary fields that have the same dict ID
        let schema = Arc::new(Schema::new(vec![
            Field::new_dictionary("dict_1", DataType::UInt16, DataType::Utf8, false),
            Field::new_dictionary("dict_2", DataType::UInt16, DataType::Utf8, false),
        ]));

        let arr_one_1: Arc<DictionaryArray<UInt16Type>> =
            Arc::new(vec!["a", "a", "b"].into_iter().collect());
        let arr_one_2: Arc<DictionaryArray<UInt16Type>> =
            Arc::new(vec!["c", "c", "d"].into_iter().collect());
        let arr_two_1: Arc<DictionaryArray<UInt16Type>> =
            Arc::new(vec!["b", "a", "c"].into_iter().collect());
        let arr_two_2: Arc<DictionaryArray<UInt16Type>> =
            Arc::new(vec!["k", "d", "e"].into_iter().collect());
        let batch1 =
            RecordBatch::try_new(schema.clone(), vec![arr_one_1.clone(), arr_one_2.clone()])
                .unwrap();
        let batch2 =
            RecordBatch::try_new(schema.clone(), vec![arr_two_1.clone(), arr_two_2.clone()])
                .unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_list_hydration() {
        let mut builder = ListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);

        let arr1 = builder.finish();

        builder.append_value(vec![Some("c"), None, Some("d")]);

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);

        let mut decoder = FlightDataDecoder::new(encoder);
        let expected_schema = Schema::new(vec![Field::new_list(
            "dict_list",
            Field::new_list_field(DataType::Utf8, true),
            true,
        )]);

        let expected_schema = Arc::new(expected_schema);

        let mut expected_arrays = vec![
            StringArray::from_iter(vec![Some("a"), None, Some("b")]),
            StringArray::from_iter(vec![Some("c"), None, Some("d")]),
        ]
        .into_iter();

        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let list_array =
                        downcast_array::<ListArray>(b.column_by_name("dict_list").unwrap());
                    let elem_array = downcast_array::<StringArray>(list_array.value(0).as_ref());

                    assert_eq!(elem_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_dictionary_list_resend() {
        let mut builder = ListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);

        let arr1 = builder.finish();

        builder.append_value(vec![Some("c"), None, Some("d")]);

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_struct_hydration() {
        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let mut struct_builder = StructBuilder::new(
            struct_fields.clone(),
            vec![Box::new(builder::ListBuilder::new(
                StringDictionaryBuilder::<UInt16Type>::new(),
            ))],
        );

        struct_builder
            .field_builder::<ListBuilder<GenericByteDictionaryBuilder<UInt16Type,GenericStringType<i32>>>>(0)
            .unwrap()
            .append_value(vec![Some("a"), None, Some("b")]);

        struct_builder.append(true);

        let arr1 = struct_builder.finish();

        struct_builder
            .field_builder::<ListBuilder<GenericByteDictionaryBuilder<UInt16Type,GenericStringType<i32>>>>(0)
            .unwrap()
            .append_value(vec![Some("c"), None, Some("d")]);
        struct_builder.append(true);

        let arr2 = struct_builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_struct(
            "struct",
            struct_fields,
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);

        let mut decoder = FlightDataDecoder::new(encoder);
        let expected_schema = Schema::new(vec![Field::new_struct(
            "struct",
            vec![Field::new_list(
                "dict_list",
                Field::new_list_field(DataType::Utf8, true),
                true,
            )],
            true,
        )]);

        let expected_schema = Arc::new(expected_schema);

        let mut expected_arrays = vec![
            StringArray::from_iter(vec![Some("a"), None, Some("b")]),
            StringArray::from_iter(vec![Some("c"), None, Some("d")]),
        ]
        .into_iter();

        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let struct_array =
                        downcast_array::<StructArray>(b.column_by_name("struct").unwrap());
                    let list_array = downcast_array::<ListArray>(struct_array.column(0));

                    let elem_array = downcast_array::<StringArray>(list_array.value(0).as_ref());

                    assert_eq!(elem_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_dictionary_struct_resend() {
        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let mut struct_builder = StructBuilder::new(
            struct_fields.clone(),
            vec![Box::new(builder::ListBuilder::new(
                StringDictionaryBuilder::<UInt16Type>::new(),
            ))],
        );

        struct_builder.field_builder::<ListBuilder<GenericByteDictionaryBuilder<UInt16Type,GenericStringType<i32>>>>(0)
            .unwrap()
            .append_value(vec![Some("a"), None, Some("b")]);
        struct_builder.append(true);

        let arr1 = struct_builder.finish();

        struct_builder.field_builder::<ListBuilder<GenericByteDictionaryBuilder<UInt16Type,GenericStringType<i32>>>>(0)
            .unwrap()
            .append_value(vec![Some("c"), None, Some("d")]);
        struct_builder.append(true);

        let arr2 = struct_builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_struct(
            "struct",
            struct_fields,
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_union_hydration() {
        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let union_fields = [
            (
                0,
                Arc::new(Field::new_list(
                    "dict_list",
                    Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
                    true,
                )),
            ),
            (
                1,
                Arc::new(Field::new_struct("struct", struct_fields.clone(), true)),
            ),
            (2, Arc::new(Field::new("string", DataType::Utf8, true))),
        ]
        .into_iter()
        .collect::<UnionFields>();

        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let mut builder = builder::ListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);

        let arr1 = builder.finish();

        let type_id_buffer = std::iter::once(0).collect::<ScalarBuffer<i8>>();
        let arr1 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                Arc::new(arr1) as Arc<dyn Array>,
                new_null_array(union_fields.iter().nth(1).unwrap().1.data_type(), 1),
                new_null_array(union_fields.iter().nth(2).unwrap().1.data_type(), 1),
            ],
        )
        .unwrap();

        builder.append_value(vec![Some("c"), None, Some("d")]);

        let arr2 = Arc::new(builder.finish());
        let arr2 = StructArray::new(struct_fields.clone().into(), vec![arr2], None);

        let type_id_buffer = std::iter::once(1).collect::<ScalarBuffer<i8>>();
        let arr2 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                new_null_array(union_fields.iter().next().unwrap().1.data_type(), 1),
                Arc::new(arr2),
                new_null_array(union_fields.iter().nth(2).unwrap().1.data_type(), 1),
            ],
        )
        .unwrap();

        let type_id_buffer = std::iter::once(2).collect::<ScalarBuffer<i8>>();
        let arr3 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                new_null_array(union_fields.iter().next().unwrap().1.data_type(), 1),
                new_null_array(union_fields.iter().nth(1).unwrap().1.data_type(), 1),
                Arc::new(StringArray::from(vec!["e"])),
            ],
        )
        .unwrap();

        let (type_ids, union_fields): (Vec<_>, Vec<_>) = union_fields
            .iter()
            .map(|(type_id, field_ref)| (type_id, (*Arc::clone(field_ref)).clone()))
            .unzip();
        let schema = Arc::new(Schema::new(vec![Field::new_union(
            "union",
            type_ids.clone(),
            union_fields.clone(),
            UnionMode::Sparse,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();
        let batch3 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr3)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2), Ok(batch3)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);

        let mut decoder = FlightDataDecoder::new(encoder);

        let hydrated_struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_list_field(DataType::Utf8, true),
            true,
        )];

        let hydrated_union_fields = vec![
            Field::new_list(
                "dict_list",
                Field::new_list_field(DataType::Utf8, true),
                true,
            ),
            Field::new_struct("struct", hydrated_struct_fields.clone(), true),
            Field::new("string", DataType::Utf8, true),
        ];

        let expected_schema = Schema::new(vec![Field::new_union(
            "union",
            type_ids.clone(),
            hydrated_union_fields,
            UnionMode::Sparse,
        )]);

        let expected_schema = Arc::new(expected_schema);

        let mut expected_arrays = vec![
            StringArray::from_iter(vec![Some("a"), None, Some("b")]),
            StringArray::from_iter(vec![Some("c"), None, Some("d")]),
            StringArray::from(vec!["e"]),
        ]
        .into_iter();

        let mut batch = 0;
        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let union_arr =
                        downcast_array::<UnionArray>(b.column_by_name("union").unwrap());

                    let elem_array = match batch {
                        0 => {
                            let list_array = downcast_array::<ListArray>(union_arr.child(0));
                            downcast_array::<StringArray>(list_array.value(0).as_ref())
                        }
                        1 => {
                            let struct_array = downcast_array::<StructArray>(union_arr.child(1));
                            let list_array = downcast_array::<ListArray>(struct_array.column(0));

                            downcast_array::<StringArray>(list_array.value(0).as_ref())
                        }
                        _ => downcast_array::<StringArray>(union_arr.child(2)),
                    };

                    batch += 1;

                    assert_eq!(elem_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_dictionary_union_resend() {
        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let union_fields = [
            (
                0,
                Arc::new(Field::new_list(
                    "dict_list",
                    Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
                    true,
                )),
            ),
            (
                1,
                Arc::new(Field::new_struct("struct", struct_fields.clone(), true)),
            ),
            (2, Arc::new(Field::new("string", DataType::Utf8, true))),
        ]
        .into_iter()
        .collect::<UnionFields>();

        let mut field_types = union_fields.iter().map(|(_, field)| field.data_type());
        let dict_list_ty = field_types.next().unwrap();
        let struct_ty = field_types.next().unwrap();
        let string_ty = field_types.next().unwrap();

        let struct_fields = vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )];

        let mut builder = builder::ListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);

        let arr1 = builder.finish();

        let type_id_buffer = std::iter::once(0).collect::<ScalarBuffer<i8>>();
        let arr1 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                Arc::new(arr1),
                new_null_array(struct_ty, 1),
                new_null_array(string_ty, 1),
            ],
        )
        .unwrap();

        builder.append_value(vec![Some("c"), None, Some("d")]);

        let arr2 = Arc::new(builder.finish());
        let arr2 = StructArray::new(struct_fields.clone().into(), vec![arr2], None);

        let type_id_buffer = std::iter::once(1).collect::<ScalarBuffer<i8>>();
        let arr2 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                new_null_array(dict_list_ty, 1),
                Arc::new(arr2),
                new_null_array(string_ty, 1),
            ],
        )
        .unwrap();

        let type_id_buffer = std::iter::once(2).collect::<ScalarBuffer<i8>>();
        let arr3 = UnionArray::try_new(
            union_fields.clone(),
            type_id_buffer,
            None,
            vec![
                new_null_array(dict_list_ty, 1),
                new_null_array(struct_ty, 1),
                Arc::new(StringArray::from(vec!["e"])),
            ],
        )
        .unwrap();

        let (type_ids, union_fields): (Vec<_>, Vec<_>) = union_fields
            .iter()
            .map(|(type_id, field_ref)| (type_id, (*Arc::clone(field_ref)).clone()))
            .unzip();
        let schema = Arc::new(Schema::new(vec![Field::new_union(
            "union",
            type_ids.clone(),
            union_fields.clone(),
            UnionMode::Sparse,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();
        let batch3 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr3)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2, batch3]).await;
    }

    #[tokio::test]
    async fn test_dictionary_map_hydration() {
        let mut builder = MapBuilder::new(
            None,
            StringDictionaryBuilder::<UInt16Type>::new(),
            StringDictionaryBuilder::<UInt16Type>::new(),
        );

        // {"k1":"a","k2":null,"k3":"b"}
        builder.keys().append_value("k1");
        builder.values().append_value("a");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("b");
        builder.append(true).unwrap();

        let arr1 = builder.finish();

        // {"k1":"c","k2":null,"k3":"d"}
        builder.keys().append_value("k1");
        builder.values().append_value("c");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("d");
        builder.append(true).unwrap();

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_map(
            "dict_map",
            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
            Field::new_dictionary(
                Field::MAP_KEY_FIELD_DEFAULT_NAME,
                DataType::UInt16,
                DataType::Utf8,
                false,
            ),
            Field::new_dictionary(
                Field::MAP_VALUE_FIELD_DEFAULT_NAME,
                DataType::UInt16,
                DataType::Utf8,
                true,
            ),
            false,
            false,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);

        let mut decoder = FlightDataDecoder::new(encoder);
        let expected_schema = Schema::new(vec![Field::new_map(
            "dict_map",
            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
            false,
            false,
        )]);

        let expected_schema = Arc::new(expected_schema);

        // array without dictionary fields
        let arr1 = MapArray::from_vec_of_maps::<StringArray, StringArray, _, _>(
            vec![Some(vec![
                ("k1", Some("a")),
                ("k2", None),
                ("k3", Some("b")),
            ])],
            false,
        );

        let arr2 = MapArray::from_vec_of_maps::<StringArray, StringArray, _, _>(
            vec![Some(vec![
                ("k1", Some("c")),
                ("k2", None),
                ("k3", Some("d")),
            ])],
            false,
        );

        let mut expected_arrays = vec![arr1, arr2].into_iter();

        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let map_array =
                        downcast_array::<MapArray>(b.column_by_name("dict_map").unwrap());

                    assert_eq!(map_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]
    async fn test_dictionary_map_resend() {
        let mut builder = MapBuilder::new(
            None,
            StringDictionaryBuilder::<UInt16Type>::new(),
            StringDictionaryBuilder::<UInt16Type>::new(),
        );

        // {"k1":"a","k2":null,"k3":"b"}
        builder.keys().append_value("k1");
        builder.values().append_value("a");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("b");
        builder.append(true).unwrap();

        let arr1 = builder.finish();

        // {"k1":"c","k2":null,"k3":"d"}
        builder.keys().append_value("k1");
        builder.values().append_value("c");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("d");
        builder.append(true).unwrap();

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_map(
            "dict_map",
            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
            Field::new_dictionary(
                Field::MAP_KEY_FIELD_DEFAULT_NAME,
                DataType::UInt16,
                DataType::Utf8,
                false,
            ),
            Field::new_dictionary(
                Field::MAP_VALUE_FIELD_DEFAULT_NAME,
                DataType::UInt16,
                DataType::Utf8,
                true,
            ),
            false,
            false,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_ree_resend() {
        let dict_values1 = vec![Some("a"), None, Some("b")]
            .into_iter()
            .collect::<DictionaryArray<Int32Type>>();
        let run_ends1 = Int32Array::from(vec![1, 2, 3]);
        let arr1 = RunArray::try_new(&run_ends1, &dict_values1).unwrap();

        let dict_values2 = vec![Some("c"), Some("a")]
            .into_iter()
            .collect::<DictionaryArray<Int32Type>>();
        let run_ends2 = Int32Array::from(vec![1, 2]);
        let arr2 = RunArray::try_new(&run_ends2, &dict_values2).unwrap();

        let schema = Arc::new(Schema::new(vec![Field::new(
            "ree",
            arr1.data_type().clone(),
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_of_struct_of_dict_resend() {
        // Dict(Int8, Struct { dict: Dict(Int32, Utf8), int: Int32 })
        // This exercises the Dictionary branch recursing into its value type
        // before assigning its own dict_id (depth-first ordering).
        let struct_fields: Vec<Field> = vec![
            Field::new_dictionary("dict", DataType::Int32, DataType::Utf8, true),
            Field::new("int", DataType::Int32, false),
        ];

        let inner_values =
            StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
        let inner_keys = Int32Array::from_iter_values([0, 1, 2, 3, 0]);
        let inner_dict = DictionaryArray::new(inner_keys, Arc::new(inner_values));
        let int_array = Int32Array::from(vec![10, 20, 30, 40, 50]);

        let struct_array = StructArray::from(vec![
            (
                Arc::new(struct_fields[0].clone()),
                Arc::new(inner_dict) as ArrayRef,
            ),
            (
                Arc::new(struct_fields[1].clone()),
                Arc::new(int_array) as ArrayRef,
            ),
        ]);

        let outer_keys = Int8Array::from_iter_values([0, 0, 1, 2]);
        let arr1 = DictionaryArray::new(outer_keys, Arc::new(struct_array));

        let inner_values2 = StringArray::from(vec![Some("x"), Some("y")]);
        let inner_keys2 = Int32Array::from_iter_values([0, 1, 0]);
        let inner_dict2 = DictionaryArray::new(inner_keys2, Arc::new(inner_values2));
        let int_array2 = Int32Array::from(vec![100, 200, 300]);

        let struct_array2 = StructArray::from(vec![
            (
                Arc::new(struct_fields[0].clone()),
                Arc::new(inner_dict2) as ArrayRef,
            ),
            (
                Arc::new(struct_fields[1].clone()),
                Arc::new(int_array2) as ArrayRef,
            ),
        ]);

        let outer_keys2 = Int8Array::from_iter_values([0, 1]);
        let arr2 = DictionaryArray::new(outer_keys2, Arc::new(struct_array2));

        let schema = Arc::new(Schema::new(vec![Field::new(
            "dict_struct",
            arr1.data_type().clone(),
            false,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    async fn verify_dictionary_list_view_resend<O: OffsetSizeTrait>() {
        let mut builder =
            GenericListViewBuilder::<O, _>::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);
        let arr1 = builder.finish();

        builder.append_value(vec![Some("c"), None, Some("d")]);
        let arr2 = builder.finish();

        let inner = Arc::new(Field::new_dictionary(
            "item",
            DataType::UInt16,
            DataType::Utf8,
            true,
        ));
        let dt = if O::IS_LARGE {
            DataType::LargeListView(inner)
        } else {
            DataType::ListView(inner)
        };
        let schema = Arc::new(Schema::new(vec![Field::new("dict_list_view", dt, true)]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    #[tokio::test]
    async fn test_dictionary_list_view_resend() {
        verify_dictionary_list_view_resend::<i32>().await;
    }

    #[tokio::test]
    async fn test_dictionary_large_list_view_resend() {
        verify_dictionary_list_view_resend::<i64>().await;
    }

    #[tokio::test]
    async fn test_dictionary_fixed_size_list_resend() {
        let mut builder =
            FixedSizeListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new(), 2);

        builder.values().append_value("a");
        builder.values().append_value("b");
        builder.append(true);
        let arr1 = builder.finish();

        builder.values().append_value("c");
        builder.values().append_value("d");
        builder.append(true);
        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_fixed_size_list(
            "dict_fsl",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            2,
            true,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema, vec![Arc::new(arr2)]).unwrap();

        verify_flight_round_trip(vec![batch1, batch2]).await;
    }

    async fn verify_flight_round_trip(batches: Vec<RecordBatch>) {
        verify_flight_round_trip_with_options(batches, IpcWriteOptions::default()).await;
    }

    /// Encode `batches` through a [`FlightDataEncoderBuilder`] using `options`, decode them
    /// again, and assert the decoded batches match the originals.
    async fn verify_flight_round_trip_with_options(
        batches: Vec<RecordBatch>,
        options: IpcWriteOptions,
    ) {
        let expected_schema = batches.first().unwrap().schema();

        let encoder = FlightDataEncoderBuilder::default()
            .with_options(options)
            .with_dictionary_handling(DictionaryHandling::Resend)
            .build(futures::stream::iter(batches.clone().into_iter().map(Ok)));

        let mut expected_batches = batches.into_iter();

        let mut decoder = FlightDataDecoder::new(encoder);
        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    let expected_batch = expected_batches.next().unwrap();
                    assert_eq!(b, expected_batch);
                }
            }
        }
    }

    #[test]
    fn test_schema_metadata_encoded() {
        let schema = Schema::new(vec![Field::new("data", DataType::Int32, false)]).with_metadata(
            HashMap::from([("some_key".to_owned(), "some_value".to_owned())]),
        );

        let mut dictionary_tracker = DictionaryTracker::new(false);

        let got = prepare_schema_for_flight(&schema, &mut dictionary_tracker, false);
        assert!(got.metadata().contains_key("some_key"));
    }

    #[tokio::test]
    async fn test_list_schema_round_trip() {
        let list_item = Field::new("list_item", DataType::Int32, false).with_metadata(
            HashMap::from([("level".to_owned(), "list_item".to_owned())]),
        );
        let mut list_builder = ListBuilder::new(Int32Builder::new()).with_field(list_item);
        list_builder.append_value([Some(1), Some(2)]);
        list_builder.append_value([Some(3)]);
        let list = Arc::new(list_builder.finish()) as ArrayRef;
        let list_field = Field::new("list", list.data_type().clone(), false)
            .with_metadata(HashMap::from([("level".to_owned(), "list".to_owned())]));

        let nested_item = Field::new("nested_item", DataType::Int32, true).with_metadata(
            HashMap::from([("level".to_owned(), "nested_item".to_owned())]),
        );
        let nested_list_builder =
            ListBuilder::new(Int32Builder::new()).with_field(nested_item.clone());
        let nested_list = Field::new_list("nested_list", nested_item, false).with_metadata(
            HashMap::from([("level".to_owned(), "nested_list".to_owned())]),
        );
        let mut large_list_builder =
            LargeListBuilder::new(nested_list_builder).with_field(nested_list);
        large_list_builder.values().append_value([Some(4), None]);
        large_list_builder.values().append_value([Some(5)]);
        large_list_builder.append(true);
        large_list_builder.append(false);
        let large_list = Arc::new(large_list_builder.finish()) as ArrayRef;
        let large_list_field =
            Field::new("large_list", large_list.data_type().clone(), true).with_metadata(
                HashMap::from([("level".to_owned(), "large_list".to_owned())]),
            );

        let fixed_size_item = Field::new("fixed_size_item", DataType::Int32, true).with_metadata(
            HashMap::from([("level".to_owned(), "fixed_size_item".to_owned())]),
        );
        let mut fixed_size_list_builder =
            FixedSizeListBuilder::new(Int32Builder::new(), 2).with_field(fixed_size_item);
        fixed_size_list_builder.values().append_value(6);
        fixed_size_list_builder.values().append_null();
        fixed_size_list_builder.append(true);
        fixed_size_list_builder.values().append_value(7);
        fixed_size_list_builder.values().append_value(8);
        fixed_size_list_builder.append(true);
        let fixed_size_list = Arc::new(fixed_size_list_builder.finish()) as ArrayRef;
        let fixed_size_list_field = Field::new(
            "fixed_size_list",
            fixed_size_list.data_type().clone(),
            false,
        )
        .with_metadata(HashMap::from([(
            "level".to_owned(),
            "fixed_size_list".to_owned(),
        )]));

        let schema = Arc::new(
            Schema::new(vec![list_field, large_list_field, fixed_size_list_field])
                .with_metadata(HashMap::from([("level".to_owned(), "schema".to_owned())])),
        );
        let batch = RecordBatch::try_new(schema, vec![list, large_list, fixed_size_list]).unwrap();

        verify_flight_round_trip(vec![batch]).await;
    }

    #[test]
    fn test_encode_no_column_batch() {
        let batch = RecordBatch::try_new_with_options(
            Arc::new(Schema::empty()),
            vec![],
            &RecordBatchOptions::new().with_row_count(Some(10)),
        )
        .expect("cannot create record batch");

        hydrate_dictionaries(&batch, batch.schema()).expect("failed to optimize");
    }

    fn make_flight_data(
        batch: &RecordBatch,
        options: &IpcWriteOptions,
    ) -> (Vec<FlightData>, FlightData) {
        flight_data_from_arrow_batch(batch, options)
    }

    fn flight_data_from_arrow_batch(
        batch: &RecordBatch,
        options: &IpcWriteOptions,
    ) -> (Vec<FlightData>, FlightData) {
        let data_gen = IpcDataGenerator::default();
        let mut dictionary_tracker = DictionaryTracker::new(false);
        let mut ipc_write_context = IpcWriteContext::default();

        let (encoded_dictionaries, encoded_batch) = data_gen
            .encode(
                batch,
                &mut dictionary_tracker,
                options,
                &mut ipc_write_context,
            )
            .expect("DictionaryTracker configured above to not error on replacement");

        let flight_dictionaries = encoded_dictionaries.into_iter().map(Into::into).collect();
        let flight_batch = encoded_batch.into();

        (flight_dictionaries, flight_batch)
    }

    #[test]
    fn test_split_batch_for_grpc_response() {
        let max_flight_data_size = 1024;

        // no split
        let c = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]);
        let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c) as ArrayRef)])
            .expect("cannot create record batch");
        let split: Vec<_> = split_batch_for_grpc_response(batch.clone(), max_flight_data_size);
        assert_eq!(split.len(), 1);
        assert_eq!(batch, split[0]);

        // split once
        let n_rows = max_flight_data_size + 1;
        assert!(n_rows % 2 == 1, "should be an odd number");
        let c = UInt8Array::from((0..n_rows).map(|i| (i % 256) as u8).collect::<Vec<_>>());
        let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c) as ArrayRef)])
            .expect("cannot create record batch");
        let split: Vec<_> = split_batch_for_grpc_response(batch.clone(), max_flight_data_size);
        assert_eq!(split.len(), 3);
        assert_eq!(
            split.iter().map(|batch| batch.num_rows()).sum::<usize>(),
            n_rows
        );
        let a = pretty_format_batches(&split).unwrap().to_string();
        let b = pretty_format_batches(&[batch]).unwrap().to_string();
        assert_eq!(a, b);
    }

    #[test]
    fn test_split_batch_for_grpc_response_sizes() {
        // 2000 8 byte entries into 2k pieces: 8 chunks of 250 rows
        verify_split(2000, 2 * 1024, vec![250, 250, 250, 250, 250, 250, 250, 250]);

        // 2000 8 byte entries into 4k pieces: 4 chunks of 500 rows
        verify_split(2000, 4 * 1024, vec![500, 500, 500, 500]);

        // 2023 8 byte entries into 3k pieces does not divide evenly
        verify_split(2023, 3 * 1024, vec![337, 337, 337, 337, 337, 337, 1]);

        // 10 8 byte entries into 1 byte pieces means each rows gets its own
        verify_split(10, 1, vec![1, 1, 1, 1, 1, 1, 1, 1, 1, 1]);

        // 10 8 byte entries into 1k byte pieces means one piece
        verify_split(10, 1024, vec![10]);
    }

    /// Creates a UInt64Array of 8 byte integers with input_rows rows
    /// `max_flight_data_size_bytes` pieces and verifies the row counts in
    /// those pieces
    fn verify_split(
        num_input_rows: u64,
        max_flight_data_size_bytes: usize,
        expected_sizes: Vec<usize>,
    ) {
        let array: UInt64Array = (0..num_input_rows).collect();

        let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(array) as ArrayRef)])
            .expect("cannot create record batch");

        let input_rows = batch.num_rows();

        let split: Vec<_> =
            split_batch_for_grpc_response(batch.clone(), max_flight_data_size_bytes);
        let sizes: Vec<_> = split.iter().map(RecordBatch::num_rows).collect();
        let output_rows: usize = sizes.iter().sum();

        assert_eq!(sizes, expected_sizes, "mismatch for {batch:?}");
        assert_eq!(input_rows, output_rows, "mismatch for {batch:?}");
    }

    // test sending record batches
    // test sending record batches with multiple different dictionaries

    #[tokio::test]
    async fn flight_data_size_even() {
        let s1 = StringArray::from_iter_values(std::iter::repeat_n(".10 bytes.", 1024));
        let i1 = Int16Array::from_iter_values(0..1024);
        let s2 = StringArray::from_iter_values(std::iter::repeat_n("6bytes", 1024));
        let i2 = Int64Array::from_iter_values(0..1024);

        let batch = RecordBatch::try_from_iter(vec![
            ("s1", Arc::new(s1) as _),
            ("i1", Arc::new(i1) as _),
            ("s2", Arc::new(s2) as _),
            ("i2", Arc::new(i2) as _),
        ])
        .unwrap();

        verify_encoded_split(batch, 88).await;
    }

    #[tokio::test]
    async fn flight_data_size_uneven_variable_lengths() {
        // each row has a longer string than the last with increasing lengths 0 --> 1024
        let array = StringArray::from_iter_values((0..1024).map(|i| "*".repeat(i)));
        let batch = RecordBatch::try_from_iter(vec![("data", Arc::new(array) as _)]).unwrap();

        // overage is much higher than ideal
        // https://github.com/apache/arrow-rs/issues/3478
        verify_encoded_split(batch, 4304).await;
    }

    #[tokio::test]
    async fn flight_data_size_large_row() {
        // batch with individual that can each exceed the batch size
        let array1 = StringArray::from_iter_values(vec![
            "*".repeat(500),
            "*".repeat(500),
            "*".repeat(500),
            "*".repeat(500),
        ]);
        let array2 = StringArray::from_iter_values(vec![
            "*".to_string(),
            "*".repeat(1000),
            "*".repeat(2000),
            "*".repeat(4000),
        ]);

        let array3 = StringArray::from_iter_values(vec![
            "*".to_string(),
            "*".to_string(),
            "*".repeat(1000),
            "*".repeat(2000),
        ]);

        let batch = RecordBatch::try_from_iter(vec![
            ("a1", Arc::new(array1) as _),
            ("a2", Arc::new(array2) as _),
            ("a3", Arc::new(array3) as _),
        ])
        .unwrap();

        // 5k over limit (which is 2x larger than limit of 5k)
        // overage is much higher than ideal
        // https://github.com/apache/arrow-rs/issues/3478
        verify_encoded_split(batch, 5784).await;
    }

    #[tokio::test]
    async fn flight_data_size_string_dictionary() {
        // Small dictionary (only 2 distinct values ==> 2 entries in dictionary)
        let array: DictionaryArray<Int32Type> = (1..1024)
            .map(|i| match i % 3 {
                0 => Some("value0"),
                1 => Some("value1"),
                _ => None,
            })
            .collect();

        let batch = RecordBatch::try_from_iter(vec![("a1", Arc::new(array) as _)]).unwrap();

        verify_encoded_split(batch, 56).await;
    }

    #[tokio::test]
    async fn flight_data_size_large_dictionary() {
        // large dictionary (all distinct values ==> 1024 entries in dictionary)
        let values: Vec<_> = (1..1024).map(|i| "**".repeat(i)).collect();

        let array: DictionaryArray<Int32Type> = values.iter().map(|s| Some(s.as_str())).collect();

        let batch = RecordBatch::try_from_iter(vec![("a1", Arc::new(array) as _)]).unwrap();

        // overage is much higher than ideal
        // https://github.com/apache/arrow-rs/issues/3478
        verify_encoded_split(batch, 3328).await;
    }

    #[tokio::test]
    async fn flight_data_size_large_dictionary_repeated_non_uniform() {
        // large dictionary (1024 distinct values) that are used throughout the array
        let values = StringArray::from_iter_values((0..1024).map(|i| "******".repeat(i)));
        let keys = Int32Array::from_iter_values((0..3000).map(|i| (3000 - i) % 1024));
        let array = DictionaryArray::new(keys, Arc::new(values));

        let batch = RecordBatch::try_from_iter(vec![("a1", Arc::new(array) as _)]).unwrap();

        // overage is much higher than ideal
        // https://github.com/apache/arrow-rs/issues/3478
        verify_encoded_split(batch, 5280).await;
    }

    #[tokio::test]
    async fn flight_data_size_multiple_dictionaries() {
        // high cardinality
        let values1: Vec<_> = (1..1024).map(|i| "**".repeat(i)).collect();
        // highish cardinality
        let values2: Vec<_> = (1..1024).map(|i| "**".repeat(i % 10)).collect();
        // medium cardinality
        let values3: Vec<_> = (1..1024).map(|i| "**".repeat(i % 100)).collect();

        let array1: DictionaryArray<Int32Type> = values1.iter().map(|s| Some(s.as_str())).collect();
        let array2: DictionaryArray<Int32Type> = values2.iter().map(|s| Some(s.as_str())).collect();
        let array3: DictionaryArray<Int32Type> = values3.iter().map(|s| Some(s.as_str())).collect();

        let batch = RecordBatch::try_from_iter(vec![
            ("a1", Arc::new(array1) as _),
            ("a2", Arc::new(array2) as _),
            ("a3", Arc::new(array3) as _),
        ])
        .unwrap();

        // overage is much higher than ideal
        // https://github.com/apache/arrow-rs/issues/3478
        verify_encoded_split(batch, 4112).await;
    }

    /// Return size, in memory of flight data
    fn flight_data_size(d: &FlightData) -> usize {
        let flight_descriptor_size = d
            .flight_descriptor
            .as_ref()
            .map(|descriptor| {
                let path_len: usize = descriptor.path.iter().map(|p| p.len()).sum();

                std::mem::size_of_val(descriptor) + descriptor.cmd.len() + path_len
            })
            .unwrap_or(0);

        flight_descriptor_size + d.app_metadata.len() + d.data_body.len() + d.data_header.len()
    }

    /// Coverage for <https://github.com/apache/arrow-rs/issues/3478>
    ///
    /// Encodes the specified batch using several values of
    /// `max_flight_data_size` between 1K to 5K and ensures that the
    /// resulting size of the flight data stays within the limit
    /// + `allowed_overage`
    ///
    /// `allowed_overage` is how far off the actual data encoding is
    /// from the target limit that was set. It is an improvement when
    /// the allowed_overage decreses.
    ///
    /// Note this overhead will likely always be greater than zero to
    /// account for encoding overhead such as IPC headers and padding.
    ///
    ///
    async fn verify_encoded_split(batch: RecordBatch, allowed_overage: usize) {
        let num_rows = batch.num_rows();

        // Track the overall required maximum overage
        let mut max_overage_seen = 0;

        for max_flight_data_size in [1024, 2021, 5000] {
            println!("Encoding {num_rows} with a maximum size of {max_flight_data_size}");

            let mut stream = FlightDataEncoderBuilder::new()
                .with_max_flight_data_size(max_flight_data_size)
                // use 8-byte alignment - default alignment is 64 which produces bigger ipc data
                .with_options(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap())
                .build(futures::stream::iter([Ok(batch.clone())]));

            let mut i = 0;
            while let Some(data) = stream.next().await.transpose().unwrap() {
                let actual_data_size = flight_data_size(&data);

                let actual_overage = actual_data_size.saturating_sub(max_flight_data_size);

                assert!(
                    actual_overage <= allowed_overage,
                    "encoded data[{i}]: actual size {actual_data_size}, \
                         actual_overage: {actual_overage} \
                         allowed_overage: {allowed_overage}"
                );

                i += 1;

                max_overage_seen = max_overage_seen.max(actual_overage)
            }
        }

        // ensure that the specified overage is exactly the maximum so
        // that when the splitting logic improves, the tests must be
        // updated to reflect the better logic
        assert_eq!(
            allowed_overage, max_overage_seen,
            "Specified overage was too high"
        );
    }
}