graphar 0.1.2

Apache GraphAr format reader/writer
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

use std::{fs::File, path::Path};

use arrow_array::RecordBatch;
use orc_rust::{ArrowReaderBuilder, ArrowWriterBuilder};

use crate::error::Result;

pub fn read(path: impl AsRef<Path>) -> Result<Vec<RecordBatch>> {
    let file = File::open(path)?;
    let reader = ArrowReaderBuilder::try_new(file)?.build();
    let mut batches = Vec::new();
    for batch in reader {
        batches.push(batch?);
    }
    Ok(batches)
}

pub fn write(path: impl AsRef<Path>, batches: &[RecordBatch]) -> Result<()> {
    if batches.is_empty() {
        return Ok(());
    }
    let file = File::create(path)?;
    let mut writer = ArrowWriterBuilder::new(file, batches[0].schema()).try_build()?;
    for batch in batches {
        writer.write(batch)?;
    }
    writer.close()?;
    Ok(())
}