refactor(udf): connect lazily to get rid of block_in_place (#12413)

Signed-off-by: Bugen Zhao <[email protected]>

src/expr/src/expr/expr_udf.rs

@@ -164,9 +164,7 @@ pub(crate) fn get_or_create_client(link: &str) -> Result<Arc<ArrowFlightUdfClien
         Ok(client)
     } else {
         // create new client
-        let client = Arc::new(tokio::task::block_in_place(|| {
-            tokio::runtime::Handle::current().block_on(ArrowFlightUdfClient::connect(link))
-        })?);
+        let client = Arc::new(ArrowFlightUdfClient::connect_lazy(link)?);
         clients.insert(link.into(), Arc::downgrade(&client));
         Ok(client)
     }

src/udf/src/error.rs

@@ -23,7 +23,7 @@ pub enum Error {
     #[error("failed to connect to UDF service: {0}")]
     Connect(#[from] tonic::transport::Error),
 
-    #[error("failed to check UDF: {0}")]
+    #[error("failed to send requests to UDF service: {0}")]
     Tonic(#[from] Box<tonic::Status>),
 
     #[error("failed to call UDF: {0}")]

src/udf/src/external.rs

@@ -38,6 +38,13 @@ impl ArrowFlightUdfClient {
         Ok(Self { client })
     }
 
+    /// Connect to a UDF service lazily (i.e. only when the first request is sent).
+    pub fn connect_lazy(addr: &str) -> Result<Self> {
+        let conn = tonic::transport::Endpoint::new(addr.to_string())?.connect_lazy();
+        let client = FlightServiceClient::new(conn);
+        Ok(Self { client })
+    }
+
     /// Check if the function is available and the schema is match.
     pub async fn check(&self, id: &str, args: &Schema, returns: &Schema) -> Result<()> {
         let descriptor = FlightDescriptor::new_path(vec![id.into()]);