handlers: fix response encoding format

src/datatypes.rs

@@ -8,8 +8,8 @@ use datafusion::common::{DFSchema, ParamValues};
 use datafusion::prelude::*;
 use datafusion::scalar::ScalarValue;
 use futures::{stream, StreamExt};
-use pgwire::api::portal::Portal;
-use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse};
+use pgwire::api::portal::{Format, Portal};
+use pgwire::api::results::{DataRowEncoder, FieldInfo, QueryResponse};
 use pgwire::api::Type;
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 
@@ -316,25 +316,32 @@ fn encode_value(
     Ok(())
 }
 
-pub(crate) fn df_schema_to_pg_fields(schema: &DFSchema) -> PgWireResult<Vec<FieldInfo>> {
+pub(crate) fn df_schema_to_pg_fields(
+    schema: &DFSchema,
+    format: &Format,
+) -> PgWireResult<Vec<FieldInfo>> {
     schema
         .fields()
         .iter()
-        .map(|f| {
+        .enumerate()
+        .map(|(idx, f)| {
             let pg_type = into_pg_type(f.data_type())?;
             Ok(FieldInfo::new(
                 f.name().into(),
                 None,
                 None,
                 pg_type,
-                FieldFormat::Text,
+                format.format_for(idx),
             ))
         })
         .collect::<PgWireResult<Vec<FieldInfo>>>()
 }
 
-pub(crate) async fn encode_dataframe<'a>(df: DataFrame) -> PgWireResult<QueryResponse<'a>> {
-    let fields = Arc::new(df_schema_to_pg_fields(df.schema())?);
+pub(crate) async fn encode_dataframe<'a>(
+    df: DataFrame,
+    format: &Format,
+) -> PgWireResult<QueryResponse<'a>> {
+    let fields = Arc::new(df_schema_to_pg_fields(df.schema(), format)?);
 
     let recordbatch_stream = df
         .execute_stream()

src/handlers.rs

@@ -4,7 +4,7 @@ use async_trait::async_trait;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::prelude::*;
-use pgwire::api::portal::Portal;
+use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DescribePortalResponse, DescribeStatementResponse, Response};
 use pgwire::api::stmt::QueryParser;
@@ -50,7 +50,7 @@ impl SimpleQueryHandler for DfSessionService {
                 .await
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
 
-            let resp = datatypes::encode_dataframe(df).await?;
+            let resp = datatypes::encode_dataframe(df, &Format::UnifiedText).await?;
             Ok(vec![Response::Query(resp)])
         } else {
             Ok(vec![Response::Error(Box::new(ErrorInfo::new(
@@ -107,7 +107,7 @@ impl ExtendedQueryHandler for DfSessionService {
         let plan = &target.statement;
 
         let schema = plan.schema();
-        let fields = datatypes::df_schema_to_pg_fields(schema.as_ref())?;
+        let fields = datatypes::df_schema_to_pg_fields(schema.as_ref(), &Format::UnifiedBinary)?;
         let params = plan
             .get_parameter_types()
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
@@ -135,8 +135,9 @@ impl ExtendedQueryHandler for DfSessionService {
         C: ClientInfo + Unpin + Send + Sync,
     {
         let plan = &target.statement.statement;
+        let format = &target.result_column_format;
         let schema = plan.schema();
-        let fields = datatypes::df_schema_to_pg_fields(schema.as_ref())?;
+        let fields = datatypes::df_schema_to_pg_fields(schema.as_ref(), format)?;
 
         Ok(DescribePortalResponse::new(fields))
     }
@@ -175,7 +176,7 @@ impl ExtendedQueryHandler for DfSessionService {
             .await
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
 
-        let resp = datatypes::encode_dataframe(dataframe).await?;
+        let resp = datatypes::encode_dataframe(dataframe, &portal.result_column_format).await?;
         Ok(Response::Query(resp))
     }
 }