arrow_array/
record_batch.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! A two-dimensional batch of column-oriented data with a defined
6//! [schema](arrow_schema::Schema).
7
8use crate::cast::AsArray;
9use crate::{Array, ArrayRef, StructArray, new_empty_array};
10use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaBuilder, SchemaRef};
11use std::ops::Index;
12use std::sync::Arc;
13
14/// Trait for types that can read `RecordBatch`'s.
15///
16/// To create from an iterator, see [RecordBatchIterator].
17pub trait RecordBatchReader: Iterator<Item = Result<RecordBatch, ArrowError>> {
18    /// Returns the schema of this `RecordBatchReader`.
19    ///
20    /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
21    /// reader should have the same schema as returned from this method.
22    fn schema(&self) -> SchemaRef;
23}
24
25impl<R: RecordBatchReader + ?Sized> RecordBatchReader for Box<R> {
26    fn schema(&self) -> SchemaRef {
27        self.as_ref().schema()
28    }
29}
30
31/// Trait for types that can write `RecordBatch`'s.
32pub trait RecordBatchWriter {
33    /// Write a single batch to the writer.
34    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError>;
35
36    /// Write footer or termination data, then mark the writer as done.
37    fn close(self) -> Result<(), ArrowError>;
38}
39
40/// Creates an array from a literal slice of values,
41/// suitable for rapid testing and development.
42///
43/// Example:
44///
45/// ```rust
46///
47/// use arrow_array::create_array;
48///
49/// let array = create_array!(Int32, [1, 2, 3, 4, 5]);
50/// let array = create_array!(Utf8, [Some("a"), Some("b"), None, Some("e")]);
51/// ```
52/// Support for limited data types is available. The macro will return a compile error if an unsupported data type is used.
53/// Presently supported data types are:
54/// - `Boolean`, `Null`
55/// - `Decimal32`, `Decimal64`, `Decimal128`, `Decimal256`
56/// - `Float16`, `Float32`, `Float64`
57/// - `Int8`, `Int16`, `Int32`, `Int64`
58/// - `UInt8`, `UInt16`, `UInt32`, `UInt64`
59/// - `IntervalDayTime`, `IntervalYearMonth`
60/// - `Second`, `Millisecond`, `Microsecond`, `Nanosecond`
61/// - `Second32`, `Millisecond32`, `Microsecond64`, `Nanosecond64`
62/// - `DurationSecond`, `DurationMillisecond`, `DurationMicrosecond`, `DurationNanosecond`
63/// - `TimestampSecond`, `TimestampMillisecond`, `TimestampMicrosecond`, `TimestampNanosecond`
64/// - `Utf8`, `Utf8View`, `LargeUtf8`, `Binary`, `LargeBinary`
65#[macro_export]
66macro_rules! create_array {
67    // `@from` is used for those types that have a common method `<type>::from`
68    (@from Boolean) => { $crate::BooleanArray };
69    (@from Int8) => { $crate::Int8Array };
70    (@from Int16) => { $crate::Int16Array };
71    (@from Int32) => { $crate::Int32Array };
72    (@from Int64) => { $crate::Int64Array };
73    (@from UInt8) => { $crate::UInt8Array };
74    (@from UInt16) => { $crate::UInt16Array };
75    (@from UInt32) => { $crate::UInt32Array };
76    (@from UInt64) => { $crate::UInt64Array };
77    (@from Float16) => { $crate::Float16Array };
78    (@from Float32) => { $crate::Float32Array };
79    (@from Float64) => { $crate::Float64Array };
80    (@from Utf8) => { $crate::StringArray };
81    (@from Utf8View) => { $crate::StringViewArray };
82    (@from LargeUtf8) => { $crate::LargeStringArray };
83    (@from IntervalDayTime) => { $crate::IntervalDayTimeArray };
84    (@from IntervalYearMonth) => { $crate::IntervalYearMonthArray };
85    (@from Second) => { $crate::TimestampSecondArray };
86    (@from Millisecond) => { $crate::TimestampMillisecondArray };
87    (@from Microsecond) => { $crate::TimestampMicrosecondArray };
88    (@from Nanosecond) => { $crate::TimestampNanosecondArray };
89    (@from Second32) => { $crate::Time32SecondArray };
90    (@from Millisecond32) => { $crate::Time32MillisecondArray };
91    (@from Microsecond64) => { $crate::Time64MicrosecondArray };
92    (@from Nanosecond64) => { $crate::Time64Nanosecond64Array };
93    (@from DurationSecond) => { $crate::DurationSecondArray };
94    (@from DurationMillisecond) => { $crate::DurationMillisecondArray };
95    (@from DurationMicrosecond) => { $crate::DurationMicrosecondArray };
96    (@from DurationNanosecond) => { $crate::DurationNanosecondArray };
97    (@from Decimal32) => { $crate::Decimal32Array };
98    (@from Decimal64) => { $crate::Decimal64Array };
99    (@from Decimal128) => { $crate::Decimal128Array };
100    (@from Decimal256) => { $crate::Decimal256Array };
101    (@from TimestampSecond) => { $crate::TimestampSecondArray };
102    (@from TimestampMillisecond) => { $crate::TimestampMillisecondArray };
103    (@from TimestampMicrosecond) => { $crate::TimestampMicrosecondArray };
104    (@from TimestampNanosecond) => { $crate::TimestampNanosecondArray };
105    (@from Date32) => { $crate::Date32Array };
106    (@from Date64) => { $crate::Date64Array };
107
108    (@from $ty: ident) => {
109        compile_error!(concat!("Unsupported data type: ", stringify!($ty)))
110    };
111
112    (Null, $size: expr) => {
113        std::sync::Arc::new($crate::NullArray::new($size))
114    };
115
116    (Binary, [$($values: expr),*]) => {
117        std::sync::Arc::new($crate::BinaryArray::from_vec(vec![$($values),*]))
118    };
119
120    (LargeBinary, [$($values: expr),*]) => {
121        std::sync::Arc::new($crate::LargeBinaryArray::from_vec(vec![$($values),*]))
122    };
123
124    ($ty: tt, [$($values: expr),*]) => {
125        std::sync::Arc::new(<$crate::create_array!(@from $ty)>::from(vec![$($values),*]))
126    };
127
128    (Binary, $values: expr) => {
129        std::sync::Arc::new($crate::BinaryArray::from_vec($values))
130    };
131
132    (LargeBinary, $values: expr) => {
133        std::sync::Arc::new($crate::LargeBinaryArray::from_vec($values))
134    };
135
136    ($ty: tt, $values: expr) => {
137        std::sync::Arc::new(<$crate::create_array!(@from $ty)>::from($values))
138    };
139}
140
141/// Creates a record batch from literal slice of values, suitable for rapid
142/// testing and development.
143///
144/// Example:
145///
146/// ```rust
147/// use arrow_array::record_batch;
148///
149/// let batch = record_batch!(
150///     ("a", Int32, [1, 2, 3]),
151///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
152///     ("c", Utf8, ["alpha", "beta", "gamma"])
153/// );
154/// ```
155///
156/// Variables and expressions are also supported:
157///
158/// ```rust
159/// use arrow_array::record_batch;
160///
161/// let values = vec![1, 2, 3];
162/// let batch = record_batch!(
163///     ("a", Int32, values),
164///     ("b", Float64, vec![Some(4.0), None, Some(5.0)])
165/// );
166/// ```
167/// Due to limitation of [`create_array!`] macro, support for limited data types is available.
168#[macro_export]
169macro_rules! record_batch {
170    ($(($name: expr, $type: ident, $($values: tt)+)),*) => {
171        $crate::RecordBatch::try_from_iter_with_nullable(vec![$(
172            (
173                ::std::string::String::from($name),
174                $crate::create_array!($type, $($values)+) as $crate::ArrayRef,
175                true,
176            ),
177        )*])
178    };
179}
180
181/// A two-dimensional batch of column-oriented data with a defined
182/// [schema](arrow_schema::Schema).
183///
184/// A `RecordBatch` is a two-dimensional dataset of a number of
185/// contiguous arrays, each the same length.
186/// A record batch has a schema which must match its arrays’
187/// datatypes.
188///
189/// Record batches are a convenient unit of work for various
190/// serialization and computation functions, possibly incremental.
191///
192/// Use the [`record_batch!`] macro to create a [`RecordBatch`] from
193/// literal slice of values, useful for rapid prototyping and testing.
194///
195/// Example:
196/// ```rust
197/// use arrow_array::record_batch;
198/// let batch = record_batch!(
199///     ("a", Int32, [1, 2, 3]),
200///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
201///     ("c", Utf8, ["alpha", "beta", "gamma"])
202/// );
203/// ```
204#[derive(Clone, Debug, PartialEq)]
205pub struct RecordBatch {
206    schema: SchemaRef,
207    columns: Vec<Arc<dyn Array>>,
208
209    /// The number of rows in this RecordBatch
210    ///
211    /// This is stored separately from the columns to handle the case of no columns
212    row_count: usize,
213}
214
215impl RecordBatch {
216    /// Creates a `RecordBatch` from a schema and columns.
217    ///
218    /// Expects the following:
219    ///
220    ///  * `!columns.is_empty()`
221    ///  * `schema.fields.len() == columns.len()`
222    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
223    ///  * `columns[i].len() == columns[j].len()`
224    ///
225    /// If the conditions are not met, an error is returned.
226    ///
227    /// # Example
228    ///
229    /// ```
230    /// # use std::sync::Arc;
231    /// # use arrow_array::{Int32Array, RecordBatch};
232    /// # use arrow_schema::{DataType, Field, Schema};
233    ///
234    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
235    /// let schema = Schema::new(vec![
236    ///     Field::new("id", DataType::Int32, false)
237    /// ]);
238    ///
239    /// let batch = RecordBatch::try_new(
240    ///     Arc::new(schema),
241    ///     vec![Arc::new(id_array)]
242    /// ).unwrap();
243    /// ```
244    pub fn try_new(schema: SchemaRef, columns: Vec<ArrayRef>) -> Result<Self, ArrowError> {
245        let options = RecordBatchOptions::new();
246        Self::try_new_impl(schema, columns, &options)
247    }
248
249    /// Creates a `RecordBatch` from a schema and columns, without validation.
250    ///
251    /// See [`Self::try_new`] for the checked version.
252    ///
253    /// # Safety
254    ///
255    /// Expects the following:
256    ///
257    ///  * `schema.fields.len() == columns.len()`
258    ///  * `schema.fields[i].data_type() == columns[i].data_type()`
259    ///  * `columns[i].len() == row_count`
260    ///
261    /// Note: if the schema does not match the underlying data exactly, it can lead to undefined
262    /// behavior, for example, via conversion to a `StructArray`, which in turn could lead
263    /// to incorrect access.
264    pub unsafe fn new_unchecked(
265        schema: SchemaRef,
266        columns: Vec<Arc<dyn Array>>,
267        row_count: usize,
268    ) -> Self {
269        Self {
270            schema,
271            columns,
272            row_count,
273        }
274    }
275
276    /// Creates a `RecordBatch` from a schema and columns, with additional options,
277    /// such as whether to strictly validate field names.
278    ///
279    /// See [`RecordBatch::try_new`] for the expected conditions.
280    pub fn try_new_with_options(
281        schema: SchemaRef,
282        columns: Vec<ArrayRef>,
283        options: &RecordBatchOptions,
284    ) -> Result<Self, ArrowError> {
285        Self::try_new_impl(schema, columns, options)
286    }
287
288    /// Creates a new empty [`RecordBatch`].
289    pub fn new_empty(schema: SchemaRef) -> Self {
290        let columns = schema
291            .fields()
292            .iter()
293            .map(|field| new_empty_array(field.data_type()))
294            .collect();
295
296        RecordBatch {
297            schema,
298            columns,
299            row_count: 0,
300        }
301    }
302
303    /// Validate the schema and columns using [`RecordBatchOptions`]. Returns an error
304    /// if any validation check fails, otherwise returns the created [`Self`]
305    fn try_new_impl(
306        schema: SchemaRef,
307        columns: Vec<ArrayRef>,
308        options: &RecordBatchOptions,
309    ) -> Result<Self, ArrowError> {
310        // check that number of fields in schema match column length
311        if schema.fields().len() != columns.len() {
312            return Err(ArrowError::InvalidArgumentError(format!(
313                "number of columns({}) must match number of fields({}) in schema",
314                columns.len(),
315                schema.fields().len(),
316            )));
317        }
318
319        let row_count = options
320            .row_count
321            .or_else(|| columns.first().map(|col| col.len()))
322            .ok_or_else(|| {
323                ArrowError::InvalidArgumentError(
324                    "must either specify a row count or at least one column".to_string(),
325                )
326            })?;
327
328        for (c, f) in columns.iter().zip(&schema.fields) {
329            if !f.is_nullable() && c.null_count() > 0 {
330                return Err(ArrowError::InvalidArgumentError(format!(
331                    "Column '{}' is declared as non-nullable but contains null values",
332                    f.name()
333                )));
334            }
335        }
336
337        // check that all columns have the same row count
338        if columns.iter().any(|c| c.len() != row_count) {
339            let err = match options.row_count {
340                Some(_) => "all columns in a record batch must have the specified row count",
341                None => "all columns in a record batch must have the same length",
342            };
343            return Err(ArrowError::InvalidArgumentError(err.to_string()));
344        }
345
346        // function for comparing column type and field type
347        // return true if 2 types are not matched
348        let type_not_match = if options.match_field_names {
349            |(_, (col_type, field_type)): &(usize, (&DataType, &DataType))| col_type != field_type
350        } else {
351            |(_, (col_type, field_type)): &(usize, (&DataType, &DataType))| {
352                !col_type.equals_datatype(field_type)
353            }
354        };
355
356        // check that all columns match the schema
357        let not_match = columns
358            .iter()
359            .zip(schema.fields().iter())
360            .map(|(col, field)| (col.data_type(), field.data_type()))
361            .enumerate()
362            .find(type_not_match);
363
364        if let Some((i, (col_type, field_type))) = not_match {
365            return Err(ArrowError::InvalidArgumentError(format!(
366                "column types must match schema types, expected {field_type} but found {col_type} at column index {i}"
367            )));
368        }
369
370        Ok(RecordBatch {
371            schema,
372            columns,
373            row_count,
374        })
375    }
376
377    /// Return the schema, columns and row count of this [`RecordBatch`]
378    pub fn into_parts(self) -> (SchemaRef, Vec<ArrayRef>, usize) {
379        (self.schema, self.columns, self.row_count)
380    }
381
382    /// Override the schema of this [`RecordBatch`]
383    ///
384    /// Returns an error if `schema` is not a superset of the current schema
385    /// as determined by [`Schema::contains`]
386    ///
387    /// See also [`Self::schema_metadata_mut`].
388    pub fn with_schema(self, schema: SchemaRef) -> Result<Self, ArrowError> {
389        if !schema.contains(self.schema.as_ref()) {
390            return Err(ArrowError::SchemaError(format!(
391                "target schema is not superset of current schema target={schema} current={}",
392                self.schema
393            )));
394        }
395
396        Ok(Self {
397            schema,
398            columns: self.columns,
399            row_count: self.row_count,
400        })
401    }
402
403    /// Returns the [`Schema`] of the record batch.
404    pub fn schema(&self) -> SchemaRef {
405        self.schema.clone()
406    }
407
408    /// Returns a reference to the [`Schema`] of the record batch.
409    pub fn schema_ref(&self) -> &SchemaRef {
410        &self.schema
411    }
412
413    /// Mutable access to the metadata of the schema.
414    ///
415    /// This allows you to modify [`Schema::metadata`] of [`Self::schema`] in a convenient and fast way.
416    ///
417    /// Note this will clone the entire underlying `Schema` object if it is currently shared
418    ///
419    /// # Example
420    /// ```
421    /// # use std::sync::Arc;
422    /// # use arrow_array::{record_batch, RecordBatch};
423    /// let mut batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
424    /// // Initially, the metadata is empty
425    /// assert!(batch.schema().metadata().get("key").is_none());
426    /// // Insert a key-value pair into the metadata
427    /// batch.schema_metadata_mut().insert("key", "value");
428    /// assert_eq!(batch.schema().metadata().get("key"), Some(&String::from("value")));
429    /// ```
430    pub fn schema_metadata_mut(&mut self) -> &mut arrow_schema::Metadata {
431        let schema = Arc::make_mut(&mut self.schema);
432        &mut schema.metadata
433    }
434
435    /// Projects the schema onto the specified columns
436    pub fn project(&self, indices: &[usize]) -> Result<RecordBatch, ArrowError> {
437        let projected_schema = self.schema.project(indices)?;
438        let batch_fields = indices
439            .iter()
440            .map(|f| {
441                self.columns.get(*f).cloned().ok_or_else(|| {
442                    ArrowError::SchemaError(format!(
443                        "project index {} out of bounds, max field {}",
444                        f,
445                        self.columns.len()
446                    ))
447                })
448            })
449            .collect::<Result<Vec<_>, _>>()?;
450
451        unsafe {
452            // Since we're starting from a valid RecordBatch and project
453            // creates a strict subset of the original, there's no need to
454            // redo the validation checks in `try_new_with_options`.
455            Ok(RecordBatch::new_unchecked(
456                SchemaRef::new(projected_schema),
457                batch_fields,
458                self.row_count,
459            ))
460        }
461    }
462
463    /// Normalize a semi-structured [`RecordBatch`] into a flat table.
464    ///
465    /// Nested [`Field`]s will generate names separated by `separator`, up to a depth of `max_level`
466    /// (unlimited if `None`).
467    ///
468    /// e.g. given a [`RecordBatch`] with schema:
469    ///
470    /// ```text
471    ///     "foo": StructArray<"bar": Utf8>
472    /// ```
473    ///
474    /// A separator of `"."` would generate a batch with the schema:
475    ///
476    /// ```text
477    ///     "foo.bar": Utf8
478    /// ```
479    ///
480    /// Note that giving a depth of `Some(0)` to `max_level` is the same as passing in `None`;
481    /// it will be treated as unlimited.
482    ///
483    /// # Example
484    ///
485    /// ```
486    /// # use std::sync::Arc;
487    /// # use arrow_array::{ArrayRef, Int64Array, StringArray, StructArray, RecordBatch};
488    /// # use arrow_schema::{DataType, Field, Fields, Schema};
489    /// #
490    /// let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""]));
491    /// let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)]));
492    ///
493    /// let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
494    /// let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
495    ///
496    /// let a = Arc::new(StructArray::from(vec![
497    ///     (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef),
498    ///     (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef),
499    /// ]));
500    ///
501    /// let schema = Schema::new(vec![
502    ///     Field::new(
503    ///         "a",
504    ///         DataType::Struct(Fields::from(vec![animals_field, n_legs_field])),
505    ///         false,
506    ///     )
507    /// ]);
508    ///
509    /// let normalized = RecordBatch::try_new(Arc::new(schema), vec![a])
510    ///     .expect("valid conversion")
511    ///     .normalize(".", None)
512    ///     .expect("valid normalization");
513    ///
514    /// let expected = RecordBatch::try_from_iter_with_nullable(vec![
515    ///     ("a.animals", animals.clone(), true),
516    ///     ("a.n_legs", n_legs.clone(), true),
517    /// ])
518    /// .expect("valid conversion");
519    ///
520    /// assert_eq!(expected, normalized);
521    /// ```
522    pub fn normalize(&self, separator: &str, max_level: Option<usize>) -> Result<Self, ArrowError> {
523        let max_level = match max_level.unwrap_or(usize::MAX) {
524            0 => usize::MAX,
525            val => val,
526        };
527        let mut stack: Vec<(usize, ArrayRef, String, FieldRef)> = self
528            .columns
529            .iter()
530            .zip(self.schema.fields())
531            .rev()
532            .map(|(c, f)| (0, c.clone(), f.name().clone(), Arc::clone(f)))
533            .collect();
534        let mut columns: Vec<ArrayRef> = Vec::new();
535        let mut fields: Vec<FieldRef> = Vec::new();
536
537        while let Some((depth, c, name, field_ref)) = stack.pop() {
538            match field_ref.data_type() {
539                DataType::Struct(_) if depth < max_level => {
540                    let (flat_fields, flat_cols) = c.as_struct().flatten();
541                    for (cff, fff) in flat_cols.into_iter().zip(flat_fields.iter()).rev() {
542                        let child_name = format!("{name}{separator}{}", fff.name());
543                        stack.push((depth + 1, cff, child_name, Arc::clone(fff)))
544                    }
545                }
546                _ => {
547                    let updated_field =
548                        Field::new(name, field_ref.data_type().clone(), field_ref.is_nullable());
549                    columns.push(c);
550                    fields.push(Arc::new(updated_field));
551                }
552            }
553        }
554        RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
555    }
556
557    /// Returns the number of columns in the record batch.
558    ///
559    /// # Example
560    ///
561    /// ```
562    /// # use std::sync::Arc;
563    /// # use arrow_array::{Int32Array, RecordBatch};
564    /// # use arrow_schema::{DataType, Field, Schema};
565    ///
566    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
567    /// let schema = Schema::new(vec![
568    ///     Field::new("id", DataType::Int32, false)
569    /// ]);
570    ///
571    /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)]).unwrap();
572    ///
573    /// assert_eq!(batch.num_columns(), 1);
574    /// ```
575    pub fn num_columns(&self) -> usize {
576        self.columns.len()
577    }
578
579    /// Returns the number of rows in each column.
580    ///
581    /// # Example
582    ///
583    /// ```
584    /// # use std::sync::Arc;
585    /// # use arrow_array::{Int32Array, RecordBatch};
586    /// # use arrow_schema::{DataType, Field, Schema};
587    ///
588    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
589    /// let schema = Schema::new(vec![
590    ///     Field::new("id", DataType::Int32, false)
591    /// ]);
592    ///
593    /// let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)]).unwrap();
594    ///
595    /// assert_eq!(batch.num_rows(), 5);
596    /// ```
597    pub fn num_rows(&self) -> usize {
598        self.row_count
599    }
600
601    /// Get a reference to a column's array by index.
602    ///
603    /// # Panics
604    ///
605    /// Panics if `index` is outside of `0..num_columns`.
606    pub fn column(&self, index: usize) -> &ArrayRef {
607        &self.columns[index]
608    }
609
610    /// Get a reference to a column's array by name.
611    pub fn column_by_name(&self, name: &str) -> Option<&ArrayRef> {
612        self.schema()
613            .column_with_name(name)
614            .map(|(index, _)| &self.columns[index])
615    }
616
617    /// Get a reference to all columns in the record batch.
618    pub fn columns(&self) -> &[ArrayRef] {
619        &self.columns[..]
620    }
621
622    /// Remove column by index and return it.
623    ///
624    /// Return the `ArrayRef` if the column is removed.
625    ///
626    /// # Panics
627    ///
628    /// Panics if `index` is out of bounds.
629    ///
630    /// # Example
631    ///
632    /// ```
633    /// use std::sync::Arc;
634    /// use arrow_array::{BooleanArray, Int32Array, RecordBatch};
635    /// use arrow_schema::{DataType, Field, Schema};
636    /// let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
637    /// let bool_array = BooleanArray::from(vec![true, false, false, true, true]);
638    /// let schema = Schema::new(vec![
639    ///     Field::new("id", DataType::Int32, false),
640    ///     Field::new("bool", DataType::Boolean, false),
641    /// ]);
642    ///
643    /// let mut batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array), Arc::new(bool_array)]).unwrap();
644    ///
645    /// let removed_column = batch.remove_column(0);
646    /// assert_eq!(removed_column.as_any().downcast_ref::<Int32Array>().unwrap(), &Int32Array::from(vec![1, 2, 3, 4, 5]));
647    /// assert_eq!(batch.num_columns(), 1);
648    /// ```
649    pub fn remove_column(&mut self, index: usize) -> ArrayRef {
650        let mut builder = SchemaBuilder::from(self.schema.as_ref());
651        builder.remove(index);
652        self.schema = Arc::new(builder.finish());
653        self.columns.remove(index)
654    }
655
656    /// Return a new RecordBatch where each column is sliced
657    /// according to `offset` and `length`
658    ///
659    /// # Panics
660    ///
661    /// Panics if `offset` with `length` is greater than column length.
662    pub fn slice(&self, offset: usize, length: usize) -> RecordBatch {
663        assert!((offset + length) <= self.num_rows());
664
665        let columns = self
666            .columns()
667            .iter()
668            .map(|column| column.slice(offset, length))
669            .collect();
670
671        Self {
672            schema: self.schema.clone(),
673            columns,
674            row_count: length,
675        }
676    }
677
678    /// Create a `RecordBatch` from an iterable list of pairs of the
679    /// form `(field_name, array)`, with the same requirements on
680    /// fields and arrays as [`RecordBatch::try_new`]. This method is
681    /// often used to create a single `RecordBatch` from arrays,
682    /// e.g. for testing.
683    ///
684    /// The resulting schema is marked as nullable for each column if
685    /// the array for that column is has any nulls. To explicitly
686    /// specify nullibility, use [`RecordBatch::try_from_iter_with_nullable`]
687    ///
688    /// Example:
689    /// ```
690    /// # use std::sync::Arc;
691    /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
692    ///
693    /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
694    /// let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
695    ///
696    /// let record_batch = RecordBatch::try_from_iter(vec![
697    ///   ("a", a),
698    ///   ("b", b),
699    /// ]);
700    /// ```
701    /// Another way to quickly create a [`RecordBatch`] is to use the [`record_batch!`] macro,
702    /// which is particularly helpful for rapid prototyping and testing.
703    ///
704    /// Example:
705    ///
706    /// ```rust
707    /// use arrow_array::record_batch;
708    /// let batch = record_batch!(
709    ///     ("a", Int32, [1, 2, 3]),
710    ///     ("b", Float64, [Some(4.0), None, Some(5.0)]),
711    ///     ("c", Utf8, ["alpha", "beta", "gamma"])
712    /// );
713    /// ```
714    pub fn try_from_iter<I, F>(value: I) -> Result<Self, ArrowError>
715    where
716        I: IntoIterator<Item = (F, ArrayRef)>,
717        F: AsRef<str>,
718    {
719        // TODO: implement `TryFrom` trait, once
720        // https://github.com/rust-lang/rust/issues/50133 is no longer an
721        // issue
722        let iter = value.into_iter().map(|(field_name, array)| {
723            let nullable = array.null_count() > 0;
724            (field_name, array, nullable)
725        });
726
727        Self::try_from_iter_with_nullable(iter)
728    }
729
730    /// Create a `RecordBatch` from an iterable list of tuples of the
731    /// form `(field_name, array, nullable)`, with the same requirements on
732    /// fields and arrays as [`RecordBatch::try_new`]. This method is often
733    /// used to create a single `RecordBatch` from arrays, e.g. for
734    /// testing.
735    ///
736    /// Example:
737    /// ```
738    /// # use std::sync::Arc;
739    /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
740    ///
741    /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
742    /// let b: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
743    ///
744    /// // Note neither `a` nor `b` has any actual nulls, but we mark
745    /// // b an nullable
746    /// let record_batch = RecordBatch::try_from_iter_with_nullable(vec![
747    ///   ("a", a, false),
748    ///   ("b", b, true),
749    /// ]);
750    /// ```
751    pub fn try_from_iter_with_nullable<I, F>(value: I) -> Result<Self, ArrowError>
752    where
753        I: IntoIterator<Item = (F, ArrayRef, bool)>,
754        F: AsRef<str>,
755    {
756        let iter = value.into_iter();
757        let capacity = iter.size_hint().0;
758        let mut schema = SchemaBuilder::with_capacity(capacity);
759        let mut columns = Vec::with_capacity(capacity);
760
761        for (field_name, array, nullable) in iter {
762            let field_name = field_name.as_ref();
763            schema.push(Field::new(field_name, array.data_type().clone(), nullable));
764            columns.push(array);
765        }
766
767        let schema = Arc::new(schema.finish());
768        RecordBatch::try_new(schema, columns)
769    }
770
771    /// Registers all buffers in this record batch with the provided [`MemoryPool`].
772    ///
773    /// This claims memory for all columns in the batch by calling [`Array::claim`]
774    /// on each column.
775    ///
776    /// [`MemoryPool`]: arrow_buffer::MemoryPool
777    /// [`Array::claim`]: crate::Array::claim
778    #[cfg(feature = "pool")]
779    pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
780        for column in self.columns() {
781            column.claim(pool);
782        }
783    }
784
785    /// Returns the total number of bytes of memory occupied physically by this batch.
786    ///
787    /// Note that this does not always correspond to the exact memory usage of a
788    /// `RecordBatch` (might overestimate), since multiple columns can share the same
789    /// buffers or slices thereof, the memory used by the shared buffers might be
790    /// counted multiple times.
791    pub fn get_array_memory_size(&self) -> usize {
792        self.columns()
793            .iter()
794            .map(|array| array.get_array_memory_size())
795            .sum()
796    }
797}
798
799/// Options that control the behaviour used when creating a [`RecordBatch`].
800#[derive(Debug)]
801#[non_exhaustive]
802pub struct RecordBatchOptions {
803    /// Match field names of structs and lists. If set to `true`, the names must match.
804    pub match_field_names: bool,
805
806    /// Optional row count, useful for specifying a row count for a RecordBatch with no columns
807    pub row_count: Option<usize>,
808}
809
810impl RecordBatchOptions {
811    /// Creates a new `RecordBatchOptions`
812    pub fn new() -> Self {
813        Self {
814            match_field_names: true,
815            row_count: None,
816        }
817    }
818    /// Sets the row_count of RecordBatchOptions and returns self
819    pub fn with_row_count(mut self, row_count: Option<usize>) -> Self {
820        self.row_count = row_count;
821        self
822    }
823    /// Sets the match_field_names of RecordBatchOptions and returns self
824    pub fn with_match_field_names(mut self, match_field_names: bool) -> Self {
825        self.match_field_names = match_field_names;
826        self
827    }
828}
829impl Default for RecordBatchOptions {
830    fn default() -> Self {
831        Self::new()
832    }
833}
834impl From<StructArray> for RecordBatch {
835    fn from(value: StructArray) -> Self {
836        let row_count = value.len();
837        let (fields, columns, nulls) = value.into_parts();
838        assert_eq!(
839            nulls.map(|n| n.null_count()).unwrap_or_default(),
840            0,
841            "Cannot convert nullable StructArray to RecordBatch, see StructArray documentation"
842        );
843
844        RecordBatch {
845            schema: Arc::new(Schema::new(fields)),
846            row_count,
847            columns,
848        }
849    }
850}
851
852impl From<&StructArray> for RecordBatch {
853    fn from(struct_array: &StructArray) -> Self {
854        struct_array.clone().into()
855    }
856}
857
858impl Index<&str> for RecordBatch {
859    type Output = ArrayRef;
860
861    /// Get a reference to a column's array by name.
862    ///
863    /// # Panics
864    ///
865    /// Panics if the name is not in the schema.
866    fn index(&self, name: &str) -> &Self::Output {
867        self.column_by_name(name).unwrap()
868    }
869}
870
871/// Generic implementation of [RecordBatchReader] that wraps an iterator.
872///
873/// # Example
874///
875/// ```
876/// # use std::sync::Arc;
877/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray, RecordBatchIterator, RecordBatchReader};
878/// #
879/// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
880/// let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
881///
882/// let record_batch = RecordBatch::try_from_iter(vec![
883///   ("a", a),
884///   ("b", b),
885/// ]).unwrap();
886///
887/// let batches: Vec<RecordBatch> = vec![record_batch.clone(), record_batch.clone()];
888///
889/// let mut reader = RecordBatchIterator::new(batches.into_iter().map(Ok), record_batch.schema());
890///
891/// assert_eq!(reader.schema(), record_batch.schema());
892/// assert_eq!(reader.next().unwrap().unwrap(), record_batch);
893/// # assert_eq!(reader.next().unwrap().unwrap(), record_batch);
894/// # assert!(reader.next().is_none());
895/// ```
896pub struct RecordBatchIterator<I>
897where
898    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
899{
900    inner: I::IntoIter,
901    inner_schema: SchemaRef,
902}
903
904impl<I> RecordBatchIterator<I>
905where
906    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
907{
908    /// Create a new [RecordBatchIterator].
909    ///
910    /// If `iter` is an infallible iterator, use `.map(Ok)`.
911    pub fn new(iter: I, schema: SchemaRef) -> Self {
912        Self {
913            inner: iter.into_iter(),
914            inner_schema: schema,
915        }
916    }
917}
918
919impl<I> Iterator for RecordBatchIterator<I>
920where
921    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
922{
923    type Item = I::Item;
924
925    fn next(&mut self) -> Option<Self::Item> {
926        self.inner.next()
927    }
928
929    fn size_hint(&self) -> (usize, Option<usize>) {
930        self.inner.size_hint()
931    }
932}
933
934impl<I> RecordBatchReader for RecordBatchIterator<I>
935where
936    I: IntoIterator<Item = Result<RecordBatch, ArrowError>>,
937{
938    fn schema(&self) -> SchemaRef {
939        self.inner_schema.clone()
940    }
941}
942
943#[cfg(test)]
944mod tests {
945    use super::*;
946    use crate::types::{Date64Type, Int32Type, Int64Type};
947    use crate::{
948        BooleanArray, Int8Array, Int32Array, Int64Array, ListArray, StringArray, StringViewArray,
949    };
950    use arrow_buffer::{Buffer, NullBuffer, ToByteSlice};
951    use arrow_data::{ArrayData, ArrayDataBuilder};
952    use arrow_schema::{Fields, TimeUnit};
953    use std::collections::HashMap;
954
955    #[test]
956    fn create_record_batch() {
957        let schema = Schema::new(vec![
958            Field::new("a", DataType::Int32, false),
959            Field::new("b", DataType::Utf8, false),
960        ]);
961
962        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
963        let b = StringArray::from(vec!["a", "b", "c", "d", "e"]);
964
965        let record_batch =
966            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
967        check_batch(record_batch, 5)
968    }
969
970    #[test]
971    fn create_string_view_record_batch() {
972        let schema = Schema::new(vec![
973            Field::new("a", DataType::Int32, false),
974            Field::new("b", DataType::Utf8View, false),
975        ]);
976
977        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
978        let b = StringViewArray::from(vec!["a", "b", "c", "d", "e"]);
979
980        let record_batch =
981            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
982
983        assert_eq!(5, record_batch.num_rows());
984        assert_eq!(2, record_batch.num_columns());
985        assert_eq!(&DataType::Int32, record_batch.schema().field(0).data_type());
986        assert_eq!(
987            &DataType::Utf8View,
988            record_batch.schema().field(1).data_type()
989        );
990        assert_eq!(5, record_batch.column(0).len());
991        assert_eq!(5, record_batch.column(1).len());
992    }
993
994    #[test]
995    fn create_binary_record_batch_from_variables() {
996        let binary_values = vec![b"a".as_slice()];
997        let large_binary_values = vec![b"xxx".as_slice()];
998
999        let record_batch = record_batch!(
1000            ("a", Binary, binary_values),
1001            ("b", LargeBinary, large_binary_values)
1002        )
1003        .unwrap();
1004
1005        assert_eq!(1, record_batch.num_rows());
1006        assert_eq!(2, record_batch.num_columns());
1007        assert_eq!(
1008            &DataType::Binary,
1009            record_batch.schema().field(0).data_type()
1010        );
1011        assert_eq!(
1012            &DataType::LargeBinary,
1013            record_batch.schema().field(1).data_type()
1014        );
1015
1016        let binary = record_batch.column(0).as_binary::<i32>();
1017        assert_eq!(b"a", binary.value(0));
1018
1019        let large_binary = record_batch.column(1).as_binary::<i64>();
1020        assert_eq!(b"xxx", large_binary.value(0));
1021    }
1022
1023    #[test]
1024    fn create_record_batch_macro() {
1025        let values = vec![7, 8, 9];
1026        let name = String::from("c");
1027
1028        let record_batch = record_batch!(
1029            ("a", Int32, [1, 2, 3]),
1030            ("b", Utf8, [Some("x"), None, Some("z")]),
1031            (name, Int64, values)
1032        )
1033        .unwrap();
1034
1035        let expected = Schema::new(vec![
1036            Field::new("a", DataType::Int32, true),
1037            Field::new("b", DataType::Utf8, true),
1038            Field::new("c", DataType::Int64, true),
1039        ]);
1040        assert_eq!(&expected, record_batch.schema().as_ref());
1041        assert_eq!(3, record_batch.num_rows());
1042        assert_eq!(
1043            &[1, 2, 3],
1044            record_batch.column(0).as_primitive::<Int32Type>().values()
1045        );
1046        let b = record_batch.column(1).as_string::<i32>();
1047        assert_eq!(
1048            vec![Some("x"), None, Some("z")],
1049            b.iter().collect::<Vec<_>>()
1050        );
1051        assert_eq!(
1052            &[7, 8, 9],
1053            record_batch.column(2).as_primitive::<Int64Type>().values()
1054        );
1055    }
1056
1057    #[test]
1058    fn create_record_batch_macro_parameterized_types() {
1059        let record_batch = record_batch!(
1060            ("ts", TimestampSecond, [1, 2]),
1061            ("d", Decimal128, [Some(100), None])
1062        )
1063        .unwrap();
1064
1065        let expected = Schema::new(vec![
1066            Field::new("ts", DataType::Timestamp(TimeUnit::Second, None), true),
1067            Field::new("d", DataType::Decimal128(38, 10), true),
1068        ]);
1069        assert_eq!(&expected, record_batch.schema().as_ref());
1070        assert_eq!(1, record_batch.column(1).null_count());
1071    }
1072
1073    #[test]
1074    fn create_record_batch_macro_date_types() {
1075        let record_batch = record_batch!(
1076            ("d32", Date32, [Some(1), None]),
1077            ("d64", Date64, [86_400_000, 172_800_000])
1078        )
1079        .unwrap();
1080
1081        let expected = Schema::new(vec![
1082            Field::new("d32", DataType::Date32, true),
1083            Field::new("d64", DataType::Date64, true),
1084        ]);
1085        assert_eq!(&expected, record_batch.schema().as_ref());
1086        assert_eq!(1, record_batch.column(0).null_count());
1087        assert_eq!(
1088            &[86_400_000, 172_800_000],
1089            record_batch.column(1).as_primitive::<Date64Type>().values()
1090        );
1091    }
1092
1093    #[test]
1094    fn create_record_batch_macro_length_mismatch() {
1095        let err = record_batch!(("a", Int32, [1, 2, 3]), ("b", Int32, [1, 2])).unwrap_err();
1096        assert!(matches!(err, ArrowError::InvalidArgumentError(_)));
1097    }
1098
1099    #[test]
1100    fn byte_size_should_not_regress() {
1101        let schema = Schema::new(vec![
1102            Field::new("a", DataType::Int32, false),
1103            Field::new("b", DataType::Utf8, false),
1104        ]);
1105
1106        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
1107        let b = StringArray::from(vec!["a", "b", "c", "d", "e"]);
1108
1109        let record_batch =
1110            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
1111        assert_eq!(record_batch.get_array_memory_size(), 364);
1112    }
1113
1114    fn check_batch(record_batch: RecordBatch, num_rows: usize) {
1115        assert_eq!(num_rows, record_batch.num_rows());
1116        assert_eq!(2, record_batch.num_columns());
1117        assert_eq!(&DataType::Int32, record_batch.schema().field(0).data_type());
1118        assert_eq!(&DataType::Utf8, record_batch.schema().field(1).data_type());
1119        assert_eq!(num_rows, record_batch.column(0).len());
1120        assert_eq!(num_rows, record_batch.column(1).len());
1121    }
1122
1123    #[test]
1124    #[should_panic(expected = "assertion failed: (offset + length) <= self.num_rows()")]
1125    fn create_record_batch_slice() {
1126        let schema = Schema::new(vec![
1127            Field::new("a", DataType::Int32, false),
1128            Field::new("b", DataType::Utf8, false),
1129        ]);
1130        let expected_schema = schema.clone();
1131
1132        let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
1133        let b = StringArray::from(vec!["a", "b", "c", "d", "e", "f", "h", "i"]);
1134
1135        let record_batch =
1136            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
1137
1138        let offset = 2;
1139        let length = 5;
1140        let record_batch_slice = record_batch.slice(offset, length);
1141
1142        assert_eq!(record_batch_slice.schema().as_ref(), &expected_schema);
1143        check_batch(record_batch_slice, 5);
1144
1145        let offset = 2;
1146        let length = 0;
1147        let record_batch_slice = record_batch.slice(offset, length);
1148
1149        assert_eq!(record_batch_slice.schema().as_ref(), &expected_schema);
1150        check_batch(record_batch_slice, 0);
1151
1152        let offset = 2;
1153        let length = 10;
1154        let _record_batch_slice = record_batch.slice(offset, length);
1155    }
1156
1157    #[test]
1158    #[should_panic(expected = "assertion failed: (offset + length) <= self.num_rows()")]
1159    fn create_record_batch_slice_empty_batch() {
1160        let schema = Schema::empty();
1161
1162        let record_batch = RecordBatch::new_empty(Arc::new(schema));
1163
1164        let offset = 0;
1165        let length = 0;
1166        let record_batch_slice = record_batch.slice(offset, length);
1167        assert_eq!(0, record_batch_slice.schema().fields().len());
1168
1169        let offset = 1;
1170        let length = 2;
1171        let _record_batch_slice = record_batch.slice(offset, length);
1172    }
1173
1174    #[test]
1175    fn create_record_batch_try_from_iter() {
1176        let a: ArrayRef = Arc::new(Int32Array::from(vec![
1177            Some(1),
1178            Some(2),
1179            None,
1180            Some(4),
1181            Some(5),
1182        ]));
1183        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));
1184
1185        let record_batch =
1186            RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).expect("valid conversion");
1187
1188        let expected_schema = Schema::new(vec![
1189            Field::new("a", DataType::Int32, true),
1190            Field::new("b", DataType::Utf8, false),
1191        ]);
1192        assert_eq!(record_batch.schema().as_ref(), &expected_schema);
1193        check_batch(record_batch, 5);
1194    }
1195
1196    #[test]
1197    fn create_record_batch_try_from_iter_with_nullable() {
1198        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
1199        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));
1200
1201        // Note there are no nulls in a or b, but we specify that b is nullable
1202        let record_batch =
1203            RecordBatch::try_from_iter_with_nullable(vec![("a", a, false), ("b", b, true)])
1204                .expect("valid conversion");
1205
1206        let expected_schema = Schema::new(vec![
1207            Field::new("a", DataType::Int32, false),
1208            Field::new("b", DataType::Utf8, true),
1209        ]);
1210        assert_eq!(record_batch.schema().as_ref(), &expected_schema);
1211        check_batch(record_batch, 5);
1212    }
1213
1214    #[test]
1215    fn create_record_batch_schema_mismatch() {
1216        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1217
1218        let a = Int64Array::from(vec![1, 2, 3, 4, 5]);
1219
1220        let err = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap_err();
1221        assert_eq!(
1222            err.to_string(),
1223            "Invalid argument error: column types must match schema types, expected Int32 but found Int64 at column index 0"
1224        );
1225    }
1226
1227    #[test]
1228    fn create_record_batch_field_name_mismatch() {
1229        let fields = vec![
1230            Field::new("a1", DataType::Int32, false),
1231            Field::new_list("a2", Field::new_list_field(DataType::Int8, false), false),
1232        ];
1233        let schema = Arc::new(Schema::new(vec![Field::new_struct("a", fields, true)]));
1234
1235        let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
1236        let a2_child = Int8Array::from(vec![1, 2, 3, 4]);
1237        let a2 = ArrayDataBuilder::new(DataType::List(Arc::new(Field::new(
1238            "array",
1239            DataType::Int8,
1240            false,
1241        ))))
1242        .add_child_data(a2_child.into_data())
1243        .len(2)
1244        .add_buffer(Buffer::from([0i32, 3, 4].to_byte_slice()))
1245        .build()
1246        .unwrap();
1247        let a2: ArrayRef = Arc::new(ListArray::from(a2));
1248        let a = ArrayDataBuilder::new(DataType::Struct(Fields::from(vec![
1249            Field::new("aa1", DataType::Int32, false),
1250            Field::new("a2", a2.data_type().clone(), false),
1251        ])))
1252        .add_child_data(a1.into_data())
1253        .add_child_data(a2.into_data())
1254        .len(2)
1255        .build()
1256        .unwrap();
1257        let a: ArrayRef = Arc::new(StructArray::from(a));
1258
1259        // creating the batch with field name validation should fail
1260        let batch = RecordBatch::try_new(schema.clone(), vec![a.clone()]);
1261        assert!(batch.is_err());
1262
1263        // creating the batch without field name validation should pass
1264        let options = RecordBatchOptions {
1265            match_field_names: false,
1266            row_count: None,
1267        };
1268        let batch = RecordBatch::try_new_with_options(schema, vec![a], &options);
1269        assert!(batch.is_ok());
1270    }
1271
1272    #[test]
1273    fn create_record_batch_record_mismatch() {
1274        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1275
1276        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
1277        let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
1278
1279        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]);
1280        assert!(batch.is_err());
1281    }
1282
1283    #[test]
1284    fn create_record_batch_from_struct_array() {
1285        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
1286        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
1287        let struct_array = StructArray::from(vec![
1288            (
1289                Arc::new(Field::new("b", DataType::Boolean, false)),
1290                boolean.clone() as ArrayRef,
1291            ),
1292            (
1293                Arc::new(Field::new("c", DataType::Int32, false)),
1294                int.clone() as ArrayRef,
1295            ),
1296        ]);
1297
1298        let batch = RecordBatch::from(&struct_array);
1299        assert_eq!(2, batch.num_columns());
1300        assert_eq!(4, batch.num_rows());
1301        assert_eq!(
1302            struct_array.data_type(),
1303            &DataType::Struct(batch.schema().fields().clone())
1304        );
1305        assert_eq!(batch.column(0).as_ref(), boolean.as_ref());
1306        assert_eq!(batch.column(1).as_ref(), int.as_ref());
1307    }
1308
1309    #[test]
1310    fn record_batch_equality() {
1311        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
1312        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
1313        let schema1 = Schema::new(vec![
1314            Field::new("id", DataType::Int32, false),
1315            Field::new("val", DataType::Int32, false),
1316        ]);
1317
1318        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1319        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
1320        let schema2 = Schema::new(vec![
1321            Field::new("id", DataType::Int32, false),
1322            Field::new("val", DataType::Int32, false),
1323        ]);
1324
1325        let batch1 = RecordBatch::try_new(
1326            Arc::new(schema1),
1327            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
1328        )
1329        .unwrap();
1330
1331        let batch2 = RecordBatch::try_new(
1332            Arc::new(schema2),
1333            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
1334        )
1335        .unwrap();
1336
1337        assert_eq!(batch1, batch2);
1338    }
1339
1340    /// validates if the record batch can be accessed using `column_name` as index i.e. `record_batch["column_name"]`
1341    #[test]
1342    fn record_batch_index_access() {
1343        let id_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1344        let val_arr = Arc::new(Int32Array::from(vec![5, 6, 7, 8]));
1345        let schema1 = Schema::new(vec![
1346            Field::new("id", DataType::Int32, false),
1347            Field::new("val", DataType::Int32, false),
1348        ]);
1349        let record_batch =
1350            RecordBatch::try_new(Arc::new(schema1), vec![id_arr.clone(), val_arr.clone()]).unwrap();
1351
1352        assert_eq!(record_batch["id"].as_ref(), id_arr.as_ref());
1353        assert_eq!(record_batch["val"].as_ref(), val_arr.as_ref());
1354    }
1355
1356    #[test]
1357    fn record_batch_vals_ne() {
1358        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
1359        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
1360        let schema1 = Schema::new(vec![
1361            Field::new("id", DataType::Int32, false),
1362            Field::new("val", DataType::Int32, false),
1363        ]);
1364
1365        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1366        let val_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1367        let schema2 = Schema::new(vec![
1368            Field::new("id", DataType::Int32, false),
1369            Field::new("val", DataType::Int32, false),
1370        ]);
1371
1372        let batch1 = RecordBatch::try_new(
1373            Arc::new(schema1),
1374            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
1375        )
1376        .unwrap();
1377
1378        let batch2 = RecordBatch::try_new(
1379            Arc::new(schema2),
1380            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
1381        )
1382        .unwrap();
1383
1384        assert_ne!(batch1, batch2);
1385    }
1386
1387    #[test]
1388    fn record_batch_column_names_ne() {
1389        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
1390        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
1391        let schema1 = Schema::new(vec![
1392            Field::new("id", DataType::Int32, false),
1393            Field::new("val", DataType::Int32, false),
1394        ]);
1395
1396        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1397        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
1398        let schema2 = Schema::new(vec![
1399            Field::new("id", DataType::Int32, false),
1400            Field::new("num", DataType::Int32, false),
1401        ]);
1402
1403        let batch1 = RecordBatch::try_new(
1404            Arc::new(schema1),
1405            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
1406        )
1407        .unwrap();
1408
1409        let batch2 = RecordBatch::try_new(
1410            Arc::new(schema2),
1411            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
1412        )
1413        .unwrap();
1414
1415        assert_ne!(batch1, batch2);
1416    }
1417
1418    #[test]
1419    fn record_batch_column_number_ne() {
1420        let id_arr1 = Int32Array::from(vec![1, 2, 3, 4]);
1421        let val_arr1 = Int32Array::from(vec![5, 6, 7, 8]);
1422        let schema1 = Schema::new(vec![
1423            Field::new("id", DataType::Int32, false),
1424            Field::new("val", DataType::Int32, false),
1425        ]);
1426
1427        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1428        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
1429        let num_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
1430        let schema2 = Schema::new(vec![
1431            Field::new("id", DataType::Int32, false),
1432            Field::new("val", DataType::Int32, false),
1433            Field::new("num", DataType::Int32, false),
1434        ]);
1435
1436        let batch1 = RecordBatch::try_new(
1437            Arc::new(schema1),
1438            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
1439        )
1440        .unwrap();
1441
1442        let batch2 = RecordBatch::try_new(
1443            Arc::new(schema2),
1444            vec![Arc::new(id_arr2), Arc::new(val_arr2), Arc::new(num_arr2)],
1445        )
1446        .unwrap();
1447
1448        assert_ne!(batch1, batch2);
1449    }
1450
1451    #[test]
1452    fn record_batch_row_count_ne() {
1453        let id_arr1 = Int32Array::from(vec![1, 2, 3]);
1454        let val_arr1 = Int32Array::from(vec![5, 6, 7]);
1455        let schema1 = Schema::new(vec![
1456            Field::new("id", DataType::Int32, false),
1457            Field::new("val", DataType::Int32, false),
1458        ]);
1459
1460        let id_arr2 = Int32Array::from(vec![1, 2, 3, 4]);
1461        let val_arr2 = Int32Array::from(vec![5, 6, 7, 8]);
1462        let schema2 = Schema::new(vec![
1463            Field::new("id", DataType::Int32, false),
1464            Field::new("num", DataType::Int32, false),
1465        ]);
1466
1467        let batch1 = RecordBatch::try_new(
1468            Arc::new(schema1),
1469            vec![Arc::new(id_arr1), Arc::new(val_arr1)],
1470        )
1471        .unwrap();
1472
1473        let batch2 = RecordBatch::try_new(
1474            Arc::new(schema2),
1475            vec![Arc::new(id_arr2), Arc::new(val_arr2)],
1476        )
1477        .unwrap();
1478
1479        assert_ne!(batch1, batch2);
1480    }
1481
1482    #[test]
1483    fn normalize_simple() {
1484        let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""]));
1485        let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)]));
1486        let year: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2022)]));
1487
1488        let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
1489        let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
1490        let year_field = Arc::new(Field::new("year", DataType::Int64, true));
1491
1492        let a = Arc::new(StructArray::from(vec![
1493            (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef),
1494            (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef),
1495            (year_field.clone(), Arc::new(year.clone()) as ArrayRef),
1496        ]));
1497
1498        let month = Arc::new(Int64Array::from(vec![Some(4), Some(6)]));
1499
1500        let schema = Schema::new(vec![
1501            Field::new(
1502                "a",
1503                DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])),
1504                false,
1505            ),
1506            Field::new("month", DataType::Int64, true),
1507        ]);
1508
1509        let normalized =
1510            RecordBatch::try_new(Arc::new(schema.clone()), vec![a.clone(), month.clone()])
1511                .expect("valid conversion")
1512                .normalize(".", Some(0))
1513                .expect("valid normalization");
1514
1515        let expected = RecordBatch::try_from_iter_with_nullable(vec![
1516            ("a.animals", animals.clone(), true),
1517            ("a.n_legs", n_legs.clone(), true),
1518            ("a.year", year.clone(), true),
1519            ("month", month.clone(), true),
1520        ])
1521        .expect("valid conversion");
1522
1523        assert_eq!(expected, normalized);
1524
1525        // check 0 and None have the same effect
1526        let normalized = RecordBatch::try_new(Arc::new(schema), vec![a, month.clone()])
1527            .expect("valid conversion")
1528            .normalize(".", None)
1529            .expect("valid normalization");
1530
1531        assert_eq!(expected, normalized);
1532    }
1533
1534    #[test]
1535    fn normalize_nested() {
1536        // Initialize schema
1537        let a = Arc::new(Field::new("a", DataType::Int64, true));
1538        let b = Arc::new(Field::new("b", DataType::Int64, false));
1539        let c = Arc::new(Field::new("c", DataType::Int64, true));
1540
1541        let one = Arc::new(Field::new(
1542            "1",
1543            DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])),
1544            false,
1545        ));
1546        let two = Arc::new(Field::new(
1547            "2",
1548            DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])),
1549            true,
1550        ));
1551
1552        let exclamation = Arc::new(Field::new(
1553            "!",
1554            DataType::Struct(Fields::from(vec![one.clone(), two.clone()])),
1555            false,
1556        ));
1557
1558        let schema = Schema::new(vec![exclamation.clone()]);
1559
1560        // Initialize fields
1561        let a_field = Int64Array::from(vec![Some(0), Some(1)]);
1562        let b_field = Int64Array::from(vec![Some(2), Some(3)]);
1563        let c_field = Int64Array::from(vec![None, Some(4)]);
1564
1565        let one_field = StructArray::from(vec![
1566            (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
1567            (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
1568            (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
1569        ]);
1570        let two_field = StructArray::from(vec![
1571            (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
1572            (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
1573            (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
1574        ]);
1575
1576        let exclamation_field = Arc::new(StructArray::from(vec![
1577            (one.clone(), Arc::new(one_field) as ArrayRef),
1578            (two.clone(), Arc::new(two_field) as ArrayRef),
1579        ]));
1580
1581        // Normalize top level
1582        let normalized =
1583            RecordBatch::try_new(Arc::new(schema.clone()), vec![exclamation_field.clone()])
1584                .expect("valid conversion")
1585                .normalize(".", Some(1))
1586                .expect("valid normalization");
1587
1588        let expected = RecordBatch::try_from_iter_with_nullable(vec![
1589            (
1590                "!.1",
1591                Arc::new(StructArray::from(vec![
1592                    (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
1593                    (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
1594                    (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
1595                ])) as ArrayRef,
1596                false,
1597            ),
1598            (
1599                "!.2",
1600                Arc::new(StructArray::from(vec![
1601                    (a.clone(), Arc::new(a_field.clone()) as ArrayRef),
1602                    (b.clone(), Arc::new(b_field.clone()) as ArrayRef),
1603                    (c.clone(), Arc::new(c_field.clone()) as ArrayRef),
1604                ])) as ArrayRef,
1605                true,
1606            ),
1607        ])
1608        .expect("valid conversion");
1609
1610        assert_eq!(expected, normalized);
1611
1612        // Normalize all levels
1613        let normalized = RecordBatch::try_new(Arc::new(schema), vec![exclamation_field])
1614            .expect("valid conversion")
1615            .normalize(".", None)
1616            .expect("valid normalization");
1617
1618        let expected = RecordBatch::try_from_iter_with_nullable(vec![
1619            ("!.1.a", Arc::new(a_field.clone()) as ArrayRef, true),
1620            ("!.1.b", Arc::new(b_field.clone()) as ArrayRef, false),
1621            ("!.1.c", Arc::new(c_field.clone()) as ArrayRef, true),
1622            ("!.2.a", Arc::new(a_field.clone()) as ArrayRef, true),
1623            ("!.2.b", Arc::new(b_field.clone()) as ArrayRef, false),
1624            ("!.2.c", Arc::new(c_field.clone()) as ArrayRef, true),
1625        ])
1626        .expect("valid conversion");
1627
1628        assert_eq!(expected, normalized);
1629    }
1630
1631    #[test]
1632    fn normalize_empty() {
1633        let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true));
1634        let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true));
1635        let year_field = Arc::new(Field::new("year", DataType::Int64, true));
1636
1637        let schema = Schema::new(vec![
1638            Field::new(
1639                "a",
1640                DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])),
1641                false,
1642            ),
1643            Field::new("month", DataType::Int64, true),
1644        ]);
1645
1646        let normalized = RecordBatch::new_empty(Arc::new(schema.clone()))
1647            .normalize(".", Some(0))
1648            .expect("valid normalization");
1649
1650        let expected = RecordBatch::new_empty(Arc::new(
1651            schema.normalize(".", Some(0)).expect("valid normalization"),
1652        ));
1653
1654        assert_eq!(expected, normalized);
1655    }
1656
1657    #[test]
1658    fn project() {
1659        let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]));
1660        let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
1661        let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));
1662
1663        let record_batch =
1664            RecordBatch::try_from_iter(vec![("a", a.clone()), ("b", b.clone()), ("c", c.clone())])
1665                .expect("valid conversion");
1666
1667        let expected =
1668            RecordBatch::try_from_iter(vec![("a", a), ("c", c)]).expect("valid conversion");
1669
1670        assert_eq!(expected, record_batch.project(&[0, 2]).unwrap());
1671    }
1672
1673    #[test]
1674    fn project_empty() {
1675        let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));
1676
1677        let record_batch =
1678            RecordBatch::try_from_iter(vec![("c", c.clone())]).expect("valid conversion");
1679
1680        let expected = RecordBatch::try_new_with_options(
1681            Arc::new(Schema::empty()),
1682            vec![],
1683            &RecordBatchOptions {
1684                match_field_names: true,
1685                row_count: Some(3),
1686            },
1687        )
1688        .expect("valid conversion");
1689
1690        assert_eq!(expected, record_batch.project(&[]).unwrap());
1691    }
1692
1693    #[test]
1694    fn test_no_column_record_batch() {
1695        let schema = Arc::new(Schema::empty());
1696
1697        let err = RecordBatch::try_new(schema.clone(), vec![]).unwrap_err();
1698        assert!(
1699            err.to_string()
1700                .contains("must either specify a row count or at least one column")
1701        );
1702
1703        let options = RecordBatchOptions::new().with_row_count(Some(10));
1704
1705        let ok = RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
1706        assert_eq!(ok.num_rows(), 10);
1707
1708        let a = ok.slice(2, 5);
1709        assert_eq!(a.num_rows(), 5);
1710
1711        let b = ok.slice(5, 0);
1712        assert_eq!(b.num_rows(), 0);
1713
1714        assert_ne!(a, b);
1715        assert_eq!(b, RecordBatch::new_empty(schema))
1716    }
1717
1718    #[test]
1719    fn test_nulls_in_non_nullable_field() {
1720        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1721        let maybe_batch = RecordBatch::try_new(
1722            schema,
1723            vec![Arc::new(Int32Array::from(vec![Some(1), None]))],
1724        );
1725        assert_eq!(
1726            "Invalid argument error: Column 'a' is declared as non-nullable but contains null values",
1727            format!("{}", maybe_batch.err().unwrap())
1728        );
1729    }
1730    #[test]
1731    fn test_record_batch_options() {
1732        let options = RecordBatchOptions::new()
1733            .with_match_field_names(false)
1734            .with_row_count(Some(20));
1735        assert!(!options.match_field_names);
1736        assert_eq!(options.row_count.unwrap(), 20)
1737    }
1738
1739    #[test]
1740    #[should_panic(expected = "Cannot convert nullable StructArray to RecordBatch")]
1741    fn test_from_struct() {
1742        let s = StructArray::from(ArrayData::new_null(
1743            // Note child is not nullable
1744            &DataType::Struct(vec![Field::new("foo", DataType::Int32, false)].into()),
1745            2,
1746        ));
1747        let _ = RecordBatch::from(s);
1748    }
1749
1750    #[test]
1751    fn test_with_schema() {
1752        let required_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1753        let required_schema = Arc::new(required_schema);
1754        let nullable_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
1755        let nullable_schema = Arc::new(nullable_schema);
1756
1757        let batch = RecordBatch::try_new(
1758            required_schema.clone(),
1759            vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as _],
1760        )
1761        .unwrap();
1762
1763        // Can add nullability
1764        let batch = batch.with_schema(nullable_schema.clone()).unwrap();
1765
1766        // Cannot remove nullability
1767        batch.clone().with_schema(required_schema).unwrap_err();
1768
1769        // Can add metadata
1770        let metadata = arrow_schema::Metadata::from([("foo", "bar")]);
1771        let metadata_schema = nullable_schema.as_ref().clone().with_metadata(metadata);
1772        let batch = batch.with_schema(Arc::new(metadata_schema)).unwrap();
1773
1774        // Cannot remove metadata
1775        batch.with_schema(nullable_schema).unwrap_err();
1776    }
1777
1778    #[test]
1779    fn test_boxed_reader() {
1780        // Make sure we can pass a boxed reader to a function generic over
1781        // RecordBatchReader.
1782        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1783        let schema = Arc::new(schema);
1784
1785        let reader = RecordBatchIterator::new(std::iter::empty(), schema);
1786        let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);
1787
1788        fn get_size(reader: impl RecordBatchReader) -> usize {
1789            reader.size_hint().0
1790        }
1791
1792        let size = get_size(reader);
1793        assert_eq!(size, 0);
1794    }
1795
1796    #[test]
1797    fn test_remove_column_maintains_schema_metadata() {
1798        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
1799        let bool_array = BooleanArray::from(vec![true, false, false, true, true]);
1800
1801        let mut metadata = HashMap::new();
1802        metadata.insert("foo".to_string(), "bar".to_string());
1803        let schema = Schema::new(vec![
1804            Field::new("id", DataType::Int32, false),
1805            Field::new("bool", DataType::Boolean, false),
1806        ])
1807        .with_metadata(metadata);
1808
1809        let mut batch = RecordBatch::try_new(
1810            Arc::new(schema),
1811            vec![Arc::new(id_array), Arc::new(bool_array)],
1812        )
1813        .unwrap();
1814
1815        let _removed_column = batch.remove_column(0);
1816        assert_eq!(batch.schema().metadata().len(), 1);
1817        assert_eq!(
1818            batch.schema().metadata().get("foo").unwrap().as_str(),
1819            "bar"
1820        );
1821    }
1822
1823    #[test]
1824    fn test_normalize_nullable_struct() {
1825        let child = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1826        let struct_nulls =
1827            NullBuffer::new(arrow_buffer::BooleanBuffer::from(vec![true, false, true]));
1828        let struct_array = Arc::new(StructArray::new(
1829            Fields::from(vec![Field::new("x", DataType::Int32, false)]),
1830            vec![child],
1831            Some(struct_nulls),
1832        )) as ArrayRef;
1833
1834        let schema = Schema::new(vec![Field::new(
1835            "s",
1836            DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)])),
1837            true,
1838        )]);
1839        let batch = RecordBatch::try_new(Arc::new(schema), vec![struct_array]).unwrap();
1840
1841        let normalized = batch.normalize(".", None).unwrap();
1842
1843        assert_eq!(normalized.num_columns(), 1);
1844        assert_eq!(normalized.schema().field(0).name(), "s.x");
1845        assert!(normalized.schema().field(0).is_nullable());
1846        let col = normalized.column(0);
1847        assert!(col.is_valid(0));
1848        assert!(col.is_null(1));
1849        assert!(col.is_valid(2));
1850    }
1851}
arrow_array/record_batch.rs

arrow_array/
record_batch.rs
