feat(query): add external_server_request_batch_rows setting (#15440)

Author: sundy-li

src/query/expression/src/function.rs

@@ -107,6 +107,7 @@ pub struct FunctionContext {
 
     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,

src/query/expression/src/utils/udf_client.rs

@@ -40,6 +40,7 @@ const MAX_DECODING_MESSAGE_SIZE: usize = 16 * 1024 * 1024 * 1024;
 #[derive(Debug, Clone)]
 pub struct UDFFlightClient {
     inner: FlightServiceClient<Channel>,
+    batch_rows: u64,
 }
 
 impl UDFFlightClient {
@@ -48,6 +49,7 @@ impl UDFFlightClient {
         addr: &str,
         conn_timeout: u64,
         request_timeout: u64,
+        batch_rows: u64,
     ) -> Result<UDFFlightClient> {
         let endpoint = Endpoint::from_shared(addr.to_string())
             .map_err(|err| {
@@ -69,7 +71,7 @@ impl UDFFlightClient {
             })?
             .max_decoding_message_size(MAX_DECODING_MESSAGE_SIZE);
 
-        Ok(UDFFlightClient { inner })
+        Ok(UDFFlightClient { inner, batch_rows })
     }
 
     fn make_request<T>(&self, t: T) -> Request<T> {
@@ -135,9 +137,16 @@ impl UDFFlightClient {
         input_batch: RecordBatch,
     ) -> Result<RecordBatch> {
         let descriptor = FlightDescriptor::new_path(vec![func_name.to_string()]);
+        let batch_rows = self.batch_rows as usize;
+        let batches = (0..input_batch.num_rows())
+            .step_by(batch_rows)
+            .map(move |start| {
+                Ok(input_batch.slice(start, batch_rows.min(input_batch.num_rows() - start)))
+            });
+
         let flight_data_stream = FlightDataEncoderBuilder::new()
             .with_flight_descriptor(Some(descriptor))
-            .build(stream::iter(vec![Ok(input_batch)]))
+            .build(stream::iter(batches))
             .map(|data| data.unwrap());
         let request = self.make_request(flight_data_stream);
         let flight_data_stream = self.inner.do_exchange(request).await?.into_inner();

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

@@ -59,6 +59,7 @@ impl AsyncTransform for TransformUdfServer {
     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
@@ -91,10 +92,15 @@ impl AsyncTransform for TransformUdfServer {
                 .to_record_batch_with_dataschema(&data_schema)
                 .map_err(|err| ErrorCode::from_string(format!("{err}")))?;
 
-            let mut client =
-                UDFFlightClient::connect(server_addr, connect_timeout, request_timeout).await?;
-            let result_batch = client.do_exchange(&func.func_name, input_batch).await?;
+            let mut client = UDFFlightClient::connect(
+                server_addr,
+                connect_timeout,
+                request_timeout,
+                request_bacth_rows,
+            )
+            .await?;
 
+            let result_batch = client.do_exchange(&func.func_name, input_batch).await?;
             let schema = DataSchema::try_from(&(*result_batch.schema()))?;
             let (result_block, result_schema) =
                 DataBlock::from_record_batch(&schema, &result_batch).map_err(|err| {

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

@@ -618,6 +618,8 @@ impl TableContext for QueryContext {
             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)?;
@@ -642,6 +644,7 @@ impl TableContext for QueryContext {
 
             external_server_connect_timeout_secs,
             external_server_request_timeout_secs,
+            external_server_request_batch_rows,
             geometry_output_format,
             parse_datetime_ignore_remainder,
         })

src/query/settings/src/settings_default.rs

@@ -630,6 +630,12 @@ impl DefaultSettings {
                     mode: SettingMode::Both,
                     range: Some(SettingRange::Numeric(0..=u64::MAX)),
                 }),
+                ("external_server_request_batch_rows", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(65536),
+                    desc: "Request batch rows to external server",
+                    mode: SettingMode::Both,
+                    range: Some(SettingRange::Numeric(1..=u64::MAX)),
+                }),
                 ("enable_parquet_prewhere", DefaultSettingValue {
                     value: UserSettingValue::UInt64(0),
                     desc: "Enables parquet prewhere",

src/query/settings/src/settings_getter_setter.rs

@@ -557,6 +557,10 @@ impl Settings {
         self.try_get_u64("external_server_request_timeout_secs")
     }
 
+    pub fn get_external_server_request_batch_rows(&self) -> Result<u64> {
+        self.try_get_u64("external_server_request_batch_rows")
+    }
+
     pub fn get_create_query_flight_client_with_current_rt(&self) -> Result<bool> {
         Ok(self.try_get_u64("create_query_flight_client_with_current_rt")? != 0)
     }

src/query/sql/src/planner/binder/udf.rs

@@ -110,6 +110,9 @@ impl Binder {
                     self.ctx
                         .get_settings()
                         .get_external_server_request_timeout_secs()?,
+                    self.ctx
+                        .get_settings()
+                        .get_external_server_request_batch_rows()?,
                 )
                 .await?;
                 client