rs-arrow-ipc-stream-count 0.1.0

use std::io;

use io::BufReader;
use io::Read;

use io::BufWriter;
use io::Write;

use arrow_schema::{DataType, Field, Schema};

use arrow_array::RecordBatch;
use arrow_array::record_batch;

use arrow_ipc::reader::StreamReader;
use arrow_ipc::writer::StreamWriter;

pub struct IpcStreamReader<R>(pub StreamReader<BufReader<R>>);

impl<R> IpcStreamReader<R>
where
    R: Read,
{
    pub fn into_bat_iter(self) -> impl Iterator<Item = Result<RecordBatch, io::Error>> {
        self.0.map(|r| r.map_err(io::Error::other))
    }
}

impl<R> IpcStreamReader<R>
where
    R: Read,
{
    pub fn from_reader(rdr: R) -> Result<Self, io::Error> {
        let srdr = StreamReader::try_new_buffered(rdr, None).map_err(io::Error::other)?;
        Ok(Self(srdr))
    }
}

#[derive(Debug, Default)]
pub struct BasicCount {
    pub num_record_batch: usize,
    pub num_rows: usize,
}

impl BasicCount {
    pub fn from_bat_iter<I>(mut i: I) -> Result<Self, io::Error>
    where
        I: Iterator<Item = Result<RecordBatch, io::Error>>,
    {
        i.try_fold(Self::default(), |state, next| {
            let rbat: RecordBatch = next?;
            let num_rows: usize = rbat.num_rows();
            Ok(Self {
                num_record_batch: state.num_record_batch + 1,
                num_rows: state.num_rows + num_rows,
            })
        })
    }
}

impl BasicCount {
    pub fn schema() -> Schema {
        Schema::new(vec![
            Field::new("num_record_batch", DataType::UInt64, false),
            Field::new("num_rows", DataType::UInt64, false),
        ])
    }

    pub fn to_stream_writer<W>(&self, mut wtr: StreamWriter<BufWriter<W>>) -> Result<(), io::Error>
    where
        W: Write,
    {
        let rbat: RecordBatch = record_batch!(
            ("num_record_batch", Int64, [self.num_record_batch as i64]),
            ("num_rows", Int64, [self.num_rows as i64])
        )
        .map_err(io::Error::other)?;

        wtr.write(&rbat).map_err(io::Error::other)?;

        wtr.flush().map_err(io::Error::other)?;
        wtr.finish().map_err(io::Error::other)
    }

    pub fn to_writer<W>(&self, wtr: W) -> Result<(), io::Error>
    where
        W: Write,
    {
        let s: Schema = Self::schema();
        let swtr = StreamWriter::try_new_buffered(wtr, &s).map_err(io::Error::other)?;
        self.to_stream_writer(swtr)?;
        Ok(())
    }
}

impl BasicCount {
    pub fn from_reader<R>(rdr: R) -> Result<Self, io::Error>
    where
        R: Read,
    {
        let irdr = IpcStreamReader::from_reader(rdr)?;
        let ibat = irdr.into_bat_iter();
        Self::from_bat_iter(ibat)
    }
}

impl BasicCount {
    pub fn reader2count2writer<R, W>(rdr: R, wtr: W) -> Result<(), io::Error>
    where
        R: Read,
        W: Write,
    {
        let bcnt: Self = Self::from_reader(rdr)?;
        bcnt.to_writer(wtr)
    }
}

impl BasicCount {
    pub fn stdin2count2stdout() -> Result<(), io::Error> {
        let r = io::stdin().lock();
        let mut w = io::stdout().lock();
        Self::reader2count2writer(r, &mut w)?;
        w.flush()
    }
}