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

Author: mach-kernel

ballista/client/src/extension.rs

@@ -300,7 +300,7 @@ impl SessionContextExt for SessionContext {
                 continue;
             }
 
-            self.register_udf(ScalarUDF::new_from_impl(RemoteScalarUDF::new(udf)))
+            self.register_udf(ScalarUDF::new_from_impl(RemoteScalarUDF::new(udf)?))
         }
 
         Ok(())

ballista/core/proto/ballista.proto

@@ -772,6 +772,12 @@ 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 {

ballista/core/src/remote_catalog/mod.rs

@@ -17,6 +17,6 @@
 //
 
 pub mod catalog_serialize_ext;
-pub mod function_serialize_ext;
+pub mod remote_function_serialize_ext;
 pub mod remote_scalar_udf;
 pub mod remote_table_provider;

…remote_catalog/function_serialize_ext.rs …catalog/remote_function_serialize_ext.rsballista/core/src/remote_catalog/function_serialize_ext.rs renamed to ballista/core/src/remote_catalog/remote_function_serialize_ext.rs ballista/core/src/remote_catalog/function_serialize_ext.rs renamed to ballista/core/src/remote_catalog/remote_function_serialize_ext.rs

@@ -16,42 +16,21 @@
 // under the License.
 
 use crate::serde::protobuf::{
-    ScalarUdfDocumentation, ScalarUdfInfo, ScalarUdfTypeSignature,
+    ScalarUdfDocumentation, ScalarUdfDocumentationArgument, ScalarUdfInfo,
+    ScalarUdfTypeSignature,
 };
-use arrow::datatypes::DataType;
 use datafusion::execution::FunctionRegistry;
 use datafusion::functions::all_default_functions;
-use datafusion::logical_expr::{ScalarUDF, UserDefinedLogicalNode};
 use datafusion::prelude::SessionContext;
 use datafusion_proto_common::ArrowType;
 use std::collections::HashSet;
 
 /// Used to serialize function shapes to ship to Ballista clients
-pub trait FunctionSerializeExt {
+pub trait RemoteFunctionSerializeExt {
     fn serialize_udfs(&self) -> Vec<ScalarUdfInfo>;
 }
 
-fn try_derive_return_type(udf: &ScalarUDF, types: &[DataType]) -> Option<DataType> {
-    if let Ok(return_type) = udf.return_type(&types) {
-        return Some(return_type);
-    }
-
-    // TODO: try to serialize these with dummy values
-    // let fields = types.iter().enumerate().map(|(i, t)| Arc::new(Field::new(
-    //     format!("Field{}", i),
-    //     t.clone(),
-    //     true
-    // ))).collect::<Vec<_>>();
-    //
-    // let return_type = f.return_field_from_args(ReturnFieldArgs {
-    //     arg_fields: &fields,
-    //     scalar_arguments: &vec![None; fields.len()],
-    // }).expect("Must have return type");
-
-    None
-}
-
-impl FunctionSerializeExt for SessionContext {
+impl RemoteFunctionSerializeExt for SessionContext {
     fn serialize_udfs(&self) -> Vec<ScalarUdfInfo> {
         let mut udfs = vec![];
 
@@ -78,26 +57,34 @@ impl FunctionSerializeExt for SessionContext {
                         .collect::<Result<Vec<ArrowType>, _>>()
                         .expect("Must serialize data types");
 
-                    if let Some(ref return_type) = try_derive_return_type(f.as_ref(), &t)
-                    {
-                        Some(ScalarUdfTypeSignature {
+                    // TODO: some functions use `ScalarUDF::return_field_from_args`, which this does not support
+                    f.return_type(&t)
+                        .ok()
+                        .and_then(|ref return_type| return_type.try_into().ok())
+                        .map(|arrow_return_type| ScalarUdfTypeSignature {
                             arity,
-                            return_type: Some(
-                                return_type
-                                    .try_into()
-                                    .expect("Must serialize return type"),
-                            ),
+                            return_type: Some(arrow_return_type),
                         })
-                    } else {
-                        None
-                    }
                 })
                 .collect::<Vec<_>>();
 
-            let docs = f.documentation().map(|d| ScalarUdfDocumentation {
-                description: d.description.clone(),
-                syntax_example: d.syntax_example.clone(),
-                sql_example: d.sql_example.clone(),
+            let docs = f.documentation().map(|d| {
+                let arguments = d
+                    .arguments
+                    .iter()
+                    .flatten()
+                    .map(|(arg, desc)| ScalarUdfDocumentationArgument {
+                        argument: arg.clone(),
+                        description: desc.clone(),
+                    })
+                    .collect::<Vec<_>>();
+
+                ScalarUdfDocumentation {
+                    description: d.description.clone(),
+                    syntax_example: d.syntax_example.clone(),
+                    sql_example: d.sql_example.clone(),
+                    arguments,
+                }
             });
 
             udfs.push(ScalarUdfInfo {

ballista/core/src/remote_catalog/remote_scalar_udf.rs

@@ -18,6 +18,7 @@
 
 use crate::serde::protobuf::ScalarUdfInfo;
 use arrow::datatypes::DataType;
+use datafusion::common::Result;
 use datafusion::common::{exec_err, plan_err, DataFusionError};
 use datafusion::logical_expr::{
     ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
@@ -55,36 +56,40 @@ impl PartialEq for RemoteScalarUDF {
 impl Eq for RemoteScalarUDF {}
 
 impl RemoteScalarUDF {
-    pub fn new(meta: ScalarUdfInfo) -> Self {
-        let arities = meta
-            .signatures
-            .iter()
-            .map(|signature| {
-                signature
-                    .arity
-                    .iter()
-                    .map(|t| t.try_into())
-                    .collect::<Result<Vec<DataType>, _>>()
-                    .expect("Must deserialize arities")
-            })
-            .collect();
+    pub fn new(meta: ScalarUdfInfo) -> Result<Self> {
+        let mut arities = vec![];
+
+        for signature in &meta.signatures {
+            let signature_types = signature
+                .arity
+                .iter()
+                .map(|t| t.try_into())
+                .collect::<Result<Vec<DataType>, _>>()?;
+
+            arities.push(signature_types);
+        }
 
         let documentation = meta.documentation.clone().map(|d| Documentation {
             doc_section: Default::default(),
             description: d.description.clone(),
             syntax_example: d.syntax_example.clone(),
             sql_example: d.sql_example.clone(),
-            arguments: None,
+            arguments: Some(
+                d.arguments
+                    .iter()
+                    .map(|a| (a.argument.clone(), a.description.clone()))
+                    .collect(),
+            ),
             alternative_syntax: None,
             related_udfs: None,
         });
 
-        Self {
+        Ok(Self {
             arities,
             documentation,
             meta,
             signature: Signature::new(TypeSignature::VariadicAny, Volatility::Volatile),
-        }
+        })
     }
 }
 
@@ -126,7 +131,7 @@ impl ScalarUDFImpl for RemoteScalarUDF {
 
     fn invoke_with_args(
         &self,
-        args: ScalarFunctionArgs,
+        _args: ScalarFunctionArgs,
     ) -> datafusion::common::Result<ColumnarValue> {
         exec_err!("This is a stub function and should never be called on the client")
     }

ballista/core/src/serde/generated/ballista.rs

@@ -1151,6 +1151,15 @@ pub struct ScalarUdfDocumentation {
     pub syntax_example: ::prost::alloc::string::String,
     #[prost(string, optional, tag = "3")]
     pub sql_example: ::core::option::Option<::prost::alloc::string::String>,
+    #[prost(message, repeated, tag = "4")]
+    pub arguments: ::prost::alloc::vec::Vec<ScalarUdfDocumentationArgument>,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct ScalarUdfDocumentationArgument {
+    #[prost(string, tag = "1")]
+    pub argument: ::prost::alloc::string::String,
+    #[prost(string, tag = "2")]
+    pub description: ::prost::alloc::string::String,
 }
 /// Generated client implementations.
 pub mod scheduler_grpc_client {

ballista/core/src/serde/mod.rs

@@ -50,7 +50,6 @@ use crate::remote_catalog::remote_table_provider::RemoteTableProvider;
 use crate::serde::protobuf::ballista_physical_plan_node::PhysicalPlanType;
 use crate::serde::scheduler::PartitionLocation;
 use datafusion::catalog::TableProvider;
-use datafusion::logical_expr::UserDefinedLogicalNode;
 pub use generated::ballista as protobuf;
 use prost::Message;
 use std::fmt::Debug;

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -44,7 +44,7 @@ use crate::cluster::{bind_task_bias, bind_task_round_robin};
 use crate::config::TaskDistributionPolicy;
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use crate::scheduler_server::SchedulerServer;
-use ballista_core::remote_catalog::function_serialize_ext::FunctionSerializeExt;
+use ballista_core::remote_catalog::remote_function_serialize_ext::RemoteFunctionSerializeExt;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tonic::{Request, Response, Status};