feat: negotiate HTTP arrow features and stringify variant schema (#19780)

Author: youngsofun

src/common/io/src/format_settings.rs

@@ -148,6 +148,8 @@ pub struct OutputFormatSettings {
     pub headers: u8,
     // used only for compat with old bendSQL driver.
     pub http_arrow_use_jsonb: bool,
+    // used only in http arrow response to gate decimal64 physical type output.
+    pub http_arrow_use_decimal64: bool,
 
     pub json_compact: bool,
     pub json_strings: bool,
@@ -165,6 +167,7 @@ impl Default for OutputFormatSettings {
             http_json_result_mode: HttpHandlerDataFormat::Display,
             headers: 0,
             http_arrow_use_jsonb: false,
+            http_arrow_use_decimal64: true,
             json_compact: false,
             json_strings: false,
             format_null_as_str: false,

src/query/service/Cargo.toml

@@ -32,6 +32,7 @@ storage-hdfs = ["opendal/services-hdfs", "databend-common-storage/storage-hdfs"]
 anyhow = { workspace = true }
 arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
+arrow-cast = { workspace = true }
 arrow-flight = { workspace = true }
 arrow-ipc = { workspace = true }
 arrow-json = { workspace = true }
@@ -185,7 +186,6 @@ uuid = { workspace = true }
 walkdir = { workspace = true }
 
 [dev-dependencies]
-arrow-cast = { workspace = true }
 databend-common-sql-test-support = { workspace = true }
 databend-storages-common-pruner = { workspace = true }
 geo = { workspace = true }

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -68,6 +68,7 @@ use super::query::HttpQuery;
 use super::query::HttpQueryRequest;
 use super::query::HttpQueryResponseInternal;
 use super::query::ResponseState;
+use super::query::execute_state::ARROW_FEATURE_NEGOTIATION_VERSION;
 use super::query::execute_state::LEGACY_ARROW_RESULT_VERSION;
 use super::query::execute_state::SERVER_MAX_ARROW_RESULT_VERSION;
 use super::query::execute_state::legacy_bendsql_python_arrow_result_version;
@@ -156,6 +157,28 @@ pub struct QueryResponseSettings {
     pub http_json_result_mode: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub arrow_result_version: Option<u64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub arrow_features: Option<ArrowFeatures>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq)]
+pub struct ArrowFeatures {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub decimal64: Option<bool>,
+}
+
+impl ArrowFeatures {
+    pub(crate) fn decimal64_enabled() -> Self {
+        Self {
+            decimal64: Some(true),
+        }
+    }
+
+    fn visible_enabled_only(&self) -> Option<Self> {
+        self.decimal64
+            .filter(|enabled| *enabled)
+            .map(|_| Self::decimal64_enabled())
+    }
 }
 
 #[derive(Serialize, Debug, Clone)]
@@ -211,6 +234,7 @@ impl QueryResponse {
                     warnings,
                     response_settings,
                     arrow_result_version: _,
+                    arrow_features: _,
                 },
         }: HttpQueryResponseInternal,
         is_final: bool,
@@ -263,7 +287,8 @@ impl QueryResponse {
                 || encoding.arrow_result_version.is_some()
         );
         let response_settings = response_settings.map(|mut settings| {
-            settings.arrow_result_version = visible_arrow_result_version(encoding);
+            settings.arrow_result_version = visible_arrow_result_version(&encoding);
+            settings.arrow_features = visible_arrow_features(&encoding);
             settings
         });
 
@@ -324,7 +349,7 @@ impl QueryResponse {
     }
 }
 
-fn visible_arrow_result_version(encoding: ResponseEncoding) -> Option<u64> {
+fn visible_arrow_result_version(encoding: &ResponseEncoding) -> Option<u64> {
     match encoding.format {
         QueryResultFormat::Json => None,
         QueryResultFormat::Arrow => encoding
@@ -333,6 +358,22 @@ fn visible_arrow_result_version(encoding: ResponseEncoding) -> Option<u64> {
     }
 }
 
+fn visible_arrow_features(encoding: &ResponseEncoding) -> Option<ArrowFeatures> {
+    if matches!(encoding.format, QueryResultFormat::Json) {
+        return None;
+    }
+
+    encoding
+        .arrow_result_version
+        .filter(|version| *version >= ARROW_FEATURE_NEGOTIATION_VERSION)
+        .and_then(|_| {
+            encoding
+                .arrow_features
+                .as_ref()
+                .and_then(ArrowFeatures::visible_enabled_only)
+        })
+}
+
 fn negotiate_arrow_result_version(
     req: &HttpQueryRequest,
     query_result_format: QueryResultFormat,
@@ -360,6 +401,28 @@ fn negotiate_arrow_result_version(
     )))
 }
 
+fn negotiate_arrow_features(
+    req: &HttpQueryRequest,
+    arrow_result_version: Option<u64>,
+) -> Option<ArrowFeatures> {
+    if arrow_result_version? < ARROW_FEATURE_NEGOTIATION_VERSION {
+        return None;
+    }
+
+    Some(
+        match req
+            .arrow_features
+            .as_ref()
+            .and_then(|features| features.decimal64)
+        {
+            Some(false) => ArrowFeatures {
+                decimal64: Some(false),
+            },
+            _ => ArrowFeatures::decimal64_enabled(),
+        },
+    )
+}
+
 fn require_negotiated_arrow_result_version(
     query_id: &str,
     query_result_format: QueryResultFormat,
@@ -382,10 +445,11 @@ pub struct StateResponse {
     pub stats: QueryStats,
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 struct ResponseEncoding {
     format: QueryResultFormat,
     arrow_result_version: Option<u64>,
+    arrow_features: Option<ArrowFeatures>,
 }
 
 impl StateResponse {
@@ -477,6 +541,7 @@ async fn query_final_handler(
                         response.state.arrow_result_version,
                     )
                     .map_err(HttpErrorCode::bad_request)?,
+                    arrow_features: response.state.arrow_features.clone(),
                 };
                 // it is safe to set these 2 fields to None, because client now check for null/None first.
                 response.session = None;
@@ -610,6 +675,7 @@ async fn query_page_handler(
                         resp.state.arrow_result_version,
                     )
                     .map_err(HttpErrorCode::bad_request)?,
+                    arrow_features: resp.state.arrow_features.clone(),
                 };
                 query
                     .update_expire_time(false, resp.is_data_drained())
@@ -699,6 +765,7 @@ pub(crate) async fn query_handler(
                             resp.state.arrow_result_version,
                         )
                         .map_err(HttpErrorCode::bad_request)?,
+                        arrow_features: resp.state.arrow_features.clone(),
                     };
                     Ok(QueryResponse::from_internal(
                         ctx.query_id.clone(),
@@ -713,8 +780,16 @@ pub(crate) async fn query_handler(
         let arrow_result_version =
             negotiate_arrow_result_version(&req, query_result_format, ctx.user_agent.as_deref())
                 .map_err(HttpErrorCode::bad_request)?;
+        let arrow_features = negotiate_arrow_features(&req, arrow_result_version);
 
-        match HttpQuery::try_create(ctx, req.clone(), arrow_result_version).await {
+        match HttpQuery::try_create(
+            ctx,
+            req.clone(),
+            arrow_result_version,
+            arrow_features.clone(),
+        )
+        .await
+        {
             Err(err) => {
                 let err = err.display_with_sql(&sql);
                 error!("Failed to start SQL query, error: {:?}", err);
@@ -774,6 +849,7 @@ pub(crate) async fn query_handler(
                     ResponseEncoding {
                         format: query_result_format,
                         arrow_result_version,
+                        arrow_features,
                     },
                 )
                 .into_response())
@@ -1200,3 +1276,51 @@ impl<'a> FromRequest<'a> for QueryResultFormat {
             .unwrap_or(QueryResultFormat::Json))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::servers::http::v1::query::HttpQueryRequest;
+
+    fn dummy_request(arrow_features: Option<ArrowFeatures>) -> HttpQueryRequest {
+        HttpQueryRequest {
+            session_id: None,
+            session: None,
+            sql: "select 1".to_string(),
+            pagination: Default::default(),
+            string_fields: true,
+            stage_attachment: None,
+            params: None,
+            arrow_result_version_max: None,
+            arrow_features,
+        }
+    }
+
+    #[test]
+    fn test_negotiate_arrow_features_preserves_explicit_decimal64_disable() {
+        let features = negotiate_arrow_features(
+            &dummy_request(Some(ArrowFeatures {
+                decimal64: Some(false),
+            })),
+            Some(ARROW_FEATURE_NEGOTIATION_VERSION),
+        );
+        assert_eq!(
+            features,
+            Some(ArrowFeatures {
+                decimal64: Some(false),
+            })
+        );
+    }
+
+    #[test]
+    fn test_visible_arrow_features_hides_disabled_entries() {
+        let encoding = ResponseEncoding {
+            format: QueryResultFormat::Arrow,
+            arrow_result_version: Some(ARROW_FEATURE_NEGOTIATION_VERSION),
+            arrow_features: Some(ArrowFeatures {
+                decimal64: Some(false),
+            }),
+        };
+        assert_eq!(visible_arrow_features(&encoding), None);
+    }
+}