feat: include optional FormatOptions in arrow encode (#229)

* feat: include optional FormatOptions in arrow encode

* chore: remove unused imports/exports

* test: update duckdb example

* feat: use FieldInfo in encode functions

* feat: simplify Encoder trait function

* fix: lint fix

* feat: include FormatOptions in encoder api

* chore: update pgwire to a released version

Author: sunng87

Cargo.lock

@@ -18,7 +18,7 @@ bytes = "1.10.1"
 chrono = { version = "0.4", features = ["std"] }
 datafusion = { version = "50", default-features = false }
 futures = "0.3"
-pgwire = { version = "0.34", default-features = false }
+pgwire = { version = "0.35", default-features = false }
 postgres-types = "0.2"
 rust_decimal = { version = "1.39", features = ["db-postgres"] }
 tokio = { version = "1", default-features = false }

Cargo.toml

@@ -19,6 +19,7 @@ use pgwire::api::stmt::{NoopQueryParser, StoredStatement};
 use pgwire::api::{ClientInfo, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::tokio::process_socket;
+use pgwire::types::format::FormatOptions;
 use tokio::net::TcpListener;
 
 pub struct DuckDBBackend {
@@ -45,7 +46,7 @@ impl AuthSource for DummyAuthSource {
 
 #[async_trait]
 impl SimpleQueryHandler for DuckDBBackend {
-    async fn do_query<C>(&self, _client: &mut C, query: &str) -> PgWireResult<Vec<Response>>
+    async fn do_query<C>(&self, client: &mut C, query: &str) -> PgWireResult<Vec<Response>>
     where
         C: ClientInfo + Unpin + Send + Sync,
     {
@@ -59,9 +60,12 @@ impl SimpleQueryHandler for DuckDBBackend {
                 .query_arrow(params![])
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
             let schema = ret.get_schema();
+            let format_options = FormatOptions::from_client_metadata(client.metadata());
+
             let header = Arc::new(arrow_schema_to_pg_fields(
                 schema.as_ref(),
                 &Format::UnifiedText,
+                Some(Arc::new(format_options)),
             )?);
 
             let header_ref = header.clone();
@@ -155,7 +159,7 @@ impl ExtendedQueryHandler for DuckDBBackend {
 
     async fn do_query<C>(
         &self,
-        _client: &mut C,
+        client: &mut C,
         portal: &Portal<Self::Statement>,
         _max_rows: usize,
     ) -> PgWireResult<Response>
@@ -178,9 +182,11 @@ impl ExtendedQueryHandler for DuckDBBackend {
                 .query_arrow(params![])
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
             let schema = ret.get_schema();
+            let format_options = FormatOptions::from_client_metadata(client.metadata());
             let header = Arc::new(arrow_schema_to_pg_fields(
                 schema.as_ref(),
                 &Format::UnifiedText,
+                Some(Arc::new(format_options)),
             )?);
 
             let header_ref = header.clone();

arrow-pg/examples/duckdb.rs

@@ -10,6 +10,7 @@ use pgwire::api::results::FieldInfo;
 use pgwire::api::Type;
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use pgwire::messages::data::DataRow;
+use pgwire::types::format::FormatOptions;
 use postgres_types::Kind;
 
 use crate::row_encoder::RowEncoder;
@@ -111,20 +112,25 @@ pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
     })
 }
 
-pub fn arrow_schema_to_pg_fields(schema: &Schema, format: &Format) -> PgWireResult<Vec<FieldInfo>> {
+pub fn arrow_schema_to_pg_fields(
+    schema: &Schema,
+    format: &Format,
+    data_format_options: Option<Arc<FormatOptions>>,
+) -> PgWireResult<Vec<FieldInfo>> {
+    let _ = data_format_options;
     schema
         .fields()
         .iter()
         .enumerate()
         .map(|(idx, f)| {
             let pg_type = into_pg_type(f.data_type())?;
-            Ok(FieldInfo::new(
-                f.name().into(),
-                None,
-                None,
-                pg_type,
-                format.format_for(idx),
-            ))
+            let mut field_info =
+                FieldInfo::new(f.name().into(), None, None, pg_type, format.format_for(idx));
+            if let Some(data_format_options) = &data_format_options {
+                field_info = field_info.with_format_options(data_format_options.clone());
+            }
+
+            Ok(field_info)
         })
         .collect::<PgWireResult<Vec<FieldInfo>>>()
 }

arrow-pg/src/datatypes.rs

@@ -13,13 +13,22 @@ use pgwire::api::results::QueryResponse;
 use pgwire::api::Type;
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::data::DataRow;
+use pgwire::types::format::FormatOptions;
 use rust_decimal::prelude::ToPrimitive;
 use rust_decimal::Decimal;
 
 use super::{arrow_schema_to_pg_fields, encode_recordbatch, into_pg_type};
 
-pub async fn encode_dataframe(df: DataFrame, format: &Format) -> PgWireResult<QueryResponse> {
-    let fields = Arc::new(arrow_schema_to_pg_fields(df.schema().as_arrow(), format)?);
+pub async fn encode_dataframe(
+    df: DataFrame,
+    format: &Format,
+    data_format_options: Option<Arc<FormatOptions>>,
+) -> PgWireResult<QueryResponse> {
+    let fields = Arc::new(arrow_schema_to_pg_fields(
+        df.schema().as_arrow(),
+        format,
+        data_format_options,
+    )?);
 
     let recordbatch_stream = df
         .execute_stream()