arrow RecordBatch構造体

RustのarrowのRecordBatchのメモ。

概要

RecordBatchは、Arrowのデータを扱う構造体。

RecordBatchを作る例

use std::sync::Arc;

use arrow::{
    array::{ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray},
    datatypes::{DataType, Field, Schema},
    error::ArrowError,
};

fn main() -> Result<(), ArrowError> {
    let field1 = Field::new("foo", DataType::Int32, false);
    let field2 = Field::new("bar", DataType::Int64, true);
    let field3 = Field::new("zzz", DataType::Utf8, true);
    let fields = vec![field1, field2, field3];
    let schema = Arc::new(Schema::new(fields));

    let array1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
    let array2: ArrayRef = Arc::new(Int64Array::from(vec![11, 22, 33, 44]));
    let array3: ArrayRef = Arc::new(StringArray::from(vec!["aaa", "bbb", "ccc", "ddd"]));
    let arrays = vec![array1, array2, array3];

    let batch = RecordBatch::try_new(schema, arrays)?;
    println!("{:?}", batch);

    Ok(())
}

RecordBatchを作るには、まずSchemaを用意する。
Schemaは複数のFieldで構成される。

Field::new()の第1引数はカラム名、第2引数はデータ型、第3引数はnullable。

データはarrow::array::Arrayで用意する。
データ型に応じて、Int32ArrayやStringArray等の構造体を使用する。

Builderの例

Int32ArrayやStringArray等は、Int32BuilderやStringBuilder等のビルダーを使って作ることも出来る。

use std::sync::Arc;

use arrow::{
    array::{ArrayRef, Int32Builder, Int64Builder, RecordBatch, StringBuilder},
    datatypes::{DataType, Field, Schema},
    error::ArrowError,
};

    let mut builder1 = Int32Builder::new();
    let mut builder2 = Int64Builder::new();
    let mut builder3 = StringBuilder::new();

    for (foo, bar, zzz) in [(1, 11, "aaa"), (2, 22, "bbb")] {
        builder1.append_value(foo);
        builder2.append_value(bar);
        builder3.append_value(zzz);
    }

    let array1: ArrayRef = Arc::new(builder1.finish());
    let array2: ArrayRef = Arc::new(builder2.finish());
    let array3: ArrayRef = Arc::new(builder3.finish());
    let arrays = vec![array1, array2, array3];

builderに対してappend_value()で値を追加していく。

append_null()でnullを追加することも出来る。