/Users/andrewlamb/Software/arrow-rs/arrow-array/src/record_batch.rs

Source
// 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.

//! A two-dimensional batch of column-oriented data with a defined
//! [schema](arrow_schema::Schema).

use crate::cast::AsArray;
use crate::{Array, ArrayRef, StructArray, new_empty_array};
use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaBuilder, SchemaRef};
use std::ops::Index;
use std::sync::Arc;

/// Trait for types that can read `RecordBatch`'s.
///
/// To create from an iterator, see [RecordBatchIterator].
pub trait RecordBatchReader: Iterator<Item = Result<RecordBatch, ArrowError>> {
    /// Returns the schema of this `RecordBatchReader`.
    ///
    /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
    /// reader should have the same schema as returned from this method.
    fn schema(&self) -> SchemaRef;
}

impl<R: RecordBatchReader + ?Sized> RecordBatchReader for Box<R> {
    fn schema(&self) -> SchemaRef {
        self.as_ref().schema()
    }
}

/// Trait for types that can write `RecordBatch`'s.
pub trait RecordBatchWriter {
    /// Write a single batch to the writer.
    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError>;

    /// Write footer or termination data, then mark the writer as done.
    fn close(self) -> Result<(), ArrowError>;
}

/// Creates an array from a literal slice of values,
/// suitable for rapid testing and development.
///
/// Example:
///
/// ```rust
///
/// use arrow_array::create_array;
///
/// let array = create_array!(Int32, [1, 2, 3, 4, 5]);
/// let array = create_array!(Utf8, [Some("a"), Some("b"), None, Some("e")]);
/// ```
/// Support for limited data types is available. The macro will return a compile error if an unsupported data type is used.
/// Presently supported data types are:
/// - `Boolean`, `Null`
/// - `Decimal32`, `Decimal64`, `Decimal128`, `Decimal256`
/// - `Float16`, `Float32`, `Float64`
/// - `Int8`, `Int16`, `Int32`, `Int64`
/// - `UInt8`, `UInt16`, `UInt32`, `UInt64`
/// - `IntervalDayTime`, `IntervalYearMonth`
/// - `Second`, `Millisecond`, `Microsecond`, `Nanosecond`
/// - `Second32`, `Millisecond32`, `Microsecond64`, `Nanosecond64`
/// - `DurationSecond`, `DurationMillisecond`, `DurationMicrosecond`, `DurationNanosecond`
/// - `TimestampSecond`, `TimestampMillisecond`, `TimestampMicrosecond`, `TimestampNanosecond`
/// - `Utf8`, `Utf8View`, `LargeUtf8`, `Binary`, `LargeBinary`
#[macro_export]
macro_rules! create_array {
    // `@from` is used for those types that have a common method `<type>::from`
    (@from Boolean) => { $crate::BooleanArray };
    (@from Int8) => { $crate::Int8Array };
    (@from Int16) => { $crate::Int16Array };
    (@from Int32) => { $crate::Int32Array };
    (@from Int64) => { $crate::Int64Array };
    (@from UInt8) => { $crate::UInt8Array };
    (@from UInt16) => { $crate::UInt16Array };
    (@from UInt32) => { $crate::UInt32Array };
    (@from UInt64) => { $crate::UInt64Array };
    (@from Float16) => { $crate::Float16Array };
    (@from Float32) => { $crate::Float32Array };
    (@from Float64) => { $crate::Float64Array };
    (@from Utf8) => { $crate::StringArray };
    (@from Utf8View) => { $crate::StringViewArray };
    (@from LargeUtf8) => { $crate::LargeStringArray };
    (@from IntervalDayTime) => { $crate::IntervalDayTimeArray };
    (@from IntervalYearMonth) => { $crate::IntervalYearMonthArray };
    (@from Second) => { $crate::TimestampSecondArray };
    (@from Millisecond) => { $crate::TimestampMillisecondArray };
    (@from Microsecond) => { $crate::TimestampMicrosecondArray };
    (@from Nanosecond) => { $crate::TimestampNanosecondArray };
    (@from Second32) => { $crate::Time32SecondArray };
    (@from Millisecond32) => { $crate::Time32MillisecondArray };
    (@from Microsecond64) => { $crate::Time64MicrosecondArray };
    (@from Nanosecond64) => { $crate::Time64Nanosecond64Array };
    (@from DurationSecond) => { $crate::DurationSecondArray };
    (@from DurationMillisecond) => { $crate::DurationMillisecondArray };
    (@from DurationMicrosecond) => { $crate::DurationMicrosecondArray };
    (@from DurationNanosecond) => { $crate::DurationNanosecondArray };
    (@from Decimal32) => { $crate::Decimal32Array };
    (@from Decimal64) => { $crate::Decimal64Array };
    (@from Decimal128) => { $crate::Decimal128Array };
    (@from Decimal256) => { $crate::Decimal256Array };
    (@from TimestampSecond) => { $crate::TimestampSecondArray };
    (@from TimestampMillisecond) => { $crate::TimestampMillisecondArray };
    (@from TimestampMicrosecond) => { $crate::TimestampMicrosecondArray };
    (@from TimestampNanosecond) => { $crate::TimestampNanosecondArray };

    (@from $ty: ident) => {
        compile_error!(concat!("Unsupported data type: ", stringify!($ty)))
    };

    (Null, $size: expr) => {
        std::sync::Arc::new($crate::NullArray::new($size))
    };

    (Binary, [$($values: expr),*]) => {
        std::sync::Arc::new($crate::BinaryArray::from_vec(vec![$($values),*]))
    };

    (LargeBinary, [$($values: expr),*]) => {
        std::sync::Arc::new($crate::LargeBinaryArray::from_vec(vec![$($values),*]))
    };

    ($ty: tt, [$($values: expr),*]) => {
        std::sync::Arc::new(<$crate::create_array!(@from $ty)>::from(vec![$($values),*]))
    };
}

/// Creates a record batch from literal slice of values, suitable for rapid
/// testing and development.
///
/// Example:
///
/// ```rust
/// use arrow_array::record_batch;
/// use arrow_schema;
///
/// let batch = record_batch!(
///     ("a", Int32, [1, 2, 3]),
///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
///     ("c", Utf8, ["alpha", "beta", "gamma"])
/// );
/// ```
/// Due to limitation of [`create_array!`] macro, support for limited data types is available.
#[macro_export]
macro_rules! record_batch {
    ($(($name: expr, $type: ident, [$($values: expr),*])),*) => {
        {
            let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![
                $(
                    arrow_schema::Field::new($name, arrow_schema::DataType::$type, true),
                )*
            ]));

            let batch = $crate::RecordBatch::try_new(
                schema,
                vec![$(
                    $crate::create_array!($type, [$($values),*]),
                )*]
            );

            batch
        }
    }
}

/// A two-dimensional batch of column-oriented data with a defined
/// [schema](arrow_schema::Schema).
///
/// A `RecordBatch` is a two-dimensional dataset of a number of
/// contiguous arrays, each the same length.
/// A record batch has a schema which must match its arrays’
/// datatypes.
///
/// Record batches are a convenient unit of work for various
/// serialization and computation functions, possibly incremental.
///
/// Use the [`record_batch!`] macro to create a [`RecordBatch`] from
/// literal slice of values, useful for rapid prototyping and testing.
///
/// Example:
/// ```rust
/// use arrow_array::record_batch;
/// let batch = record_batch!(
///     ("a", Int32, [1, 2, 3]),
///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
///     ("c", Utf8, ["alpha", "beta", "gamma"])
/// );
/// ```
#[derive(Clone, Debug, PartialEq)]
pub struct RecordBatch {
    schema: SchemaRef,
    columns: Vec<Arc<dyn Array>>,

    /// The number of rows in this RecordBatch
    ///
    /// This is stored separately from the columns to handle the case of no columns
    row_count: usize,
}

impl RecordBatch {
    /// Creates a `RecordBatch` from a schema and columns.
    ///
    /// Expects the following:
    ///
    ///  * `!columns.is_empty()`
    ///  * `schema.fields.len() == columns.len()`
    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
    ///  * `columns[i].len() == columns[j].len()`
    ///
    /// If the conditions are not met, an error is returned.
    ///
    /// # Example
    ///
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{Int32Array, RecordBatch};
    /// # use arrow_schema::{DataType, Field, Schema};
    ///
    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
    /// let schema = Schema::new(vec![
    ///     Field::new("id", DataType::Int32, false)
    /// ]);
    ///
    /// let batch = RecordBatch::try_new(
    ///     Arc::new(schema),
    ///     vec![Arc::new(id_array)]
    /// ).unwrap();
    /// ```
    pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self, ArrowError> {
        let options = RecordBatchOptions::new();
        Self::try_new_impl(schema, columns, &options)
    }

    /// Creates a `RecordBatch` from a schema and columns, without validation.
    ///
    /// See [`Self::try_new`] for the checked version.
    ///
    /// # Safety
    ///
    /// Expects the following:
    ///
    ///  * `schema.fields.len() == columns.len()`
    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
    ///  * `columns[i].len() == row_count`
    ///
    /// Note: if the schema does not match the underlying data exactly, it can lead to undefined
    /// behavior, for example, via conversion to a `StructArray`, which in turn could lead
    /// to incorrect access.
    pub unsafe fn new_unchecked(
        schema: SchemaRef,
        columns: Vec<Arc<dyn Array>>,
        row_count: usize,
    ) -> Self {
        Self {
            schema,
            columns,
            row_count,
        }
    }

    /// Creates a `RecordBatch` from a schema and columns, with additional options,
    /// such as whether to strictly validate field names.
    ///
    /// See [`RecordBatch::try_new`] for the expected conditions.
    pub fn try_new_with_options(
        schema: SchemaRef,
        columns: Vec<ArrayRef>,
        options: &RecordBatchOptions,
    ) -> Result<Self, ArrowError> {
        Self::try_new_impl(schema, columns, options)
    }

    /// Creates a new empty [`RecordBatch`].
    pub fn new_empty(schema: SchemaRef) -> Self {
        let columns = schema
            .fields()
            .iter()
            .map(|field| new_empty_array(field.data_type()))
            .collect();

        RecordBatch {
            schema,
            columns,
            row_count: 0,
        }
    }

    /// Validate the schema and columns using [`RecordBatchOptions`]. Returns an error
    /// if any validation check fails, otherwise returns the created [`Self`]
    fn try_new_impl(
        schema: SchemaRef,
        columns: Vec<ArrayRef>,
        options: &RecordBatchOptions,
    ) -> Result<Self, ArrowError> {
        // check that number of fields in schema match column length
        if schema.fields().len() != columns.len() {
            return Err(ArrowError::InvalidArgumentError(format!(
                "number of columns({}) must match number of fields({}) in schema",
                columns.len(),
                schema.fields().len(),
            )));
        }

        let row_count = options
            .row_count
            .or_else(|| columns.first()128.map128(|col| col128.len128()))
            .ok_or_else(|| {0
                ArrowError::InvalidArgumentError(
                    "must either specify a row count or at least one column".to_string(),
                )
            })?;

        for (c, f) in columns.iter()278.zip278(&schema.fields278) {
            if !f.is_nullable() && c.null_count() > 060 {
                return Err(ArrowError::InvalidArgumentError(format!(
                    "Column '{}' is declared as non-nullable but contains null values",
                    f.name()
                )));
            }
        }

        // check that all columns have the same row count
        if columns.iter()278.any278(|c| c.len() != row_count) {
            let err = match options.row_count {
                Some(_) => "all columns in a record batch must have the specified row count",
                None => "all columns in a record batch must have the same length",
            };
            return Err(ArrowError::InvalidArgumentError(err.to_string()));
        }

        // function for comparing column type and field type
        // return true if 2 types are not matched
        let type_not_match = if options.match_field_names {
            |(_, (col_type, field_type)): &(usize, (&DataType, &DataType))| col_type != field_type
        } else {
            |(_, (col_type, field_type)): &(usize, (&DataType, &DataType))| {
                !col_type.equals_datatype(field_type)
            }
        };

        // check that all columns match the schema
        let not_match = columns
            .iter()
            .zip(schema.fields().iter())
            .map278(|(col, field)| (col.data_type(), field.data_type()))
            .enumerate()
            .find(type_not_match);

        if let Some((i0, (col_type0, field_type0))) = not_match {
            return Err(ArrowError::InvalidArgumentError(format!(
                "column types must match schema types, expected {field_type} but found {col_type} at column index {i}"
            )));
        }

        Ok(RecordBatch {
            schema,
            columns,
            row_count,
        })
    }

    /// Return the schema, columns and row count of this [`RecordBatch`]
    pub fn into_parts(self) -> (SchemaRef, Vec<ArrayRef>, usize) {
        (self.schema, self.columns, self.row_count)
    }

    /// Override the schema of this [`RecordBatch`]
    ///
    /// Returns an error if `schema` is not a superset of the current schema
    /// as determined by [`Schema::contains`]
    ///
    /// See also [`Self::schema_metadata_mut`].
    pub fn with_schema(self, schema: SchemaRef) -> Result<Self, ArrowError> {
        if !schema.contains(self.schema.as_ref()) {
            return Err(ArrowError::SchemaError(format!(
                "target schema is not superset of current schema target={schema} current={}",
                self.schema
            )));
        }

        Ok(Self {
            schema,
            columns: self.columns,
            row_count: self.row_count,
        })
    }

    /// Returns the [`Schema`] of the record batch.
    pub fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    /// Returns a reference to the [`Schema`] of the record batch.
    pub fn schema_ref(&self) -> &SchemaRef {
        &self.schema
    }

    /// Mutable access to the metadata of the schema.
    ///
    /// This allows you to modify [`Schema::metadata`] of [`Self::schema`] in a convenient and fast way.
    ///
    /// Note this will clone the entire underlying `Schema` object if it is currently shared
    ///
    /// # Example
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{record_batch, RecordBatch};
    /// let mut batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
    /// // Initially, the metadata is empty
    /// assert!(batch.schema().metadata().get("key").is_none());
    /// // Insert a key-value pair into the metadata
    /// batch.schema_metadata_mut().insert("key".into(), "value".into());
    /// assert_eq!(batch.schema().metadata().get("key"), Some(&String::from("value")));
    /// ```
    pub fn schema_metadata_mut(&mut self) -> &mut std::collections::HashMap<String, String> {
        let schema = Arc::make_mut(&mut self.schema);
        &mut schema.metadata
    }

    /// Projects the schema onto the specified columns
    pub fn project(&self, indices: &[usize]) -> Result<RecordBatch, ArrowError> {
        let projected_schema = self.schema.project(indices)?;
        let batch_fields = indices
            .iter()
            .map(|f| {
                self.columns.get(*f).cloned().ok_or_else(|| {
                    ArrowError::SchemaError(format!(
                        "project index {} out of bounds, max field {}",
                        f,
                        self.columns.len()
                    ))
                })
            })
            .collect::<Result<Vec<_>, _>>()?;

        unsafe {
            // Since we're starting from a valid RecordBatch and project
            // creates a strict subset of the original, there's no need to
            // redo the validation checks in `try_new_with_options`.
            Ok(RecordBatch::new_unchecked(
                SchemaRef::new(projected_schema),
                batch_fields,
                self.row_count,
            ))
        }
    }

    /// Normalize a semi-structured [`RecordBatch`] into a flat table.
    ///
    /// Nested [`Field`]s will generate names separated by `separator`, up to a depth of `max_level`
    /// (unlimited if `None`).
    ///
    /// e.g. given a [`RecordBatch`] with schema:
    ///
    /// ```text
    ///     "foo": StructArray<"bar": Utf8>
    /// ```
    ///
    /// A separator of `"."` would generate a batch with the schema:
    ///
    /// ```text
    ///     "foo.bar": Utf8
    /// ```
    ///
    /// Note that giving a depth of `Some(0)` to `max_level` is the same as passing in `None`;
    /// it will be treated as unlimited.
    ///
    /// # Example
    ///
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{ArrayRef, Int64Array, StringArray, StructArray, RecordBatch};
    /// # use arrow_schema::{DataType, Field, Fields, Schema};
    /// #
    /// let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""]));
    /// let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)]));
    ///
    /// let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
    /// let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
    ///
    /// let a = Arc::new(StructArray::from(vec![
    ///     (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef),
    ///     (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef),
    /// ]));
    ///
    /// let schema = Schema::new(vec![
    ///     Field::new(
    ///         "a",
    ///         DataType::Struct(Fields::from(vec![animals_field, n_legs_field])),
    ///         false,
    ///     )
    /// ]);
    ///
    /// let normalized = RecordBatch::try_new(Arc::new(schema), vec![a])
    ///     .expect("valid conversion")
    ///     .normalize(".", None)
    ///     .expect("valid normalization");
    ///
    /// let expected = RecordBatch::try_from_iter_with_nullable(vec![
    ///     ("a.animals", animals.clone(), true),
    ///     ("a.n_legs", n_legs.clone(), true),
    /// ])
    /// .expect("valid conversion");
    ///
    /// assert_eq!(expected, normalized);
    /// ```
    pub fn normalize(&self, separator: &str, max_level: Option<usize>) -> Result<Self, ArrowError> {
        let max_level = match max_level.unwrap_or(usize::MAX) {
            0 => usize::MAX,
            val => val,
        };
        let mut stack: Vec<(usize, &ArrayRef, Vec<&str>, &FieldRef)> = self
            .columns
            .iter()
            .zip(self.schema.fields())
            .rev()
            .map(|(c, f)| {
                let name_vec: Vec<&str> = vec![f.name()];
                (0, c, name_vec, f)
            })
            .collect();
        let mut columns: Vec<ArrayRef> = Vec::new();
        let mut fields: Vec<FieldRef> = Vec::new();

        while let Some((depth, c, name, field_ref)) = stack.pop() {
            match field_ref.data_type() {
                DataType::Struct(ff) if depth < max_level => {
                    // Need to zip these in reverse to maintain original order
                    for (cff, fff) in c.as_struct().columns().iter().zip(ff.into_iter()).rev() {
                        let mut name = name.clone();
                        name.push(separator);
                        name.push(fff.name());
                        stack.push((depth + 1, cff, name, fff))
                    }
                }
                _ => {
                    let updated_field = Field::new(
                        name.concat(),
                        field_ref.data_type().clone(),
                        field_ref.is_nullable(),
                    );
                    columns.push(c.clone());
                    fields.push(Arc::new(updated_field));
                }
            }
        }
        RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
    }

    /// Returns the number of columns in the record batch.
    ///
    /// # Example
    ///
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{Int32Array, RecordBatch};
    /// # use arrow_schema::{DataType, Field, Schema};
    ///
    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
    /// let schema = Schema::new(vec![
    ///     Field::new("id", DataType::Int32, false)
    /// ]);
    ///
    /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)]).unwrap();
    ///
    /// assert_eq!(batch.num_columns(), 1);
    /// ```
    pub fn num_columns(&self) -> usize {
        self.columns.len()
    }

    /// Returns the number of rows in each column.
    ///
    /// # Example
    ///
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{Int32Array, RecordBatch};
    /// # use arrow_schema::{DataType, Field, Schema};
    ///
    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
    /// let schema = Schema::new(vec![
    ///     Field::new("id", DataType::Int32, false)
    /// ]);
    ///
    /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)]).unwrap();
    ///
    /// assert_eq!(batch.num_rows(), 5);
    /// ```
    pub fn num_rows(&self) -> usize {
        self.row_count
    }

    /// Get a reference to a column's array by index.
    ///
    /// # Panics
    ///
    /// Panics if `index` is outside of `0..num_columns`.
    pub fn column(&self, index: usize) -> &ArrayRef {
        &self.columns[index]
    }

    /// Get a reference to a column's array by name.
    pub fn column_by_name(&self, name: &str) -> Option<&ArrayRef> {
        self.schema()
            .column_with_name(name)
            .map(|(index, _)| &self.columns[index])
    }

    /// Get a reference to all columns in the record batch.
    pub fn columns(&self) -> &[ArrayRef] {
        &self.columns[..]
    }

    /// Remove column by index and return it.
    ///
    /// Return the `ArrayRef` if the column is removed.
    ///
    /// # Panics
    ///
    /// Panics if `index`` out of bounds.
    ///
    /// # Example
    ///
    /// ```
    /// use std::sync::Arc;
    /// use arrow_array::{BooleanArray, Int32Array, RecordBatch};
    /// use arrow_schema::{DataType, Field, Schema};
    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
    /// let bool_array = BooleanArray::from(vec![true, false, false, true, true]);
    /// let schema = Schema::new(vec![
    ///     Field::new("id", DataType::Int32, false),
    ///     Field::new("bool", DataType::Boolean, false),
    /// ]);
    ///
    /// let mut batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array), Arc::new(bool_array)]).unwrap();
    ///
    /// let removed_column = batch.remove_column(0);
    /// assert_eq!(removed_column.as_any().downcast_ref::<Int32Array>().unwrap(), &Int32Array::from(vec![1, 2, 3, 4, 5]));
    /// assert_eq!(batch.num_columns(), 1);
    /// ```
    pub fn remove_column(&mut self, index: usize) -> ArrayRef {
        let mut builder = SchemaBuilder::from(self.schema.as_ref());
        builder.remove(index);
        self.schema = Arc::new(builder.finish());
        self.columns.remove(index)
    }

    /// Return a new RecordBatch where each column is sliced
    /// according to `offset` and `length`
    ///
    /// # Panics
    ///
    /// Panics if `offset` with `length` is greater than column length.
    pub fn slice(&self, offset: usize, length: usize) -> RecordBatch {
        assert!((offset + length) <= self.num_rows());

        let columns = self
            .columns()
            .iter()
            .map77(|column| column.slice(offset, length))
            .collect();

        Self {
            schema: self.schema.clone(),
            columns,
            row_count: length,
        }
    }

    /// Create a `RecordBatch` from an iterable list of pairs of the
    /// form `(field_name, array)`, with the same requirements on
    /// fields and arrays as [`RecordBatch::try_new`]. This method is
    /// often used to create a single `RecordBatch` from arrays,
    /// e.g. for testing.
    ///
    /// The resulting schema is marked as nullable for each column if
    /// the array for that column is has any nulls. To explicitly
    /// specify nullibility, use [`RecordBatch::try_from_iter_with_nullable`]
    ///
    /// Example:
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
    ///
    /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
    /// let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
    ///
    /// let record_batch = RecordBatch::try_from_iter(vec![
    ///   ("a", a),
    ///   ("b", b),
    /// ]);
    /// ```
    /// Another way to quickly create a [`RecordBatch`] is to use the [`record_batch!`] macro,
    /// which is particularly helpful for rapid prototyping and testing.
    ///
    /// Example:
    ///
    /// ```rust
    /// use arrow_array::record_batch;
    /// let batch = record_batch!(
    ///     ("a", Int32, [1, 2, 3]),
    ///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
    ///     ("c", Utf8, ["alpha", "beta", "gamma"])
    /// );
    /// ```
    pub fn try_from_iter<I, F>(value: I) -> Result<Self, ArrowError>
    where
        I: IntoIterator<Item = (F, ArrayRef)>,
        F: AsRef<str>,
    {
        // TODO: implement `TryFrom` trait, once
        // https://github.com/rust-lang/rust/issues/50133 is no longer an
        // issue
        let iter41 = value41.into_iter41().map41(|(field_name, array)| {
            let nullable = array.null_count() > 0;
            (field_name, array, nullable)
        });

        Self::try_from_iter_with_nullable(iter)
    }

    /// Create a `RecordBatch` from an iterable list of tuples of the
    /// form `(field_name, array, nullable)`, with the same requirements on
    /// fields and arrays as [`RecordBatch::try_new`]. This method is often
    /// used to create a single `RecordBatch` from arrays, e.g. for
    /// testing.
    ///
    /// Example:
    /// ```
    /// # use std::sync::Arc;
    /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
    ///
    /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
    /// let b: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
    ///
    /// // Note neither `a` nor `b` has any actual nulls, but we mark
    /// // b an nullable
    /// let record_batch = RecordBatch::try_from_iter_with_nullable(vec![
    ///   ("a", a, false),
    ///   ("b", b, true),
    /// ]);
    /// ```
    pub fn try_from_iter_with_nullable<I, F>(value: I) -> Result<Self, ArrowError>
    where
        I: IntoIterator<Item = (F, ArrayRef, bool)>,
        F: AsRef<str>,
    {
        let iter = value.into_iter();
        let capacity = iter.size_hint().0;
        let mut schema = SchemaBuilder::with_capacity(capacity);
        let mut columns = Vec::with_capacity(capacity);

        for (field_name161, array161, nullable161) in iter {
            let field_name = field_name.as_ref();
            schema.push(Field::new(field_name, array.data_type().clone(), nullable));
            columns.push(array);
        }

        let schema = Arc::new(schema.finish());
        RecordBatch::try_new(schema, columns)
    }

    /// Returns the total number of bytes of memory occupied physically by this batch.
    ///
    /// Note that this does not always correspond to the exact memory usage of a
    /// `RecordBatch` (might overestimate), since multiple columns can share the same
    /// buffers or slices thereof, the memory used by the shared buffers might be
    /// counted multiple times.
    pub fn get_array_memory_size(&self) -> usize {
        self.columns()
            .iter()
            .map(|array| array.get_array_memory_size())
            .sum()
    }
}

/// Options that control the behaviour used when creating a [`RecordBatch`].
#[derive(Debug)]
#[non_exhaustive]
pub struct RecordBatchOptions {
    /// Match field names of structs and lists. If set to `true`, the names must match.
    pub match_field_names: bool,

    /// Optional row count, useful for specifying a row count for a RecordBatch with no columns
    pub row_count: Option<usize>,
}

impl RecordBatchOptions {
    /// Creates a new `RecordBatchOptions`
    pub fn new() -> Self {
        Self {
            match_field_names: true,
            row_count: None,
        }
    }
    /// Sets the row_count of RecordBatchOptions and returns self
    pub fn with_row_count(mut self, row_count: Option<usize>) -> Self {
        self.row_count = row_count;
        self
    }
    /// Sets the match_field_names of RecordBatchOptions and returns self
    pub fn with_match_field_names(mut self, match_field_names: bool) -> Self {
        self.match_field_names = match_field_names;
        self
    }
}
impl Default for RecordBatchOptions {
    fn default() -> Self {
        Self::new()
    }
}
impl From<StructArray> for RecordBatch {
    fn from(value: StructArray) -> Self {
        let row_count = value.len();
        let (fields, columns, nulls) = value.into_parts();
        assert_eq!(
            nulls.map(|n| n.null_count()).unwrap_or_default(),
            0,
            "Cannot convert nullable StructArray to RecordBatch, see StructArray documentation"
        );

        RecordBatch {
            schema: Arc::new(Schema::new(fields)),
            row_count,
            columns,
        }
    }
}

impl From<&StructArray> for RecordBatch {
    fn from(struct_array: &StructArray) -> Self {
        struct_array.clone().into()
    }
}

impl Index<&str> for RecordBatch {
    type Output = ArrayRef;

    /// Get a reference to a column's array by name.
    ///
    /// # Panics
    ///
    /// Panics if the name is not in the schema.
    fn index(&self, name: &str) -> &Self::Output {
        self.column_by_name(name).unwrap()
    }
}

/// Generic implementation of [RecordBatchReader] that wraps an iterator.
///
/// # Example
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray, RecordBatchIterator, RecordBatchReader};
/// #
/// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
/// let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
///
/// let record_batch = RecordBatch::try_from_iter(vec![
///   ("a", a),
///   ("b", b),
/// ]).unwrap();
///
/// let batches: Vec<RecordBatch> = vec![record_batch.clone(), record_batch.clone()];
///
/// let mut reader = RecordBatchIterator::new(batches.into_iter().map(Ok), record_batch.schema());
///
/// assert_eq!(reader.schema(), record_batch.schema());
/// assert_eq!(reader.next().unwrap().unwrap(), record_batch);
/// # assert_eq!(reader.next().unwrap().unwrap(), record_batch);
/// # assert!(reader.next().is_none());
/// ```
pub struct RecordBatchIterator<I>
where
    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
{
    inner: I::IntoIter,
    inner_schema: SchemaRef,
}

impl<I> RecordBatchIterator<I>
where
    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
{
    /// Create a new [RecordBatchIterator].
    ///
    /// If `iter` is an infallible iterator, use `.map(Ok)`.
    pub fn new(iter: I, schema: SchemaRef) -> Self {
        Self {
            inner: iter.into_iter(),
            inner_schema: schema,
        }
    }
}

impl<I> Iterator for RecordBatchIterator<I>
where
    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

impl<I> RecordBatchReader for RecordBatchIterator<I>
where
    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
{
    fn schema(&self) -> SchemaRef {
        self.inner_schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        BooleanArray, Int8Array, Int32Array, Int64Array, ListArray, StringArray, StringViewArray,
    };
    use arrow_buffer::{Buffer, ToByteSlice};
    use arrow_data::{ArrayData, ArrayDataBuilder};
    use arrow_schema::Fields;
    use std::collections::HashMap;

    #[test]
    fn create_record_batch() {
        let schema = Schema::new(vec![
            Field::new("a", DataType::Int32, false),
            Field::new("b", DataType::Utf8, false),
        ]);

        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let b = StringArray::from(vec!["a", "b", "c", "d", "e"]);

        let record_batch =
            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
        check_batch(record_batch, 5)
    }

    #[test]
    fn create_string_view_record_batch() {
        let schema = Schema::new(vec![
            Field::new("a", DataType::Int32, false),
            Field::new("b", DataType::Utf8View, false),
        ]);

        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let b = StringViewArray::from(vec!["a", "b", "c", "d", "e"]);

        let record_batch =
            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();

        assert_eq!(5, record_batch.num_rows());
        assert_eq!(2, record_batch.num_columns());
        assert_eq!(&DataType::Int32, record_batch.schema().field(0).data_type());
        assert_eq!(
            &DataType::Utf8View,
            record_batch.schema().field(1).data_type()
        );
        assert_eq!(5, record_batch.column(0).len());
        assert_eq!(5, record_batch.column(1).len());
    }

    #[test]
    fn byte_size_should_not_regress() {
        let schema = Schema::new(vec![
            Field::new("a", DataType::Int32, false),
            Field::new("b", DataType::Utf8, false),
        ]);

        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let b = StringArray::from(vec!["a", "b", "c", "d", "e"]);

        let record_batch =
            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
        assert_eq!(record_batch.get_array_memory_size(), 364);
    }

    fn check_batch(record_batch: RecordBatch, num_rows: usize) {
        assert_eq!(num_rows, record_batch.num_rows());
        assert_eq!(2, record_batch.num_columns());
        assert_eq!(&DataType::Int32, record_batch.schema().field(0).data_type());
        assert_eq!(&DataType::Utf8, record_batch.schema().field(1).data_type());
        assert_eq!(num_rows, record_batch.column(0).len());
        assert_eq!(num_rows, record_batch.column(1).len());
    }

    #[test]
    #[should_panic(expected = "assertion failed: (offset + length) <= self.num_rows()")]
    fn create_record_batch_slice() {
        let schema = Schema::new(vec![
            Field::new("a", DataType::Int32, false),
            Field::new("b", DataType::Utf8, false),
        ]);
        let expected_schema = schema.clone();

        let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
        let b = StringArray::from(vec!["a", "b", "c", "d", "e", "f", "h", "i"]);

        let record_batch =
            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();

        let offset = 2;
        let length = 5;
        let record_batch_slice = record_batch.slice(offset, length);

        assert_eq!(record_batch_slice.schema().as_ref(), &expected_schema);
        check_batch(record_batch_slice, 5);

        let offset = 2;
        let length = 0;
        let record_batch_slice = record_batch.slice(offset, length);

        assert_eq!(record_batch_slice.schema().as_ref(), &expected_schema);
        check_batch(record_batch_slice, 0);

        let offset = 2;
        let length = 10;
        let _record_batch_slice = record_batch.slice(offset, length);
    }

    #[test]
    #[should_panic(expected = "assertion failed: (offset + length) <= self.num_rows()")]
    fn create_record_batch_slice_empty_batch() {
        let schema = Schema::empty();

        let record_batch = RecordBatch::new_empty(Arc::new(schema));

        let offset = 0;
        let length = 0;
        let record_batch_slice = record_batch.slice(offset, length);
        assert_eq!(0, record_batch_slice.schema().fields().len());

        let offset = 1;
        let length = 2;
        let _record_batch_slice = record_batch.slice(offset, length);
    }

    #[test]
    fn create_record_batch_try_from_iter() {
        let a: ArrayRef = Arc::new(Int32Array::from(vec![
            Some(1),
            Some(2),
            None,
            Some(4),
            Some(5),
        ]));
        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));

        let record_batch =
            RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).expect("valid conversion");

        let expected_schema = Schema::new(vec![
            Field::new("a", DataType::Int32, true),
            Field::new("b", DataType::Utf8, false),
        ]);
        assert_eq!(record_batch.schema().as_ref(), &expected_schema);
        check_batch(record_batch, 5);
    }

    #[test]
    fn create_record_batch_try_from_iter_with_nullable() {
        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));

        // Note there are no nulls in a or b, but we specify that b is nullable
        let record_batch =
            RecordBatch::try_from_iter_with_nullable(vec![("a", a, false), ("b", b, true)])
                .expect("valid conversion");

        let expected_schema = Schema::new(vec![
            Field::new("a", DataType::Int32, false),
            Field::new("b", DataType::Utf8, true),
        ]);
        assert_eq!(record_batch.schema().as_ref(), &expected_schema);
        check_batch(record_batch, 5);
    }

    #[test]
    fn create_record_batch_schema_mismatch() {
        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

        let a = Int64Array::from(vec![1, 2, 3, 4, 5]);

        let err = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap_err();
        assert_eq!(
            err.to_string(),
            "Invalid argument error: column types must match schema types, expected Int32 but found Int64 at column index 0"
        );
    }

    #[test]
    fn create_record_batch_field_name_mismatch() {
        let fields = vec![
            Field::new("a1", DataType::Int32, false),
            Field::new_list("a2", Field::new_list_field(DataType::Int8, false), false),
        ];
        let schema = Arc::new(Schema::new(vec![Field::new_struct("a", fields, true)]));

        let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
        let a2_child = Int8Array::from(vec![1, 2, 3, 4]);
        let a2 = ArrayDataBuilder::new(DataType::List(Arc::new(Field::new(
            "array",
            DataType::Int8,
            false,
        ))))
        .add_child_data(a2_child.into_data())
        .len(2)
        .add_buffer(Buffer::from([0i32, 3, 4].to_byte_slice()))
        .build()
        .unwrap();
        let a2: ArrayRef = Arc::new(ListArray::from(a2));
        let a = ArrayDataBuilder::new(DataType::Struct(Fields::from(vec![
            Field::new("aa1", DataType::Int32, false),
            Field::new("a2", a2.data_type().clone(), false),
        ])))
        .add_child_data(a1.into_data())
        .add_child_data(a2.into_data())
        .len(2)
        .build()
        .unwrap();
        let a: ArrayRef = Arc::new(StructArray::from(a));

        // creating the batch with field name validation should fail
        let batch = RecordBatch::try_new(schema.clone(), vec![a.clone()]);
        assert!(batch.is_err());

        // creating the batch without field name validation should pass
        let options = RecordBatchOptions {
            match_field_names: false,
            row_count: None,
        };
        let batch = RecordBatch::try_new_with_options(schema, vec![a], &options);
        assert!(batch.is_ok());
    }

    #[test]
    fn create_record_batch_record_mismatch() {
        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let b = Int32Array::from(vec![1, 2, 3, 4, 5]);

        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]);
        assert!(batch.is_err());
    }

    #[test]
    fn create_record_batch_from_struct_array() {
        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
        let struct_array = StructArray::from(vec![
            (
                Arc::new(Field::new("b", DataType::Boolean, false)),
                boolean.clone() as ArrayRef,
            ),
            (
                Arc::new(Field::new("c", DataType::Int32, false)),
                int.clone() as ArrayRef,
            ),
        ]);

        let batch = RecordBatch::from(&struct_array);
        assert_eq!(2, batch.num_columns());
        assert_eq!(4, batch.num_rows());
        assert_eq!(
            struct_array.data_type(),
            &DataType::Struct(batch.schema().fields().clone())
        );
        assert_eq!(batch.column(0).as_ref(), boolean.as_ref());
        assert_eq!(batch.column(1).as_ref(), int.as_ref());
    }

    #[test]
    fn record_batch_equality() {
        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema2 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let batch1 = RecordBatch::try_new(
            Arc::new(schema1),
            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            Arc::new(schema2),
            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
        )
        .unwrap();

        assert_eq!(batch1, batch2);
    }

    /// validates if the record batch can be accessed using `column_name` as index i.e. `record_batch["column_name"]`
    #[test]
    fn record_batch_index_access() {
        let id_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
        let val_arr = Arc::new(Int32Array::from(vec![5, 6, 7, 8]));
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);
        let record_batch =
            RecordBatch::try_new(Arc::new(schema1), vec![id_arr.clone(), val_arr.clone()]).unwrap();

        assert_eq!(record_batch["id"].as_ref(), id_arr.as_ref());
        assert_eq!(record_batch["val"].as_ref(), val_arr.as_ref());
    }

    #[test]
    fn record_batch_vals_ne() {
        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let schema2 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let batch1 = RecordBatch::try_new(
            Arc::new(schema1),
            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            Arc::new(schema2),
            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
        )
        .unwrap();

        assert_ne!(batch1, batch2);
    }

    #[test]
    fn record_batch_column_names_ne() {
        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema2 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
        ]);

        let batch1 = RecordBatch::try_new(
            Arc::new(schema1),
            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            Arc::new(schema2),
            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
        )
        .unwrap();

        assert_ne!(batch1, batch2);
    }

    #[test]
    fn record_batch_column_number_ne() {
        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
        let num_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema2 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
        ]);

        let batch1 = RecordBatch::try_new(
            Arc::new(schema1),
            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            Arc::new(schema2),
            vec![Arc::new(id_arr2), Arc::new(val_arr2), Arc::new(num_arr2)],
        )
        .unwrap();

        assert_ne!(batch1, batch2);
    }

    #[test]
    fn record_batch_row_count_ne() {
        let id_arr1 = Int32Array::from(vec![1, 2, 3]);
        let val_arr1 = Int32Array::from(vec![5, 6, 7]);
        let schema1 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("val", DataType::Int32, false),
        ]);

        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
        let schema2 = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("num", DataType::Int32, false),
        ]);

        let batch1 = RecordBatch::try_new(
            Arc::new(schema1),
            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
        )
        .unwrap();

        let batch2 = RecordBatch::try_new(
            Arc::new(schema2),
            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
        )
        .unwrap();

        assert_ne!(batch1, batch2);
    }

    #[test]
    fn normalize_simple() {
        let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""]));
        let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)]));
        let year: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2022)]));

        let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
        let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
        let year_field = Arc::new(Field::new("year", DataType::Int64, true));

        let a = Arc::new(StructArray::from(vec![
            (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef),
            (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef),
            (year_field.clone(), Arc::new(year.clone()) as ArrayRef),
        ]));

        let month = Arc::new(Int64Array::from(vec![Some(4), Some(6)]));

        let schema = Schema::new(vec![
            Field::new(
                "a",
                DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])),
                false,
            ),
            Field::new("month", DataType::Int64, true),
        ]);

        let normalized =
            RecordBatch::try_new(Arc::new(schema.clone()), vec![a.clone(), month.clone()])
                .expect("valid conversion")
                .normalize(".", Some(0))
                .expect("valid normalization");

        let expected = RecordBatch::try_from_iter_with_nullable(vec![
            ("a.animals", animals.clone(), true),
            ("a.n_legs", n_legs.clone(), true),
            ("a.year", year.clone(), true),
            ("month", month.clone(), true),
        ])
        .expect("valid conversion");

        assert_eq!(expected, normalized);

        // check 0 and None have the same effect
        let normalized = RecordBatch::try_new(Arc::new(schema), vec![a, month.clone()])
            .expect("valid conversion")
            .normalize(".", None)
            .expect("valid normalization");

        assert_eq!(expected, normalized);
    }

    #[test]
    fn normalize_nested() {
        // Initialize schema
        let a = Arc::new(Field::new("a", DataType::Int64, true));
        let b = Arc::new(Field::new("b", DataType::Int64, false));
        let c = Arc::new(Field::new("c", DataType::Int64, true));

        let one = Arc::new(Field::new(
            "1",
            DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])),
            false,
        ));
        let two = Arc::new(Field::new(
            "2",
            DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])),
            true,
        ));

        let exclamation = Arc::new(Field::new(
            "!",
            DataType::Struct(Fields::from(vec![one.clone(), two.clone()])),
            false,
        ));

        let schema = Schema::new(vec![exclamation.clone()]);

        // Initialize fields
        let a_field = Int64Array::from(vec![Some(0), Some(1)]);
        let b_field = Int64Array::from(vec![Some(2), Some(3)]);
        let c_field = Int64Array::from(vec![None, Some(4)]);

        let one_field = StructArray::from(vec![
            (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
            (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
            (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
        ]);
        let two_field = StructArray::from(vec![
            (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
            (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
            (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
        ]);

        let exclamation_field = Arc::new(StructArray::from(vec![
            (one.clone(), Arc::new(one_field) as ArrayRef),
            (two.clone(), Arc::new(two_field) as ArrayRef),
        ]));

        // Normalize top level
        let normalized =
            RecordBatch::try_new(Arc::new(schema.clone()), vec![exclamation_field.clone()])
                .expect("valid conversion")
                .normalize(".", Some(1))
                .expect("valid normalization");

        let expected = RecordBatch::try_from_iter_with_nullable(vec![
            (
                "!.1",
                Arc::new(StructArray::from(vec![
                    (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
                    (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
                    (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
                ])) as ArrayRef,
                false,
            ),
            (
                "!.2",
                Arc::new(StructArray::from(vec![
                    (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
                    (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
                    (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
                ])) as ArrayRef,
                true,
            ),
        ])
        .expect("valid conversion");

        assert_eq!(expected, normalized);

        // Normalize all levels
        let normalized = RecordBatch::try_new(Arc::new(schema), vec![exclamation_field])
            .expect("valid conversion")
            .normalize(".", None)
            .expect("valid normalization");

        let expected = RecordBatch::try_from_iter_with_nullable(vec![
            ("!.1.a", Arc::new(a_field.clone()) as ArrayRef, true),
            ("!.1.b", Arc::new(b_field.clone()) as ArrayRef, false),
            ("!.1.c", Arc::new(c_field.clone()) as ArrayRef, true),
            ("!.2.a", Arc::new(a_field.clone()) as ArrayRef, true),
            ("!.2.b", Arc::new(b_field.clone()) as ArrayRef, false),
            ("!.2.c", Arc::new(c_field.clone()) as ArrayRef, true),
        ])
        .expect("valid conversion");

        assert_eq!(expected, normalized);
    }

    #[test]
    fn normalize_empty() {
        let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
        let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
        let year_field = Arc::new(Field::new("year", DataType::Int64, true));

        let schema = Schema::new(vec![
            Field::new(
                "a",
                DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])),
                false,
            ),
            Field::new("month", DataType::Int64, true),
        ]);

        let normalized = RecordBatch::new_empty(Arc::new(schema.clone()))
            .normalize(".", Some(0))
            .expect("valid normalization");

        let expected = RecordBatch::new_empty(Arc::new(
            schema.normalize(".", Some(0)).expect("valid normalization"),
        ));

        assert_eq!(expected, normalized);
    }

    #[test]
    fn project() {
        let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]));
        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
        let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));

        let record_batch =
            RecordBatch::try_from_iter(vec![("a", a.clone()), ("b", b.clone()), ("c", c.clone())])
                .expect("valid conversion");

        let expected =
            RecordBatch::try_from_iter(vec![("a", a), ("c", c)]).expect("valid conversion");

        assert_eq!(expected, record_batch.project(&[0, 2]).unwrap());
    }

    #[test]
    fn project_empty() {
        let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));

        let record_batch =
            RecordBatch::try_from_iter(vec![("c", c.clone())]).expect("valid conversion");

        let expected = RecordBatch::try_new_with_options(
            Arc::new(Schema::empty()),
            vec![],
            &RecordBatchOptions {
                match_field_names: true,
                row_count: Some(3),
            },
        )
        .expect("valid conversion");

        assert_eq!(expected, record_batch.project(&[]).unwrap());
    }

    #[test]
    fn test_no_column_record_batch() {
        let schema = Arc::new(Schema::empty());

        let err = RecordBatch::try_new(schema.clone(), vec![]).unwrap_err();
        assert!(
            err.to_string()
                .contains("must either specify a row count or at least one column")
        );

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

        let ok = RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
        assert_eq!(ok.num_rows(), 10);

        let a = ok.slice(2, 5);
        assert_eq!(a.num_rows(), 5);

        let b = ok.slice(5, 0);
        assert_eq!(b.num_rows(), 0);

        assert_ne!(a, b);
        assert_eq!(b, RecordBatch::new_empty(schema))
    }

    #[test]
    fn test_nulls_in_non_nullable_field() {
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        let maybe_batch = RecordBatch::try_new(
            schema,
            vec![Arc::new(Int32Array::from(vec![Some(1), None]))],
        );
        assert_eq!(
            "Invalid argument error: Column 'a' is declared as non-nullable but contains null values",
            format!("{}", maybe_batch.err().unwrap())
        );
    }
    #[test]
    fn test_record_batch_options() {
        let options = RecordBatchOptions::new()
            .with_match_field_names(false)
            .with_row_count(Some(20));
        assert!(!options.match_field_names);
        assert_eq!(options.row_count.unwrap(), 20)
    }

    #[test]
    #[should_panic(expected = "Cannot convert nullable StructArray to RecordBatch")]
    fn test_from_struct() {
        let s = StructArray::from(ArrayData::new_null(
            // Note child is not nullable
            &DataType::Struct(vec![Field::new("foo", DataType::Int32, false)].into()),
            2,
        ));
        let _ = RecordBatch::from(s);
    }

    #[test]
    fn test_with_schema() {
        let required_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
        let required_schema = Arc::new(required_schema);
        let nullable_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
        let nullable_schema = Arc::new(nullable_schema);

        let batch = RecordBatch::try_new(
            required_schema.clone(),
            vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as _],
        )
        .unwrap();

        // Can add nullability
        let batch = batch.with_schema(nullable_schema.clone()).unwrap();

        // Cannot remove nullability
        batch.clone().with_schema(required_schema).unwrap_err();

        // Can add metadata
        let metadata = vec![("foo".to_string(), "bar".to_string())]
            .into_iter()
            .collect();
        let metadata_schema = nullable_schema.as_ref().clone().with_metadata(metadata);
        let batch = batch.with_schema(Arc::new(metadata_schema)).unwrap();

        // Cannot remove metadata
        batch.with_schema(nullable_schema).unwrap_err();
    }

    #[test]
    fn test_boxed_reader() {
        // Make sure we can pass a boxed reader to a function generic over
        // RecordBatchReader.
        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
        let schema = Arc::new(schema);

        let reader = RecordBatchIterator::new(std::iter::empty(), schema);
        let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);

        fn get_size(reader: impl RecordBatchReader) -> usize {
            reader.size_hint().0
        }

        let size = get_size(reader);
        assert_eq!(size, 0);
    }

    #[test]
    fn test_remove_column_maintains_schema_metadata() {
        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
        let bool_array = BooleanArray::from(vec![true, false, false, true, true]);

        let mut metadata = HashMap::new();
        metadata.insert("foo".to_string(), "bar".to_string());
        let schema = Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("bool", DataType::Boolean, false),
        ])
        .with_metadata(metadata);

        let mut batch = RecordBatch::try_new(
            Arc::new(schema),
            vec![Arc::new(id_array), Arc::new(bool_array)],
        )
        .unwrap();

        let _removed_column = batch.remove_column(0);
        assert_eq!(batch.schema().metadata().len(), 1);
        assert_eq!(
            batch.schema().metadata().get("foo").unwrap().as_str(),
            "bar"
        );
    }
}

Coverage Report

Created: 2025-11-17 14:14