feat: support arrow_result_version. (#769)

youngsofun · web-flow · commit 5949cee13879 · 2026-04-22T00:13:33.000+08:00
diff --git a/.github/workflows/bindings.python.yml b/.github/workflows/bindings.python.yml
@@ -138,6 +138,9 @@ jobs:
       - name: Prepare
         working-directory: tests
         run: make up
+        env:
+          DATABEND_META_VERSION: v1.2.898-nightly
+          DATABEND_QUERY_VERSION: v1.2.898-nightly
       - name: Download artifact
         uses: actions/download-artifact@v4
         with:
diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -360,11 +360,11 @@ pub async fn main() -> Result<()> {
                         println!("Authenticate failed wrong password user {user}");
                         return Ok(());
                     }
-                    databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(ipc_err)) => {
-                        if ipc_err.contains("Unauthenticated") {
-                            println!("Authenticate failed wrong password user {user}");
-                            return Ok(());
-                        }
+                    databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(ipc_err))
+                        if ipc_err.contains("Unauthenticated") =>
+                    {
+                        println!("Authenticate failed wrong password user {user}");
+                        return Ok(());
                     }
                     _ => {}
                 }
diff --git a/cli/src/session.rs b/cli/src/session.rs
@@ -107,19 +107,15 @@ impl Session {
                         }
                         databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(
                             ref ipc_err,
-                        )) => {
-                            if ipc_err.contains("Unauthenticated")
-                                || ipc_err.contains("Connection refused")
-                            {
-                                return Err(err.into());
-                            }
+                        )) if (ipc_err.contains("Unauthenticated")
+                            || ipc_err.contains("Connection refused")) =>
+                        {
+                            return Err(err.into());
                         }
                         databend_driver::Error::Api(databend_client::Error::Request(
                             ref resp_err,
-                        )) => {
-                            if resp_err.contains("error sending request for url") {
-                                return Err(err.into());
-                            }
+                        )) if resp_err.contains("error sending request for url") => {
+                            return Err(err.into());
                         }
                         _ => {}
                     }
diff --git a/core/src/client.rs b/core/src/client.rs
@@ -598,7 +598,7 @@ impl APIClient {
         // body
         let session_state = self.session_state();
         let need_sticky = session_state.need_sticky.unwrap_or(false);
-        let req = QueryRequest::new(sql)
+        let mut req = QueryRequest::new(sql)
             .with_pagination(self.make_pagination())
             .with_session(Some(session_state))
             .with_stage_attachment(stage_attachment_config);
@@ -609,6 +609,7 @@ impl APIClient {
         if self.capability.arrow_data && self.query_result_format == "arrow" && !force_json_body {
             debug!("accept arrow data");
             headers.insert(ACCEPT, HeaderValue::from_static(CONTENT_TYPE_ARROW_OR_JSON));
+            req = req.with_arrow();
         }
 
         if need_sticky {
diff --git a/core/src/lib.rs b/core/src/lib.rs
@@ -44,5 +44,6 @@ pub use presign::PresignedResponse;
 pub use response::QueryStats;
 pub use response::SchemaField;
 pub use settings::GeometryDataType;
+pub use settings::QueryResultFormatSettings;
 pub use settings::ResultFormatSettings;
 pub use stage::StageLocation;
diff --git a/core/src/pages.rs b/core/src/pages.rs
@@ -16,12 +16,11 @@ use crate::client::QueryState;
 use crate::error::Result;
 use crate::response::QueryResponse;
 use crate::schema::Schema;
-use crate::settings::ResultFormatSettings;
+use crate::settings::{QueryResultFormatSettings, ResultFormatSettings};
 use crate::{APIClient, Error, QueryStats, SchemaField};
 use arrow_array::RecordBatch;
 use log::debug;
 use parking_lot::Mutex;
-use std::collections::BTreeMap;
 use std::future::Future;
 use std::mem;
 use std::pin::Pin;
@@ -36,7 +35,7 @@ pub struct Page {
     pub data: Vec<Vec<Option<String>>>,
     pub batches: Vec<RecordBatch>,
     pub stats: QueryStats,
-    pub settings: Option<BTreeMap<String, String>>,
+    pub settings: Option<QueryResultFormatSettings>,
 }
 
 impl Page {
@@ -51,6 +50,9 @@ impl Page {
     }
 
     pub fn update(&mut self, p: Page) {
+        if self.settings.is_none() {
+            self.settings = p.settings.clone();
+        }
         self.raw_schema = p.raw_schema;
         if self.data.is_empty() {
             self.data = p.data
@@ -130,7 +132,7 @@ impl Pages {
                     s.try_into()
                         .map_err(|e| Error::Decode(format!("fail to decode string schema: {e}")))?
                 };
-                let settings = ResultFormatSettings::from_map(&page.settings)?;
+                let settings = ResultFormatSettings::try_from(&page.settings)?;
 
                 self.add_back(page);
                 let last_access_time = self.last_access_time.clone();
diff --git a/core/src/request.rs b/core/src/request.rs
@@ -26,6 +26,8 @@ pub struct QueryRequest<'a> {
     pagination: Option<PaginationConfig>,
     #[serde(skip_serializing_if = "Option::is_none")]
     stage_attachment: Option<StageAttachmentConfig<'a>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    arrow_result_version_max: Option<i64>,
 }
 
 #[derive(Serialize, Debug)]
@@ -54,6 +56,7 @@ impl<'r, 't: 'r> QueryRequest<'r> {
             sql,
             pagination: None,
             stage_attachment: None,
+            arrow_result_version_max: None,
         }
     }
 
@@ -74,6 +77,11 @@ impl<'r, 't: 'r> QueryRequest<'r> {
         self.stage_attachment = stage_attachment;
         self
     }
+
+    pub fn with_arrow(mut self) -> Self {
+        self.arrow_result_version_max = Some(2);
+        self
+    }
 }
 
 #[cfg(test)]
diff --git a/core/src/response.rs b/core/src/response.rs
@@ -14,8 +14,8 @@
 
 use crate::error_code::ErrorCode;
 use crate::session::SessionState;
+use crate::settings::QueryResultFormatSettings;
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
 
 #[derive(Deserialize, Debug, Default)]
 pub struct QueryStats {
@@ -78,7 +78,7 @@ pub struct QueryResponse {
     pub schema: Vec<SchemaField>,
     pub data: Vec<Vec<Option<String>>>,
     pub state: String,
-    pub settings: Option<BTreeMap<String, String>>,
+    pub settings: Option<QueryResultFormatSettings>,
     pub error: Option<ErrorCode>,
     // make it optional for backward compatibility
     pub warnings: Option<Vec<String>>,
diff --git a/core/src/schema.rs b/core/src/schema.rs
@@ -471,21 +471,17 @@ fn parse_type_desc(s: &str) -> Result<TypeDesc<'_>> {
                     start = i + 1;
                 }
             }
-            ',' => {
-                if depth == 1 {
-                    let s = &s[start..i];
-                    args.push(parse_type_desc(s)?);
-                    start = i + 1;
-                }
+            ',' if depth == 1 => {
+                let s = &s[start..i];
+                args.push(parse_type_desc(s)?);
+                start = i + 1;
             }
-            ' ' => {
-                if depth == 0 {
-                    let s = &s[start..i];
-                    if !s.is_empty() {
-                        name = s;
-                    }
-                    start = i + 1;
+            ' ' if depth == 0 => {
+                let s = &s[start..i];
+                if !s.is_empty() {
+                    name = s;
                 }
+                start = i + 1;
             }
             _ => {}
         }
diff --git a/core/src/settings.rs b/core/src/settings.rs
@@ -16,51 +16,65 @@ use crate::error::Result;
 use crate::Error;
 use jiff::tz::TimeZone;
 use serde::Deserialize;
-use std::collections::BTreeMap;
 use std::str::FromStr;
 
 #[derive(Debug, Clone)]
 pub struct ResultFormatSettings {
     pub geometry_output_format: GeometryDataType,
     pub timezone: TimeZone,
+    pub arrow_result_version: Option<i64>,
+    pub binary_output_format: BinaryFormat,
 }
 
 impl Default for ResultFormatSettings {
     fn default() -> Self {
         Self {
             geometry_output_format: GeometryDataType::default(),
+            binary_output_format: BinaryFormat::default(),
             timezone: TimeZone::UTC,
+            arrow_result_version: None,
         }
     }
 }
 
-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 => TimeZone::UTC,
-                    Some(t) => TimeZone::get(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,
-                })
-            }
-        }
+impl TryFrom<&Option<QueryResultFormatSettings>> for ResultFormatSettings {
+    type Error = Error;
+
+    fn try_from(settings: &Option<QueryResultFormatSettings>) -> Result<Self> {
+        let settings = settings.clone().unwrap_or_default();
+        let timezone = match settings.timezone {
+            None => TimeZone::UTC,
+            Some(t) => TimeZone::get(&t).map_err(|e| Error::Decode(e.to_string()))?,
+        };
+
+        let geometry_output_format = match settings.geometry_output_format {
+            None => GeometryDataType::default(),
+            Some(t) => GeometryDataType::from_str(&t).map_err(|e| Error::Decode(e.to_string()))?,
+        };
+
+        let binary_output_format = match settings.binary_output_format {
+            None => BinaryFormat::default(),
+            Some(t) => BinaryFormat::from_str(&t).map_err(|e| Error::Decode(e.to_string()))?,
+        };
+
+        Ok(Self {
+            geometry_output_format,
+            timezone,
+            arrow_result_version: settings.arrow_result_version,
+            binary_output_format,
+        })
     }
 }
 
-#[derive(Debug, Clone, Copy, Default, Deserialize)]
+#[derive(Debug, Clone, Default, Deserialize)]
+pub struct QueryResultFormatSettings {
+    pub timezone: Option<String>,
+    pub geometry_output_format: Option<String>,
+    pub arrow_result_version: Option<i64>,
+    pub binary_output_format: Option<String>,
+}
+
+#[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)]
 pub enum GeometryDataType {
     WKB,
     WKT,
@@ -84,3 +98,72 @@ impl FromStr for GeometryDataType {
         }
     }
 }
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum BinaryFormat {
+    #[default]
+    Hex,
+    Base64,
+    Utf8,
+    Utf8Lossy,
+}
+
+impl FromStr for BinaryFormat {
+    type Err = Error;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s.to_ascii_lowercase().as_str() {
+            "hex" => Ok(BinaryFormat::Hex),
+            "base64" => Ok(BinaryFormat::Base64),
+            "utf-8" | "utf8" => Ok(BinaryFormat::Utf8),
+            "utf-8-lossy" | "utf8-lossy" => Ok(BinaryFormat::Utf8Lossy),
+            other => Err(Error::Decode(format!(
+                "Invalid binary format '{other}', valid values: HEX | BASE64 | UTF-8 | UTF-8-LOSSY"
+            ))),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn deserialize_query_result_format_settings_from_strings() {
+        let settings: QueryResultFormatSettings = serde_json::from_str(
+            r#"{
+                "timezone": "Asia/Shanghai",
+                "geometry_output_format": "wkt",
+                "arrow_result_version": 2,
+                "binary_output_format": "utf-8"
+            }"#,
+        )
+        .unwrap();
+
+        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
+        assert_eq!(settings.geometry_output_format, GeometryDataType::WKT);
+        assert_eq!(settings.arrow_result_version, Some(2));
+        assert_eq!(settings.binary_output_format, BinaryFormat::Utf8);
+        assert_eq!(settings.timezone.iana_name(), Some("Asia/Shanghai"));
+    }
+
+    #[test]
+    fn deserialize_query_result_format_settings_with_defaults() {
+        let settings: QueryResultFormatSettings = serde_json::from_str(r#"{}"#).unwrap();
+
+        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
+        assert_eq!(settings.geometry_output_format, GeometryDataType::default());
+        assert_eq!(settings.arrow_result_version, None);
+        assert_eq!(settings.binary_output_format, BinaryFormat::default());
+        assert_eq!(settings.timezone.iana_name(), Some("UTC"));
+    }
+
+    #[test]
+    fn deserialize_query_result_format_settings_accepts_numeric_arrow_version() {
+        let settings: QueryResultFormatSettings =
+            serde_json::from_str(r#"{"arrow_result_version": 2}"#).unwrap();
+
+        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
+        assert_eq!(settings.arrow_result_version, Some(2));
+    }
+}
diff --git a/driver/src/placeholder.rs b/driver/src/placeholder.rs
@@ -78,7 +78,7 @@ impl PlaceholderVisitor {
 
     pub fn replace_sql(&mut self, params: &Params, stmt: &Statement, sql: &str) -> String {
         stmt.drive(self);
-        self.placeholders.sort_by(|l, r| l.start.cmp(&r.start));
+        self.placeholders.sort_by_key(|l| l.start);
 
         let mut results = vec![];
 
@@ -96,7 +96,7 @@ impl PlaceholderVisitor {
 
         let mut sql = sql.to_string();
         if !results.is_empty() {
-            results.sort_by(|a, b| a.1.start.cmp(&b.1.start));
+            results.sort_by_key(|a| a.1.start);
             for (value, r) in results.iter().rev() {
                 let start = r.start as usize;
                 let end = r.end as usize;
@@ -105,8 +105,7 @@ impl PlaceholderVisitor {
         }
 
         if !self.column_positions.is_empty() {
-            self.column_positions
-                .sort_by(|a, b| a.1.start.cmp(&b.1.start));
+            self.column_positions.sort_by_key(|a| a.1.start);
 
             for (index, r) in self.column_positions.iter().rev() {
                 if let Some(value) = params.get_by_index(*index) {
diff --git a/sql/src/raw_rows.rs b/sql/src/raw_rows.rs
@@ -76,7 +76,7 @@ impl TryFrom<(SchemaRef, Vec<Option<String>>, &ResultFormatSettings)> for RawRow
         (schema, data, settings): (SchemaRef, Vec<Option<String>>, &ResultFormatSettings),
     ) -> Result<Self> {
         let mut values: Vec<Value> = Vec::with_capacity(data.len());
-        for (field, val) in schema.fields().iter().zip(data.clone().into_iter()) {
+        for (field, val) in schema.fields().iter().zip(data.clone()) {
             values.push(Value::try_from((&field.data_type, val, settings))?);
         }
 
diff --git a/sql/src/rows.rs b/sql/src/rows.rs
diff --git a/sql/src/value/arrow_decoder.rs b/sql/src/value/arrow_decoder.rs

Original file line number	Diff line number	Diff line change
`@@ -360,11 +360,11 @@ pub async fn main() -> Result<()> {`
`360`	`360`	`println!("Authenticate failed wrong password user {user}");`
`361`	`361`	`return Ok(());`
`362`	`362`	`}`
`363`		`- databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(ipc_err)) => {`
`364`		`- if ipc_err.contains("Unauthenticated") {`
`365`		`- println!("Authenticate failed wrong password user {user}");`
`366`		`- return Ok(());`
`367`		`- }`
	`363`	`+ databend_driver::Error::Arrow(arrow::error::ArrowError::IpcError(ipc_err))`
	`364`	`+ if ipc_err.contains("Unauthenticated") =>`
	`365`	`+ {`
	`366`	`+ println!("Authenticate failed wrong password user {user}");`
	`367`	`+ return Ok(());`
`368`	`368`	`}`
`369`	`369`	`_ => {}`
`370`	`370`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,8 @@ pub struct QueryRequest<'a> {`
`26`	`26`	`pagination: Option<PaginationConfig>,`
`27`	`27`	`#[serde(skip_serializing_if = "Option::is_none")]`
`28`	`28`	`stage_attachment: Option<StageAttachmentConfig<'a>>,`
	`29`	`+ #[serde(skip_serializing_if = "Option::is_none")]`
	`30`	`+ arrow_result_version_max: Option<i64>,`
`29`	`31`	`}`
`30`	`32`
`31`	`33`	`#[derive(Serialize, Debug)]`
`@@ -54,6 +56,7 @@ impl<'r, 't: 'r> QueryRequest<'r> {`
`54`	`56`	`sql,`
`55`	`57`	`pagination: None,`
`56`	`58`	`stage_attachment: None,`
	`59`	`+ arrow_result_version_max: None,`
`57`	`60`	`}`
`58`	`61`	`}`
`59`	`62`
`@@ -74,6 +77,11 @@ impl<'r, 't: 'r> QueryRequest<'r> {`
`74`	`77`	`self.stage_attachment = stage_attachment;`
`75`	`78`	`self`
`76`	`79`	`}`
	`80`	`+`
	`81`	`+ pub fn with_arrow(mut self) -> Self {`
	`82`	`+ self.arrow_result_version_max = Some(2);`
	`83`	`+ self`
	`84`	`+ }`
`77`	`85`	`}`
`78`	`86`
`79`	`87`	`#[cfg(test)]`
Original file line number	Diff line number	Diff line change
`@@ -471,21 +471,17 @@ fn parse_type_desc(s: &str) -> Result<TypeDesc<'_>> {`
`471`	`471`	`start = i + 1;`
`472`	`472`	`}`
`473`	`473`	`}`
`474`		`- ',' => {`
`475`		`- if depth == 1 {`
`476`		`- let s = &s[start..i];`
`477`		`- args.push(parse_type_desc(s)?);`
`478`		`- start = i + 1;`
`479`		`- }`
	`474`	`+ ',' if depth == 1 => {`
	`475`	`+ let s = &s[start..i];`
	`476`	`+ args.push(parse_type_desc(s)?);`
	`477`	`+ start = i + 1;`
`480`	`478`	`}`
`481`		`- ' ' => {`
`482`		`- if depth == 0 {`
`483`		`- let s = &s[start..i];`
`484`		`- if !s.is_empty() {`
`485`		`- name = s;`
`486`		`- }`
`487`		`- start = i + 1;`
	`479`	`+ ' ' if depth == 0 => {`
	`480`	`+ let s = &s[start..i];`
	`481`	`+ if !s.is_empty() {`
	`482`	`+ name = s;`
`488`	`483`	`}`
	`484`	`+ start = i + 1;`
`489`	`485`	`}`
`490`	`486`	`_ => {}`
`491`	`487`	`}`