arrow_pg/
row_encoder.rs

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

arrow_pg/
row_encoder.rs
