feat: include optional FormatOptions in arrow encode

Author: sunng87

Cargo.lock

@@ -19,7 +19,8 @@ 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.34", default-features = false }
+pgwire = { git = "https://github.com/sunng87/pgwire", rev = "d89089947a56ebfe2e89631925facdd7a85c25b4", default-features = false }
 postgres-types = "0.2"
 rust_decimal = { version = "1.39", features = ["db-postgres"] }
 tokio = { version = "1", default-features = false }

Cargo.toml

@@ -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()

arrow-pg/src/datatypes/df.rs

@@ -6,6 +6,12 @@
 pub mod datatypes;
 pub mod encoder;
 mod error;
+pub mod format;
 pub mod list_encoder;
 pub mod row_encoder;
 pub mod struct_encoder;
+
+#[cfg(feature = "datafusion")]
+pub use datatypes::df::encode_dataframe;
+
+pub use datatypes::encode_recordbatch;

arrow-pg/src/lib.rs

@@ -22,6 +22,7 @@ use pgwire::api::stmt::StoredStatement;
 use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::response::TransactionStatus;
+use pgwire::types::format::FormatOptions;
 
 use crate::auth::AuthManager;
 use crate::client;
@@ -355,7 +356,10 @@ impl SimpleQueryHandler for DfSessionService {
                 results.push(resp);
             } else {
                 // For non-INSERT queries, return a regular Query response
-                let resp = df::encode_dataframe(df, &Format::UnifiedText).await?;
+                let format_options =
+                    Arc::new(FormatOptions::from_client_metadata(client.metadata()));
+                let resp =
+                    df::encode_dataframe(df, &Format::UnifiedText, Some(format_options)).await?;
                 results.push(Response::Query(resp));
             }
         }
@@ -382,7 +386,8 @@ impl ExtendedQueryHandler for DfSessionService {
     {
         if let (_, Some((_, plan))) = &target.statement {
             let schema = plan.schema();
-            let fields = arrow_schema_to_pg_fields(schema.as_arrow(), &Format::UnifiedBinary)?;
+            let fields =
+                arrow_schema_to_pg_fields(schema.as_arrow(), &Format::UnifiedBinary, None)?;
             let params = plan
                 .get_parameter_types()
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
@@ -415,7 +420,7 @@ impl ExtendedQueryHandler for DfSessionService {
         if let (_, Some((_, plan))) = &target.statement.statement {
             let format = &target.result_column_format;
             let schema = plan.schema();
-            let fields = arrow_schema_to_pg_fields(schema.as_arrow(), format)?;
+            let fields = arrow_schema_to_pg_fields(schema.as_arrow(), format, None)?;
 
             Ok(DescribePortalResponse::new(fields))
         } else {
@@ -543,7 +548,14 @@ impl ExtendedQueryHandler for DfSessionService {
                 Ok(resp)
             } else {
                 // For non-INSERT queries, return a regular Query response
-                let resp = df::encode_dataframe(dataframe, &portal.result_column_format).await?;
+                let format_options =
+                    Arc::new(FormatOptions::from_client_metadata(client.metadata()));
+                let resp = df::encode_dataframe(
+                    dataframe,
+                    &portal.result_column_format,
+                    Some(format_options),
+                )
+                .await?;
                 Ok(Response::Query(resp))
             }
         } else {