rs_arrow_ipc_stream_count/
lib.rs

1use std::io;
2
3use io::BufReader;
4use io::Read;
5
6use io::BufWriter;
7use io::Write;
8
9use arrow_schema::{DataType, Field, Schema};
10
11use arrow_array::RecordBatch;
12use arrow_array::record_batch;
13
14use arrow_ipc::reader::StreamReader;
15use arrow_ipc::writer::StreamWriter;
16
17pub struct IpcStreamReader<R>(pub StreamReader<BufReader<R>>);
18
19impl<R> IpcStreamReader<R>
20where
21    R: Read,
22{
23    pub fn into_bat_iter(self) -> impl Iterator<Item = Result<RecordBatch, io::Error>> {
24        self.0.map(|r| r.map_err(io::Error::other))
25    }
26}
27
28impl<R> IpcStreamReader<R>
29where
30    R: Read,
31{
32    pub fn from_reader(rdr: R) -> Result<Self, io::Error> {
33        let srdr = StreamReader::try_new_buffered(rdr, None).map_err(io::Error::other)?;
34        Ok(Self(srdr))
35    }
36}
37
38#[derive(Debug, Default)]
39pub struct BasicCount {
40    pub num_record_batch: usize,
41    pub num_rows: usize,
42}
43
44impl BasicCount {
45    pub fn from_bat_iter<I>(mut i: I) -> Result<Self, io::Error>
46    where
47        I: Iterator<Item = Result<RecordBatch, io::Error>>,
48    {
49        i.try_fold(Self::default(), |state, next| {
50            let rbat: RecordBatch = next?;
51            let num_rows: usize = rbat.num_rows();
52            Ok(Self {
53                num_record_batch: state.num_record_batch + 1,
54                num_rows: state.num_rows + num_rows,
55            })
56        })
57    }
58}
59
60impl BasicCount {
61    pub fn schema() -> Schema {
62        Schema::new(vec![
63            Field::new("num_record_batch", DataType::UInt64, false),
64            Field::new("num_rows", DataType::UInt64, false),
65        ])
66    }
67
68    pub fn to_stream_writer<W>(&self, mut wtr: StreamWriter<BufWriter<W>>) -> Result<(), io::Error>
69    where
70        W: Write,
71    {
72        let rbat: RecordBatch = record_batch!(
73            ("num_record_batch", Int64, [self.num_record_batch as i64]),
74            ("num_rows", Int64, [self.num_rows as i64])
75        )
76        .map_err(io::Error::other)?;
77
78        wtr.write(&rbat).map_err(io::Error::other)?;
79
80        wtr.flush().map_err(io::Error::other)?;
81        wtr.finish().map_err(io::Error::other)
82    }
83
84    pub fn to_writer<W>(&self, wtr: W) -> Result<(), io::Error>
85    where
86        W: Write,
87    {
88        let s: Schema = Self::schema();
89        let swtr = StreamWriter::try_new_buffered(wtr, &s).map_err(io::Error::other)?;
90        self.to_stream_writer(swtr)?;
91        Ok(())
92    }
93}
94
95impl BasicCount {
96    pub fn from_reader<R>(rdr: R) -> Result<Self, io::Error>
97    where
98        R: Read,
99    {
100        let irdr = IpcStreamReader::from_reader(rdr)?;
101        let ibat = irdr.into_bat_iter();
102        Self::from_bat_iter(ibat)
103    }
104}
105
106impl BasicCount {
107    pub fn reader2count2writer<R, W>(rdr: R, wtr: W) -> Result<(), io::Error>
108    where
109        R: Read,
110        W: Write,
111    {
112        let bcnt: Self = Self::from_reader(rdr)?;
113        bcnt.to_writer(wtr)
114    }
115}
116
117impl BasicCount {
118    pub fn stdin2count2stdout() -> Result<(), io::Error> {
119        let r = io::stdin().lock();
120        let mut w = io::stdout().lock();
121        Self::reader2count2writer(r, &mut w)?;
122        w.flush()
123    }
124}
rs_arrow_ipc_stream_count/lib.rs

rs_arrow_ipc_stream_count/
lib.rs
