rs_arrow_ipc_stream_cat/
lib.rs

1use std::io;
2
3use io::BufReader;
4use io::Read;
5
6use arrow_array::RecordBatch;
7
8use arrow_ipc::reader::StreamReader;
9
10use arrow::util::pretty::print_batches;
11
12pub fn show_records<I>(rbats: I) -> Result<(), io::Error>
13where
14    I: Iterator<Item = Result<RecordBatch, io::Error>>,
15{
16    print_batches(
17        &rbats
18            .map(|r| r.map_err(|e| io::Error::other(format!("show-record error: {e}"))))
19            .collect::<Result<Vec<_>, _>>()?,
20    )
21    .map_err(|e| io::Error::other(format!("show-record error: {e}")))
22}
23
24pub fn stream2records_buf<R>(
25    rdr: BufReader<R>,
26    projection: Option<Vec<usize>>,
27) -> Result<impl Iterator<Item = Result<RecordBatch, io::Error>>, io::Error>
28where
29    R: Read,
30{
31    let srdr = StreamReader::try_new_buffered(rdr, projection)
32        .map_err(|e| io::Error::other(format!("buffered stream reader open error: {e}")))?;
33    Ok(srdr
34        .map(|r| r.map_err(|e| io::Error::other(format!("buffered stream reader map error: {e}")))))
35}
36
37pub fn stream2records<R>(
38    rdr: R,
39    projection: Option<Vec<usize>>,
40) -> Result<impl Iterator<Item = Result<RecordBatch, io::Error>>, io::Error>
41where
42    R: Read,
43{
44    let srdr = StreamReader::try_new(rdr, projection)
45        .map_err(|e| io::Error::other(format!("stream reader open error: {e}")))?;
46    Ok(srdr.map(|r| r.map_err(|e| io::Error::other(format!("stream reader map error: {e}")))))
47}
48
49pub fn stdin2records(
50    projection: Option<Vec<usize>>,
51) -> Result<impl Iterator<Item = Result<RecordBatch, io::Error>>, io::Error> {
52    stream2records(io::stdin().lock(), projection)
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use arrow_array::{Int32Array, RecordBatch};
59    use arrow_ipc::writer::StreamWriter;
60    use std::sync::Arc;
61
62    #[test]
63    fn test_show_records() {
64        // Create a RecordBatch
65        let schema = arrow_schema::Schema::new(vec![arrow_schema::Field::new(
66            "a",
67            arrow_schema::DataType::Int32,
68            false,
69        )]);
70        let a = Int32Array::from(vec![1, 2, 3]);
71        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
72
73        // Create an in-memory IPC stream
74        let mut stream = Vec::new();
75        let mut writer = StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
76        writer.write(&batch).unwrap();
77        writer.finish().unwrap();
78
79        // Read the batches from the stream
80        let records = stream2records(stream.as_slice(), None).unwrap();
81
82        // Show the records (this will just print to stdout, but we can check if it panics)
83        show_records(records).unwrap();
84    }
85}
rs_arrow_ipc_stream_cat/lib.rs

rs_arrow_ipc_stream_cat/
lib.rs
