feat: use FieldInfo in encode functions

Author: sunng87

Cargo.lock

@@ -20,7 +20,7 @@ chrono = { version = "0.4", features = ["std"] }
 datafusion = { version = "50", default-features = false }
 futures = "0.3"
 #pgwire = { version = "0.34", default-features = false }
-pgwire = { git = "https://github.com/sunng87/pgwire", rev = "d89089947a56ebfe2e89631925facdd7a85c25b4", default-features = false }
+pgwire = { git = "https://github.com/sunng87/pgwire", rev = "37a32a05d2aed55bd013f3c6f93d786368350e0b", default-features = false }
 postgres-types = "0.2"
 rust_decimal = { version = "1.39", features = ["db-postgres"] }
 tokio = { version = "1", default-features = false }

Cargo.toml

@@ -12,6 +12,7 @@ use chrono::{NaiveDate, NaiveDateTime};
 use datafusion::arrow::{array::*, datatypes::*};
 use pgwire::api::results::DataRowEncoder;
 use pgwire::api::results::FieldFormat;
+use pgwire::api::results::FieldInfo;
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use pgwire::types::ToSqlText;
 use postgres_types::{ToSql, Type};
@@ -291,9 +292,11 @@ pub fn encode_value<T: Encoder>(
     encoder: &mut T,
     arr: &Arc<dyn Array>,
     idx: usize,
-    type_: &Type,
-    format: FieldFormat,
+    pg_field: &FieldInfo,
 ) -> PgWireResult<()> {
+    let type_ = pg_field.datatype();
+    let format = pg_field.format();
+
     match arr.data_type() {
         DataType::Null => encoder.encode_field_with_type_and_format(&None::<i8>, type_, format)?,
         DataType::Boolean => {
@@ -413,6 +416,7 @@ pub fn encode_value<T: Encoder>(
                     let value = ts_array
                         .value_as_datetime_with_tz(idx, tz)
                         .map(|d| d.fixed_offset());
+
                     encoder.encode_field_with_type_and_format(&value, type_, format)?;
                 } else {
                     let value = ts_array.value_as_datetime(idx);
@@ -494,7 +498,7 @@ pub fn encode_value<T: Encoder>(
                 return encoder.encode_field_with_type_and_format(&None::<&[i8]>, type_, format);
             }
             let array = arr.as_any().downcast_ref::<ListArray>().unwrap().value(idx);
-            let value = encode_list(array, type_, format)?;
+            let value = encode_list(array, pg_field)?;
             encoder.encode_field_with_type_and_format(&value, type_, format)?
         }
         DataType::Struct(_) => {
@@ -506,7 +510,7 @@ pub fn encode_value<T: Encoder>(
                     ))));
                 }
             };
-            let value = encode_struct(arr, idx, fields, format)?;
+            let value = encode_struct(arr, idx, fields, pg_field)?;
             encoder.encode_field_with_type_and_format(&value, type_, format)?
         }
         DataType::Dictionary(_, value_type) => {
@@ -537,7 +541,7 @@ pub fn encode_value<T: Encoder>(
                     ))
                 })?;
 
-            encode_value(encoder, values, idx, type_, format)?
+            encode_value(encoder, values, idx, pg_field)?
         }
         _ => {
             return Err(PgWireError::ApiError(ToSqlError::from(format!(
@@ -588,7 +592,8 @@ mod tests {
 
         let mut encoder = MockEncoder::default();
 
-        let result = encode_value(&mut encoder, &dict_arr, 2, &Type::TEXT, FieldFormat::Text);
+        let pg_field = FieldInfo::new("x".to_string(), None, None, Type::TEXT, FieldFormat::Text);
+        let result = encode_value(&mut encoder, &dict_arr, 2, &pg_field);
 
         assert!(result.is_ok());
 

arrow-pg/src/encoder.rs

@@ -37,7 +37,7 @@ use datafusion::arrow::{
 
 use bytes::{BufMut, BytesMut};
 use chrono::{DateTime, TimeZone, Utc};
-use pgwire::api::results::FieldFormat;
+use pgwire::api::results::{FieldFormat, FieldInfo};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::types::{ToSqlText, QUOTE_ESCAPE};
 use postgres_types::{ToSql, Type};
@@ -106,11 +106,10 @@ fn encode_field<T: ToSql + ToSqlText>(
     Ok(EncodedValue { bytes })
 }
 
-pub(crate) fn encode_list(
-    arr: Arc<dyn Array>,
-    type_: &Type,
-    format: FieldFormat,
-) -> PgWireResult<EncodedValue> {
+pub(crate) fn encode_list(arr: Arc<dyn Array>, pg_field: &FieldInfo) -> PgWireResult<EncodedValue> {
+    let type_ = pg_field.datatype();
+    let format = pg_field.format();
+
     match arr.data_type() {
         DataType::Null => {
             let mut bytes = BytesMut::new();
@@ -406,7 +405,7 @@ pub(crate) fn encode_list(
             .map_err(ToSqlError::from)?;
 
             let values: PgWireResult<Vec<_>> = (0..arr.len())
-                .map(|row| encode_struct(&arr, row, fields, format))
+                .map(|row| encode_struct(&arr, row, fields, pg_field))
                 .map(|x| {
                     if matches!(format, FieldFormat::Text) {
                         x.map(|opt| {

arrow-pg/src/list_encoder.rs

@@ -37,9 +37,8 @@ impl RowEncoder {
         for col in 0..self.rb.num_columns() {
             let array = self.rb.column(col);
             let field = &self.fields[col];
-            let type_ = field.datatype();
-            let format = field.format();
-            encode_value(&mut encoder, array, self.curr_idx, type_, format).unwrap();
+
+            encode_value(&mut encoder, array, self.curr_idx, &field).unwrap();
         }
         self.curr_idx += 1;
         Some(encoder.finish())

arrow-pg/src/row_encoder.rs

@@ -6,7 +6,7 @@ use arrow::array::{Array, StructArray};
 use datafusion::arrow::array::{Array, StructArray};
 
 use bytes::{BufMut, BytesMut};
-use pgwire::api::results::FieldFormat;
+use pgwire::api::results::{FieldFormat, FieldInfo};
 use pgwire::error::PgWireResult;
 use pgwire::types::{ToSqlText, QUOTE_CHECK, QUOTE_ESCAPE};
 use postgres_types::{Field, IsNull, ToSql, Type};
@@ -17,7 +17,7 @@ pub(crate) fn encode_struct(
     arr: &Arc<dyn Array>,
     idx: usize,
     fields: &[Field],
-    format: FieldFormat,
+    parent_pg_field_info: &FieldInfo,
 ) -> PgWireResult<Option<EncodedValue>> {
     let arr = arr.as_any().downcast_ref::<StructArray>().unwrap();
     if arr.is_null(idx) {
@@ -27,7 +27,19 @@ pub(crate) fn encode_struct(
     for (i, arr) in arr.columns().iter().enumerate() {
         let field = &fields[i];
         let type_ = field.type_();
-        encode_value(&mut row_encoder, arr, idx, type_, format).unwrap();
+
+        let mut pg_field = FieldInfo::new(
+            field.name().to_string(),
+            None,
+            None,
+            type_.clone(),
+            parent_pg_field_info.format(),
+        );
+        if let Some(format_options) = parent_pg_field_info.format_options() {
+            pg_field = pg_field.with_format_options(format_options);
+        }
+
+        encode_value(&mut row_encoder, arr, idx, &pg_field).unwrap();
     }
     Ok(Some(EncodedValue {
         bytes: row_encoder.row_buffer,