link the config to expr_udf

risingwavelabs · Feb 20, 2024 · fa024ca · fa024ca
1 parent 89a5c86
commit fa024ca
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/expr/core/src/expr/expr_udf.rs b/src/expr/core/src/expr/expr_udf.rs
@@ -55,7 +55,7 @@ pub struct UserDefinedFunction {
     /// See <https://github.com/risingwavelabs/risingwave/issues/13791>.
     disable_retry_count: AtomicU8,
     /// Always retry. Overrides `disable_retry_count`.
-    always_retry: bool,
+    always_retry_on_network_error: bool,
 }
 
 const INITIAL_RETRY_COUNT: u8 = 16;
@@ -128,9 +128,9 @@ impl UserDefinedFunction {
             UdfImpl::JavaScript(runtime) => runtime.call(&self.identifier, &input)?,
             UdfImpl::External(client) => {
                 let disable_retry_count = self.disable_retry_count.load(Ordering::Relaxed);
-                let result = if self.always_retry {
+                let result = if self.always_retry_on_network_error {
                     client
-                        .call_with_always_retry(&self.identifier, input)
+                        .call_with_always_retry_on_network_error(&self.identifier, input)
                         .instrument_await(self.span.clone())
                         .await
                 } else if disable_retry_count != 0 {
@@ -256,7 +256,7 @@ impl Build for UserDefinedFunction {
             identifier: identifier.clone(),
             span: format!("udf_call({})", identifier).into(),
             disable_retry_count: AtomicU8::new(0),
-            always_retry: true,
+            always_retry_on_network_error: udf.always_retry_on_network_error,
         })
     }
 }

diff --git a/src/expr/udf/src/external.rs b/src/expr/udf/src/external.rs
@@ -173,7 +173,7 @@ impl ArrowFlightUdfClient {
     }
 
     /// Always retry on connection
-    pub async fn call_with_always_retry(
+    pub async fn call_with_always_retry_on_network_error(
         &self,
         id: &str,
         input: RecordBatch,