Scheduler UDF sync -> client planning with stubs (#4)

* [1/?] attempt to serialize bare min data for stub udf planning in client

* [2/?] fix return type issue, basic case working

* fix error handling, serialize fn arguments for docs, add note about return field

* tomlfmt

* lint

* more lintg

Author: mach-kernel

ballista/client/Cargo.toml

@@ -36,6 +36,7 @@ datafusion = { workspace = true }
 log = { workspace = true }
 
 tokio = { workspace = true }
+tonic = { workspace = true }
 url = { workspace = true }
 
 [dev-dependencies]

ballista/client/src/extension.rs

@@ -16,10 +16,19 @@
 // under the License.
 
 pub use ballista_core::extension::{SessionConfigExt, SessionStateExt};
+use ballista_core::remote_catalog::remote_scalar_udf::RemoteScalarUDF;
+use ballista_core::remote_catalog::remote_table_provider::RemoteTableProvider;
 use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::serde::protobuf::{GetCatalogParams, GetRemoteFunctionsParams};
+use datafusion::catalog::{
+    CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider,
+};
+use datafusion::execution::FunctionRegistry;
+use datafusion::logical_expr::ScalarUDF;
 use datafusion::{
     error::DataFusionError, execution::SessionState, prelude::SessionContext,
 };
+use std::sync::Arc;
 use url::Url;
 
 const DEFAULT_SCHEDULER_PORT: u16 = 50050;
@@ -86,6 +95,19 @@ pub trait SessionContextExt {
         url: &str,
         state: SessionState,
     ) -> datafusion::error::Result<SessionContext>;
+
+    /// Populates the local catalog with metadata from the remote scheduler.
+    /// This allows catalog queries like SHOW TABLES to work on the client.
+    async fn populate_catalog_from_scheduler(
+        &self,
+        scheduler_url: &str,
+    ) -> datafusion::error::Result<()>;
+
+    /// Populates local context with functions from the scheduler.
+    async fn populate_functions_from_scheduler(
+        &self,
+        scheduler_url: &str,
+    ) -> datafusion::error::Result<()>;
 }
 
 #[async_trait::async_trait]
@@ -100,14 +122,21 @@ impl SessionContextExt for SessionContext {
             scheduler_url.clone()
         );
 
-        let session_state = state.upgrade_for_ballista(scheduler_url)?;
+        let session_state = state.upgrade_for_ballista(scheduler_url.clone())?;
 
         log::info!(
             "Server side SessionContext created with session id: {}",
             session_state.session_id()
         );
 
-        Ok(SessionContext::new_with_state(session_state))
+        let ctx = SessionContext::new_with_state(session_state);
+
+        // Populate local catalog from scheduler
+        ctx.populate_catalog_from_scheduler(&scheduler_url).await?;
+        ctx.populate_functions_from_scheduler(&scheduler_url)
+            .await?;
+
+        Ok(ctx)
     }
 
     async fn remote(url: &str) -> datafusion::error::Result<SessionContext> {
@@ -117,13 +146,20 @@ impl SessionContextExt for SessionContext {
             scheduler_url.clone()
         );
 
-        let session_state = SessionState::new_ballista_state(scheduler_url)?;
+        let session_state = SessionState::new_ballista_state(scheduler_url.clone())?;
         log::info!(
             "Server side SessionContext created with session id: {}",
             session_state.session_id()
         );
 
-        Ok(SessionContext::new_with_state(session_state))
+        let ctx = SessionContext::new_with_state(session_state);
+
+        // Populate local catalog from scheduler
+        ctx.populate_catalog_from_scheduler(&scheduler_url).await?;
+        ctx.populate_functions_from_scheduler(&scheduler_url)
+            .await?;
+
+        Ok(ctx)
     }
 
     #[cfg(feature = "standalone")]
@@ -157,6 +193,118 @@ impl SessionContextExt for SessionContext {
 
         Ok(SessionContext::new_with_state(session_state))
     }
+
+    async fn populate_catalog_from_scheduler(
+        &self,
+        scheduler_url: &str,
+    ) -> datafusion::error::Result<()> {
+        let mut client = SchedulerGrpcClient::connect(scheduler_url.to_string())
+            .await
+            .map_err(|e| {
+                DataFusionError::External(
+                    format!("Failed to connect to scheduler: {}", e).into(),
+                )
+            })?;
+
+        let request = tonic::Request::new(GetCatalogParams {
+            session_id: self.state().session_id().to_string(),
+        });
+
+        let response = client.get_catalog(request).await.map_err(|e| {
+            DataFusionError::External(format!("Failed to fetch catalog: {}", e).into())
+        })?;
+
+        let catalog_result = response.into_inner();
+
+        log::info!(
+            "Received {} catalogs from scheduler",
+            catalog_result.catalogs.len()
+        );
+
+        for catalog_info in catalog_result.catalogs {
+            let catalog_name = catalog_info.catalog_name;
+
+            let catalog: Arc<dyn CatalogProvider> =
+                if let Some(_existing_catalog) = self.catalog(&catalog_name) {
+                    continue;
+                } else {
+                    let new_catalog: Arc<dyn CatalogProvider> =
+                        Arc::new(MemoryCatalogProvider::new());
+                    self.register_catalog(&catalog_name, Arc::clone(&new_catalog));
+                    new_catalog
+                };
+
+            for schema_info in catalog_info.schemas {
+                let schema_name = schema_info.schema_name;
+                let schema = if let Some(_existing_schema) = catalog.schema(&schema_name)
+                {
+                    continue;
+                } else {
+                    let new_schema: Arc<dyn SchemaProvider> =
+                        Arc::new(MemorySchemaProvider::new());
+                    catalog.register_schema(&schema_name, Arc::clone(&new_schema))?;
+                    new_schema
+                };
+
+                // Bind `RemoteTableProvider` tables
+                for table_info in schema_info.tables {
+                    if let Some(proto_schema) = table_info.schema {
+                        let arrow_schema: datafusion::arrow::datatypes::Schema =
+                            (&proto_schema)
+                                .try_into()
+                                .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+                        let stub_table = RemoteTableProvider::new(
+                            &catalog_name,
+                            &schema_name,
+                            &table_info.table_name,
+                            Arc::new(arrow_schema),
+                        );
+
+                        schema.register_table(
+                            table_info.table_name.clone(),
+                            Arc::new(stub_table),
+                        )?;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn populate_functions_from_scheduler(
+        &self,
+        scheduler_url: &str,
+    ) -> datafusion::common::Result<()> {
+        let mut client = SchedulerGrpcClient::connect(scheduler_url.to_string())
+            .await
+            .map_err(|e| {
+                DataFusionError::External(
+                    format!("Failed to connect to scheduler: {}", e).into(),
+                )
+            })?;
+
+        let request = tonic::Request::new(GetRemoteFunctionsParams {
+            session_id: self.state().session_id().to_string(),
+        });
+
+        let response = client.get_remote_functions(request).await.map_err(|e| {
+            DataFusionError::External(format!("Failed to fetch catalog: {}", e).into())
+        })?;
+
+        let remote_functions = response.into_inner();
+
+        for udf in remote_functions.udfs {
+            if self.state().udf(&udf.name).is_ok() {
+                continue;
+            }
+
+            self.register_udf(ScalarUDF::new_from_impl(RemoteScalarUDF::new(udf)?))
+        }
+
+        Ok(())
+    }
 }
 
 struct Extension {}

ballista/core/proto/ballista.proto

@@ -710,6 +710,75 @@ message RunningTaskInfo {
   uint32 partition_id = 4;;
 }
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Catalog Metadata
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+message GetCatalogParams {
+  string session_id = 1;
+}
+
+message GetCatalogResult {
+  repeated CatalogInfo catalogs = 1;
+}
+
+message CatalogInfo {
+  string catalog_name = 1;
+  repeated SchemaInfo schemas = 2;
+}
+
+message SchemaInfo {
+  string schema_name = 1;
+  repeated TableInfo tables = 2;
+}
+
+message TableInfo {
+  string table_name = 1;
+  datafusion_common.Schema schema = 2;
+}
+
+message RemoteTableProviderNode {
+  string catalog_name = 1;
+  string schema_name = 2;
+  string table_name = 3;
+  datafusion_common.Schema schema = 4;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Remote UDF Metadata
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+message GetRemoteFunctionsParams {
+  string session_id = 1;
+}
+
+message GetRemoteFunctionsResult {
+  repeated ScalarUDFInfo udfs = 1;
+}
+
+message ScalarUDFInfo {
+  string name = 1;
+  optional ScalarUDFDocumentation documentation = 2;
+  repeated ScalarUDFTypeSignature signatures = 3;
+}
+
+message ScalarUDFTypeSignature {
+  repeated datafusion_common.ArrowType arity = 1;
+  datafusion_common.ArrowType return_type = 2;
+}
+
+message ScalarUDFDocumentation {
+  string description = 1;
+  string syntax_example = 2;
+  optional string sql_example = 3;
+  repeated ScalarUDFDocumentationArgument arguments = 4;
+}
+
+message ScalarUDFDocumentationArgument {
+  string argument = 1;
+  string description = 2;
+}
+
 service SchedulerGrpc {
   // Executors must poll the scheduler for heartbeat and to receive tasks
   rpc PollWork (PollWorkParams) returns (PollWorkResult) {}
@@ -736,6 +805,12 @@ service SchedulerGrpc {
   rpc CancelJob (CancelJobParams) returns (CancelJobResult) {}
 
   rpc CleanJobData (CleanJobDataParams) returns (CleanJobDataResult) {}
+
+  // Get catalog metadata for a session
+  rpc GetCatalog (GetCatalogParams) returns (GetCatalogResult) {}
+
+  // Get catalog metadata for a session
+  rpc GetRemoteFunctions (GetRemoteFunctionsParams) returns (GetRemoteFunctionsResult) {}
 }
 
 service ExecutorGrpc {

ballista/core/src/remote_catalog/catalog_serialize_ext.rs

@@ -0,0 +1,96 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::serde::generated::ballista::{CatalogInfo, SchemaInfo, TableInfo};
+use datafusion::catalog::CatalogProvider;
+use datafusion::prelude::SessionContext;
+use futures::stream;
+use futures::StreamExt;
+use std::sync::Arc;
+
+/// Used to serialize catalog schemas and names to ship to Ballista clients
+#[async_trait::async_trait]
+pub trait CatalogSerializeExt {
+    async fn serialize_catalogs(&self) -> Vec<CatalogInfo>;
+
+    async fn serialize_catalog(&self, name: &str) -> Option<CatalogInfo>;
+
+    async fn serialize_schema(
+        &self,
+        name: &str,
+        catalog: &Arc<dyn CatalogProvider>,
+    ) -> Option<SchemaInfo>;
+}
+
+#[async_trait::async_trait]
+impl CatalogSerializeExt for SessionContext {
+    async fn serialize_catalogs(&self) -> Vec<CatalogInfo> {
+        let catalog_names = self.state().catalog_list().catalog_names();
+
+        stream::iter(catalog_names.iter())
+            .filter_map(|catalog_name| self.serialize_catalog(catalog_name))
+            .collect::<Vec<_>>()
+            .await
+    }
+
+    async fn serialize_catalog(&self, name: &str) -> Option<CatalogInfo> {
+        let catalog = self.catalog(name)?;
+
+        let schemas = stream::iter(catalog.schema_names().iter())
+            .filter_map(|schema_name| self.serialize_schema(schema_name, &catalog))
+            .collect::<Vec<_>>()
+            .await;
+
+        Some(CatalogInfo {
+            catalog_name: name.to_string(),
+            schemas,
+        })
+    }
+
+    async fn serialize_schema(
+        &self,
+        name: &str,
+        catalog: &Arc<dyn CatalogProvider>,
+    ) -> Option<SchemaInfo> {
+        let schema = catalog.schema(name)?;
+
+        let tables = stream::iter(schema.table_names())
+            .filter_map(|table_name| async {
+                let table_lookup = schema.table(&table_name).await;
+                table_lookup.ok().and_then(|maybe_table| {
+                    maybe_table.map(|provider| (table_name, provider))
+                })
+            })
+            .map(|(table_name, provider)| TableInfo {
+                table_name,
+                schema: Some(
+                    provider
+                        .schema()
+                        .as_ref()
+                        .try_into()
+                        .expect("Must serialize schema"),
+                ),
+            })
+            .collect::<Vec<_>>()
+            .await;
+
+        Some(SchemaInfo {
+            schema_name: name.to_string(),
+            tables,
+        })
+    }
+}