refactor(query): add setting external_server_request_retry_times (#16307)

* refactor(query): add setting external_server_request_retry_times and profile the retry counts

* safety index_unchecked for nullable

* safety index_unchecked for nullable

* update

* update

* update

Author: sundy-li

src/common/base/src/runtime/profile/profiles.rs

@@ -45,6 +45,7 @@ pub enum ProfileStatisticsName {
     SpillReadTime,
     RuntimeFilterPruneParts,
     MemoryUsage,
+    ExternalServerRetryCount,
 }
 
 #[derive(Clone, Hash, Eq, PartialEq, serde::Serialize, serde::Deserialize, Debug)]
@@ -242,7 +243,14 @@ pub fn get_statistics_desc() -> Arc<BTreeMap<ProfileStatisticsName, ProfileDesc>
                 index: ProfileStatisticsName::MemoryUsage as usize,
                 unit: StatisticsUnit::Bytes,
                 plain_statistics: false,
-            })
+            }),
+            (ProfileStatisticsName::ExternalServerRetryCount, ProfileDesc {
+                display_name: "external server retry count",
+                desc: "The count of external server retry times",
+                index: ProfileStatisticsName::ExternalServerRetryCount as usize,
+                unit: StatisticsUnit::Count,
+                plain_statistics: true,
+            }),
         ]))
     }).clone()
 }

src/query/expression/src/function.rs

@@ -109,10 +109,6 @@ pub struct FunctionContext {
     pub openai_api_embedding_model: String,
     pub openai_api_completion_model: String,
 
-    pub external_server_connect_timeout_secs: u64,
-    pub external_server_request_timeout_secs: u64,
-    pub external_server_request_batch_rows: u64,
-
     pub geometry_output_format: GeometryDataType,
     pub parse_datetime_ignore_remainder: bool,
     pub enable_dst_hour_fix: bool,
@@ -133,9 +129,7 @@ impl Default for FunctionContext {
             openai_api_version: "".to_string(),
             openai_api_embedding_model: "".to_string(),
             openai_api_completion_model: "".to_string(),
-            external_server_connect_timeout_secs: 0,
-            external_server_request_timeout_secs: 0,
-            external_server_request_batch_rows: 0,
+
             geometry_output_format: Default::default(),
             parse_datetime_ignore_remainder: false,
             enable_dst_hour_fix: false,

src/query/expression/src/types/nullable.rs

@@ -258,8 +258,8 @@ impl<T: ValueType> NullableColumn<T> {
     // we should better use new to create a new instance to ensure the validity and column are consistent
     // todo: make column and validity private
     pub fn new(column: T::Column, validity: Bitmap) -> Self {
-        assert_eq!(T::column_len(&column), validity.len());
-        assert!(!matches!(
+        debug_assert_eq!(T::column_len(&column), validity.len());
+        debug_assert!(!matches!(
             T::upcast_column(column.clone()),
             Column::Nullable(_)
         ));
@@ -290,11 +290,15 @@ impl<T: ValueType> NullableColumn<T> {
     ///
     /// Calling this method with an out-of-bounds index is *[undefined behavior]*
     pub unsafe fn index_unchecked(&self, index: usize) -> Option<T::ScalarRef<'_>> {
-        debug_assert!(index < self.validity.len());
-
-        match self.validity.get_bit_unchecked(index) {
-            true => Some(T::index_column(&self.column, index).unwrap()),
-            false => None,
+        // we need to check the validity firstly
+        // cause `self.validity.get_bit_unchecked` may check the index from buffer address with `true` result
+        if index < self.validity.len() {
+            match self.validity.get_bit_unchecked(index) {
+                true => Some(T::index_column_unchecked(&self.column, index)),
+                false => None,
+            }
+        } else {
+            None
         }
     }
 

src/query/functions/src/scalars/variant.rs

@@ -904,10 +904,10 @@ pub fn register(registry: &mut FunctionRegistry) {
                         };
                         let new_col = cast_scalars_to_variants(col.iter(), ctx.func_ctx.tz);
                         if let Some(validity) = validity {
-                            Value::Column(Column::Nullable(Box::new(NullableColumn {
+                            Value::Column(NullableColumn::new_column(
+                                Column::Variant(new_col),
                                 validity,
-                                column: Column::Variant(new_col),
-                            })))
+                            ))
                         } else {
                             Value::Column(Column::Variant(new_col))
                         }
@@ -945,10 +945,7 @@ pub fn register(registry: &mut FunctionRegistry) {
                     _ => Bitmap::new_constant(true, col.len()),
                 };
                 let new_col = cast_scalars_to_variants(col.iter(), ctx.func_ctx.tz);
-                Value::Column(NullableColumn {
-                    validity,
-                    column: new_col,
-                })
+                Value::Column(NullableColumn::new(new_col, validity))
             }
         },
     );

src/query/functions/tests/it/scalars/tuple.rs

@@ -61,12 +61,12 @@ fn test_get(file: &mut impl Write) {
     )]);
     run_ast(file, "col.1", &[(
         "col",
-        Column::Nullable(Box::new(NullableColumn {
-            column: Column::Tuple(vec![StringType::from_data_with_validity(
+        NullableColumn::new_column(
+            Column::Tuple(vec![StringType::from_data_with_validity(
                 vec!["a", "b", "c", "d"],
                 vec![true, true, false, false],
             )]),
-            validity: vec![true, false, true, false].into(),
-        })),
+            vec![true, false, true, false].into(),
+        ),
     )]);
 }

src/query/pipeline/transforms/src/processors/transforms/transform_retry_async.rs

@@ -20,6 +20,8 @@ use super::AsyncTransform;
 
 pub trait AsyncRetry: AsyncTransform {
     fn retry_on(&self, err: &databend_common_exception::ErrorCode) -> bool;
+    // record some log when retrying
+    fn retry_hook(&self);
     fn retry_strategy(&self) -> RetryStrategy;
 }
 
@@ -67,6 +69,7 @@ impl<T: AsyncRetry + 'static> AsyncTransform for AsyncRetryWrapper<T> {
                         tokio::time::sleep(duration).await;
                     }
                     self.retries += 1;
+                    self.t.retry_hook();
                 }
             }
         }

src/query/service/src/pipelines/builders/builder_udf.rs

@@ -42,11 +42,7 @@ impl PipelineBuilder {
             })
         } else {
             self.main_pipeline.try_add_async_transformer(|| {
-                Ok(TransformUdfServer::new_retry_wrapper(
-                    self.ctx.clone(),
-                    self.func_ctx.clone(),
-                    udf.udf_funcs.clone(),
-                ))
+                TransformUdfServer::new_retry_wrapper(self.ctx.clone(), udf.udf_funcs.clone())
             })
         }
     }

src/query/service/src/pipelines/processors/transforms/transform_udf_server.rs

@@ -14,6 +14,8 @@
 
 use std::sync::Arc;
 
+use databend_common_base::runtime::profile::Profile;
+use databend_common_base::runtime::profile::ProfileStatisticsName;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -24,7 +26,6 @@ use databend_common_expression::BlockEntry;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchema;
-use databend_common_expression::FunctionContext;
 use databend_common_pipeline_transforms::processors::AsyncRetry;
 use databend_common_pipeline_transforms::processors::AsyncRetryWrapper;
 use databend_common_pipeline_transforms::processors::AsyncTransform;
@@ -35,22 +36,34 @@ use crate::sessions::QueryContext;
 
 pub struct TransformUdfServer {
     ctx: Arc<QueryContext>,
-    func_ctx: FunctionContext,
     funcs: Vec<UdfFunctionDesc>,
+    connect_timeout: u64,
+    request_timeout: u64,
+    request_bacth_rows: u64,
+    retry_times: u64,
 }
 
 impl TransformUdfServer {
     pub fn new_retry_wrapper(
         ctx: Arc<QueryContext>,
-        func_ctx: FunctionContext,
         funcs: Vec<UdfFunctionDesc>,
-    ) -> AsyncRetryWrapper<Self> {
+    ) -> Result<AsyncRetryWrapper<Self>> {
+        let settings = ctx.get_settings();
+        let connect_timeout = settings.get_external_server_connect_timeout_secs()?;
+        let request_timeout = settings.get_external_server_request_timeout_secs()?;
+        let request_bacth_rows = settings.get_external_server_request_batch_rows()?;
+        let retry_times = settings.get_external_server_request_retry_times()?;
+
         let s = Self {
             ctx,
-            func_ctx,
             funcs,
+
+            connect_timeout,
+            request_timeout,
+            request_bacth_rows,
+            retry_times,
         };
-        AsyncRetryWrapper::create(s)
+        Ok(AsyncRetryWrapper::create(s))
     }
 }
 
@@ -61,10 +74,14 @@ impl AsyncRetry for TransformUdfServer {
 
     fn retry_strategy(&self) -> RetryStrategy {
         RetryStrategy {
-            retry_times: 64,
+            retry_times: self.retry_times as usize,
             retry_sleep_duration: Some(tokio::time::Duration::from_millis(500)),
         }
     }
+
+    fn retry_hook(&self) {
+        Profile::record_usize_profile(ProfileStatisticsName::ExternalServerRetryCount, 1);
+    }
 }
 
 #[async_trait::async_trait]
@@ -73,9 +90,6 @@ impl AsyncTransform for TransformUdfServer {
 
     #[async_backtrace::framed]
     async fn transform(&mut self, mut data_block: DataBlock) -> Result<DataBlock> {
-        let connect_timeout = self.func_ctx.external_server_connect_timeout_secs;
-        let request_timeout = self.func_ctx.external_server_request_timeout_secs;
-        let request_bacth_rows = self.func_ctx.external_server_request_batch_rows;
         for func in &self.funcs {
             let server_addr = func.udf_type.as_server().unwrap();
             // construct input record_batch
@@ -110,9 +124,9 @@ impl AsyncTransform for TransformUdfServer {
 
             let mut client = UDFFlightClient::connect(
                 server_addr,
-                connect_timeout,
-                request_timeout,
-                request_bacth_rows,
+                self.connect_timeout,
+                self.request_timeout,
+                self.request_bacth_rows,
             )
             .await?
             .with_tenant(self.ctx.get_tenant().tenant_name())?

src/query/service/src/pipelines/processors/transforms/window/transform_window.rs

@@ -1003,10 +1003,12 @@ where T: Number + ResultTypeOfUnary
     fn add_block(&mut self, data: Option<DataBlock>) -> Result<()> {
         if let Some(data) = data {
             let num_rows = data.num_rows();
-            self.blocks.push_back(WindowBlock {
-                block: data.convert_to_full(),
-                builder: ColumnBuilder::with_capacity(&self.func.return_type()?, num_rows),
-            });
+            if num_rows != 0 {
+                self.blocks.push_back(WindowBlock {
+                    block: data.convert_to_full(),
+                    builder: ColumnBuilder::with_capacity(&self.func.return_type()?, num_rows),
+                });
+            }
         }
 
         // Each loop will do:

src/query/service/src/sessions/query_ctx.rs

@@ -695,12 +695,6 @@ impl TableContext for QueryContext {
 
     fn get_function_context(&self) -> Result<FunctionContext> {
         let settings = self.get_settings();
-        let external_server_connect_timeout_secs =
-            settings.get_external_server_connect_timeout_secs()?;
-        let external_server_request_timeout_secs =
-            settings.get_external_server_request_timeout_secs()?;
-        let external_server_request_batch_rows =
-            settings.get_external_server_request_batch_rows()?;
 
         let tz = settings.get_timezone()?;
         let tz = TzFactory::instance().get_by_name(&tz)?;
@@ -728,9 +722,6 @@ impl TableContext for QueryContext {
             openai_api_embedding_model: query_config.openai_api_embedding_model.clone(),
             openai_api_completion_model: query_config.openai_api_completion_model.clone(),
 
-            external_server_connect_timeout_secs,
-            external_server_request_timeout_secs,
-            external_server_request_batch_rows,
             geometry_output_format,
             parse_datetime_ignore_remainder,
             enable_dst_hour_fix,