arrow_pg/
row_encoder.rs

1use std::sync::Arc;
2
3use arrow::array::RecordBatch;
4use pgwire::{
5    api::results::{DataRowEncoder, FieldInfo},
6    error::PgWireResult,
7    messages::data::DataRow,
8};
9
10use crate::encoder::encode_value;
11
12pub struct RowEncoder {
13    rb: RecordBatch,
14    curr_idx: usize,
15    fields: Arc<Vec<FieldInfo>>,
16}
17
18impl RowEncoder {
19    pub fn new(rb: RecordBatch, fields: Arc<Vec<FieldInfo>>) -> Self {
20        assert_eq!(rb.num_columns(), fields.len());
21        Self {
22            rb,
23            fields,
24            curr_idx: 0,
25        }
26    }
27
28    pub fn next_row(&mut self) -> Option<PgWireResult<DataRow>> {
29        if self.curr_idx == self.rb.num_rows() {
30            return None;
31        }
32        let mut encoder = DataRowEncoder::new(self.fields.clone());
33        for col in 0..self.rb.num_columns() {
34            let array = self.rb.column(col);
35            let field = &self.fields[col];
36            let type_ = field.datatype();
37            let format = field.format();
38            encode_value(&mut encoder, array, self.curr_idx, type_, format).unwrap();
39        }
40        self.curr_idx += 1;
41        Some(encoder.finish())
42    }
43}
arrow_pg/row_encoder.rs

arrow_pg/
row_encoder.rs
