feat: arrow body format support settings geometry_output_format. (#716)

* ci: make string format of decimal same with http handler.

* feat: arrow body format support settings geometry_output_format.

* fix

* fix

* fix

Author: youngsofun

bindings/nodejs/tests/binding.js

@@ -216,6 +216,7 @@ Then("Select types should be expected native types", async function () {
       },
     ]);
   }
+
   // TimestampTz
   if (!(DRIVER_VERSION > [0, 30, 3] && DB_VERSION >= [1, 2, 836])) {
     const row = await this.conn.queryRow(`SELECT to_datetime_tz('2024-04-16 12:34:56.789 +0800'))`);

bindings/python/tests/blocking/steps/binding.py

@@ -185,6 +185,16 @@ def _(context):
                 f"timestamp_tz: {row.values()[0]} {exp_bug}"
             )
 
+    if DRIVER_VERSION > (0, 31, 0) and DB_VERSION > (1, 2, 841):
+        row = context.conn.query_row("SELECT st_point(60,37)")
+        assert row.values()[0] == '{"type": "Point", "coordinates": [60,37]}', (
+            f"geography: {row.values()}"
+        )
+        row = context.conn.query_row(
+            "settings(geometry_output_format='EWKT') SELECT st_point(60,37)"
+        )
+        assert row.values()[0] == "POINT(60 37)", f"geography: {row.values()}"
+
 
 @then("Select numbers should iterate all rows")
 def _(context):

core/src/lib.rs

@@ -30,6 +30,7 @@ mod session;
 mod stage;
 
 pub mod schema;
+mod settings;
 
 pub use auth::SensitiveString;
 pub use client::APIClient;
@@ -41,4 +42,6 @@ pub use presign::presign_upload_to_stage;
 pub use presign::PresignedResponse;
 pub use response::QueryStats;
 pub use response::SchemaField;
+pub use settings::GeometryDataType;
+pub use settings::ResultFormatSettings;
 pub use stage::StageLocation;

core/src/pages.rs

@@ -16,16 +16,15 @@ use crate::client::QueryState;
 use crate::error::Result;
 use crate::response::QueryResponse;
 use crate::schema::Schema;
+use crate::settings::ResultFormatSettings;
 use crate::{APIClient, Error, QueryStats, SchemaField};
 use arrow_array::RecordBatch;
-use chrono_tz::Tz;
 use log::debug;
 use parking_lot::Mutex;
 use std::collections::BTreeMap;
 use std::future::Future;
 use std::mem;
 use std::pin::Pin;
-use std::str::FromStr;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Instant;
@@ -110,7 +109,10 @@ impl Pages {
         self.first_page = Some(page);
     }
 
-    pub async fn wait_for_schema(mut self, need_progress: bool) -> Result<(Self, Schema, Tz)> {
+    pub async fn wait_for_schema(
+        mut self,
+        need_progress: bool,
+    ) -> Result<(Self, Schema, ResultFormatSettings)> {
         while let Some(page) = self.next().await {
             let page = page?;
             if !page.raw_schema.is_empty()
@@ -128,13 +130,8 @@ impl Pages {
                     s.try_into()
                         .map_err(|e| Error::Decode(format!("fail to decode string schema: {e}")))?
                 };
-                let utc = "UTC".to_owned();
-                let timezone = page
-                    .settings
-                    .as_ref()
-                    .and_then(|m| m.get("timezone"))
-                    .unwrap_or(&utc);
-                let timezone = Tz::from_str(timezone).map_err(|e| Error::Decode(e.to_string()))?;
+                let settings = ResultFormatSettings::from_map(&page.settings)?;
+
                 self.add_back(page);
                 let last_access_time = self.last_access_time.clone();
                 if let Some(node_id) = &self.node_id {
@@ -146,10 +143,10 @@ impl Pages {
                     self.client
                         .register_query_for_heartbeat(&self.query_id, state)
                 }
-                return Ok((self, schema, timezone));
+                return Ok((self, schema, settings));
             }
         }
-        Ok((self, Schema::default(), Tz::UTC))
+        Ok((self, Schema::default(), ResultFormatSettings::default()))
     }
 }
 

core/src/settings.rs

@@ -0,0 +1,77 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::error::Result;
+use crate::Error;
+use chrono_tz::Tz;
+use serde::Deserialize;
+use std::collections::BTreeMap;
+use std::str::FromStr;
+
+#[derive(Debug, Clone, Default, Copy)]
+pub struct ResultFormatSettings {
+    pub geometry_output_format: GeometryDataType,
+    pub timezone: Tz,
+}
+
+impl ResultFormatSettings {
+    pub fn from_map(settings: &Option<BTreeMap<String, String>>) -> Result<Self> {
+        match settings {
+            None => Ok(Default::default()),
+            Some(settings) => {
+                let timezone = match settings.get("timezone") {
+                    None => Tz::default(),
+                    Some(t) => Tz::from_str(t).map_err(|e| Error::Decode(e.to_string()))?,
+                };
+
+                let geometry_output_format = match settings.get("geometry_output_format") {
+                    None => GeometryDataType::default(),
+                    Some(t) => {
+                        GeometryDataType::from_str(t).map_err(|e| Error::Decode(e.to_string()))?
+                    }
+                };
+
+                Ok(Self {
+                    timezone,
+                    geometry_output_format,
+                })
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, Default, Deserialize)]
+pub enum GeometryDataType {
+    WKB,
+    WKT,
+    EWKB,
+    EWKT,
+    #[default]
+    GEOJSON,
+}
+
+impl FromStr for GeometryDataType {
+    type Err = Error;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s.to_uppercase().as_str() {
+            "WKB" => Ok(GeometryDataType::WKB),
+            "WKT" => Ok(GeometryDataType::WKT),
+            "EWKB" => Ok(GeometryDataType::EWKB),
+            "EWKT" => Ok(GeometryDataType::EWKT),
+            "GEOJSON" => Ok(GeometryDataType::GEOJSON),
+            _ => Err(Error::Decode("Invalid geometry type format".to_string())),
+        }
+    }
+}

driver/src/flight_sql.rs

@@ -25,7 +25,6 @@ use arrow_flight::sql::client::FlightSqlServiceClient;
 use arrow_flight::utils::flight_data_to_arrow_batch;
 use arrow_schema::SchemaRef as ArrowSchemaRef;
 use async_trait::async_trait;
-use chrono_tz::Tz;
 use percent_encoding::percent_decode_str;
 use tokio::sync::Mutex;
 use tokio_stream::{Stream, StreamExt};
@@ -34,9 +33,9 @@ use url::Url;
 
 use crate::client::LoadMethod;
 use crate::conn::{ConnectionInfo, IConnection, Reader};
-use databend_client::presign_upload_to_stage;
 use databend_client::schema::Schema;
 use databend_client::SensitiveString;
+use databend_client::{presign_upload_to_stage, ResultFormatSettings};
 use databend_driver_core::error::{Error, Result};
 use databend_driver_core::rows::{
     Row, RowIterator, RowStatsIterator, RowWithStats, Rows, ServerStats,
@@ -373,7 +372,7 @@ impl Stream for FlightSQLRows {
                         self.schema.clone(),
                         &dicitionaries_by_id,
                     )?;
-                    let rows = Rows::try_from((batch, Tz::UTC))?;
+                    let rows = Rows::try_from((batch, ResultFormatSettings::default()))?;
                     self.rows.extend(rows);
                     self.poll_next(cx)
                 }

driver/src/rest_api.rs

@@ -29,8 +29,8 @@ use tokio_stream::Stream;
 use crate::client::LoadMethod;
 use crate::conn::{ConnectionInfo, IConnection, Reader};
 use databend_client::schema::{Schema, SchemaRef};
-use databend_client::APIClient;
 use databend_client::Pages;
+use databend_client::{APIClient, ResultFormatSettings};
 use databend_driver_core::error::{Error, Result};
 use databend_driver_core::raw_rows::{RawRow, RawRowIterator, RawRowWithStats};
 use databend_driver_core::rows::{
@@ -296,7 +296,7 @@ pub struct RestAPIRows<T> {
     pages: Pages,
 
     schema: SchemaRef,
-    timezone: Tz,
+    settings: ResultFormatSettings,
 
     data: VecDeque<Vec<Option<String>>>,
     rows: VecDeque<Row>,
@@ -308,11 +308,11 @@ pub struct RestAPIRows<T> {
 
 impl<T> RestAPIRows<T> {
     async fn from_pages(pages: Pages) -> Result<(Schema, Self)> {
-        let (pages, schema, timezone) = pages.wait_for_schema(true).await?;
+        let (pages, schema, settings) = pages.wait_for_schema(true).await?;
         let rows = Self {
             pages,
             schema: Arc::new(schema.clone()),
-            timezone,
+            settings,
             data: Default::default(),
             rows: Default::default(),
             stats: None,
@@ -333,7 +333,7 @@ impl<T: FromRowStats + std::marker::Unpin> Stream for RestAPIRows<T> {
         // Therefore, we could guarantee the `/final` called before the last row.
         if self.data.len() > 1 {
             if let Some(row) = self.data.pop_front() {
-                let row = T::try_from_raw_row(row, self.schema.clone(), self.timezone)?;
+                let row = T::try_from_raw_row(row, self.schema.clone(), self.settings.timezone)?;
                 return Poll::Ready(Some(Ok(row)));
             }
         } else if self.rows.len() > 1 {
@@ -357,7 +357,7 @@ impl<T: FromRowStats + std::marker::Unpin> Stream for RestAPIRows<T> {
                     self.data.append(&mut new_data);
                 } else {
                     for batch in page.batches.into_iter() {
-                        let rows = Rows::try_from((batch, self.timezone))?;
+                        let rows = Rows::try_from((batch, self.settings))?;
                         self.rows.extend(rows);
                     }
                 }
@@ -369,7 +369,8 @@ impl<T: FromRowStats + std::marker::Unpin> Stream for RestAPIRows<T> {
                     let row = T::from_row(row);
                     Poll::Ready(Some(Ok(row)))
                 } else if let Some(row) = self.data.pop_front() {
-                    let row = T::try_from_raw_row(row, self.schema.clone(), self.timezone)?;
+                    let row =
+                        T::try_from_raw_row(row, self.schema.clone(), self.settings.timezone)?;
                     Poll::Ready(Some(Ok(row)))
                 } else {
                     Poll::Ready(None)

sql/Cargo.toml

@@ -16,16 +16,17 @@ flight-sql = ["dep:tonic"]
 [dependencies]
 arrow = { workspace = true }
 arrow-array = { workspace = true }
-arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 chrono = { workspace = true }
 chrono-tz = { workspace = true }
+ethnum = "1.5.1"
 databend-client = { workspace = true }
 jsonb = { workspace = true }
 tokio-stream = { workspace = true }
 tonic = { workspace = true, optional = true }
 
-geozero = { version = "0.14.0", features = ["with-wkb"] }
+geo = { version = "0.28.0", features = ["use-serde"] }
+geozero = { version = "0.14.0", features = ["with-geo", "with-geojson", "with-wkb", "with-wkt"] }
 glob = "0.3"
 hex = "0.4.3"
 itertools = "0.14"

sql/src/rows.rs

@@ -24,6 +24,7 @@ use crate::value::Value;
 use arrow::record_batch::RecordBatch;
 use chrono_tz::Tz;
 use databend_client::schema::SchemaRef;
+use databend_client::ResultFormatSettings;
 
 #[derive(Clone, Debug)]
 pub enum RowWithStats {
@@ -182,9 +183,9 @@ impl Rows {
     }
 }
 
-impl TryFrom<(RecordBatch, Tz)> for Rows {
+impl TryFrom<(RecordBatch, ResultFormatSettings)> for Rows {
     type Error = Error;
-    fn try_from((batch, ltz): (RecordBatch, Tz)) -> Result<Self> {
+    fn try_from((batch, settings): (RecordBatch, ResultFormatSettings)) -> Result<Self> {
         let batch_schema = batch.schema();
         let schema = SchemaRef::new(batch_schema.clone().try_into()?);
         let mut rows: Vec<Row> = Vec::new();
@@ -193,7 +194,7 @@ impl TryFrom<(RecordBatch, Tz)> for Rows {
             for j in 0..batch_schema.fields().len() {
                 let v = batch.column(j);
                 let field = batch_schema.field(j);
-                let value = Value::try_from((field, v, i, ltz))?;
+                let value = Value::try_from((field, v, i, settings))?;
                 values.push(value);
             }
             rows.push(Row::new(schema.clone(), values));