refactor: Move BallistaRegistry to better location

ballista/core/src/lib.rs

@@ -32,6 +32,7 @@ pub mod consistent_hash;
 pub mod error;
 pub mod event_loop;
 pub mod execution_plans;
+pub mod registry;
 pub mod utils;
 
 #[macro_use]

ballista/core/src/registry.rs

@@ -0,0 +1,112 @@
+// 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 datafusion::common::DataFusionError;
+use datafusion::execution::{FunctionRegistry, SessionState};
+use datafusion::functions::all_default_functions;
+use datafusion::functions_aggregate::all_default_aggregate_functions;
+use datafusion::functions_window::all_default_window_functions;
+use datafusion::logical_expr::planner::ExprPlanner;
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
+
+#[derive(Debug)]
+pub struct BallistaFunctionRegistry {
+    pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+    pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    pub window_functions: HashMap<String, Arc<WindowUDF>>,
+}
+
+impl Default for BallistaFunctionRegistry {
+    fn default() -> Self {
+        let scalar_functions = all_default_functions()
+            .into_iter()
+            .map(|f| (f.name().to_string(), f))
+            .collect();
+
+        let aggregate_functions = all_default_aggregate_functions()
+            .into_iter()
+            .map(|f| (f.name().to_string(), f))
+            .collect();
+
+        let window_functions = all_default_window_functions()
+            .into_iter()
+            .map(|f| (f.name().to_string(), f))
+            .collect();
+
+        Self {
+            scalar_functions,
+            aggregate_functions,
+            window_functions,
+        }
+    }
+}
+
+impl FunctionRegistry for BallistaFunctionRegistry {
+    fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
+        vec![]
+    }
+
+    fn udfs(&self) -> HashSet<String> {
+        self.scalar_functions.keys().cloned().collect()
+    }
+
+    fn udf(&self, name: &str) -> datafusion::common::Result<Arc<ScalarUDF>> {
+        let result = self.scalar_functions.get(name);
+
+        result.cloned().ok_or_else(|| {
+            DataFusionError::Internal(format!(
+                "There is no UDF named \"{name}\" in the TaskContext"
+            ))
+        })
+    }
+
+    fn udaf(&self, name: &str) -> datafusion::common::Result<Arc<AggregateUDF>> {
+        let result = self.aggregate_functions.get(name);
+
+        result.cloned().ok_or_else(|| {
+            DataFusionError::Internal(format!(
+                "There is no UDAF named \"{name}\" in the TaskContext"
+            ))
+        })
+    }
+
+    fn udwf(&self, name: &str) -> datafusion::common::Result<Arc<WindowUDF>> {
+        let result = self.window_functions.get(name);
+
+        result.cloned().ok_or_else(|| {
+            DataFusionError::Internal(format!(
+                "There is no UDWF named \"{name}\" in the TaskContext"
+            ))
+        })
+    }
+}
+
+impl From<&SessionState> for BallistaFunctionRegistry {
+    fn from(state: &SessionState) -> Self {
+        let scalar_functions = state.scalar_functions().clone();
+        let aggregate_functions = state.aggregate_functions().clone();
+        let window_functions = state.window_functions().clone();
+
+        Self {
+            scalar_functions,
+            aggregate_functions,
+            window_functions,
+        }
+    }
+}

ballista/core/src/serde/scheduler/mod.rs

@@ -15,27 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::collections::HashSet;
-use std::fmt::Debug;
-use std::{collections::HashMap, fmt, sync::Arc};
-
+use crate::error::BallistaError;
+use crate::registry::BallistaFunctionRegistry;
 use datafusion::arrow::array::{
     ArrayBuilder, StructArray, StructBuilder, UInt64Array, UInt64Builder,
 };
 use datafusion::arrow::datatypes::{DataType, Field};
-use datafusion::common::DataFusionError;
-use datafusion::execution::{FunctionRegistry, SessionState};
-use datafusion::functions::all_default_functions;
-use datafusion::functions_aggregate::all_default_aggregate_functions;
-use datafusion::functions_window::all_default_window_functions;
-use datafusion::logical_expr::planner::ExprPlanner;
-use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::Partitioning;
 use datafusion::prelude::SessionConfig;
 use serde::Serialize;
-
-use crate::error::BallistaError;
+use std::fmt::Debug;
+use std::{collections::HashMap, fmt, sync::Arc};
 
 pub mod from_proto;
 pub mod to_proto;
@@ -295,89 +286,3 @@ pub struct TaskDefinition {
     pub session_config: SessionConfig,
     pub function_registry: Arc<BallistaFunctionRegistry>,
 }
-
-#[derive(Debug)]
-pub struct BallistaFunctionRegistry {
-    pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
-    pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
-    pub window_functions: HashMap<String, Arc<WindowUDF>>,
-}
-
-impl Default for BallistaFunctionRegistry {
-    fn default() -> Self {
-        let scalar_functions = all_default_functions()
-            .into_iter()
-            .map(|f| (f.name().to_string(), f))
-            .collect();
-
-        let aggregate_functions = all_default_aggregate_functions()
-            .into_iter()
-            .map(|f| (f.name().to_string(), f))
-            .collect();
-
-        let window_functions = all_default_window_functions()
-            .into_iter()
-            .map(|f| (f.name().to_string(), f))
-            .collect();
-
-        Self {
-            scalar_functions,
-            aggregate_functions,
-            window_functions,
-        }
-    }
-}
-
-impl FunctionRegistry for BallistaFunctionRegistry {
-    fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
-        vec![]
-    }
-
-    fn udfs(&self) -> HashSet<String> {
-        self.scalar_functions.keys().cloned().collect()
-    }
-
-    fn udf(&self, name: &str) -> datafusion::common::Result<Arc<ScalarUDF>> {
-        let result = self.scalar_functions.get(name);
-
-        result.cloned().ok_or_else(|| {
-            DataFusionError::Internal(format!(
-                "There is no UDF named \"{name}\" in the TaskContext"
-            ))
-        })
-    }
-
-    fn udaf(&self, name: &str) -> datafusion::common::Result<Arc<AggregateUDF>> {
-        let result = self.aggregate_functions.get(name);
-
-        result.cloned().ok_or_else(|| {
-            DataFusionError::Internal(format!(
-                "There is no UDAF named \"{name}\" in the TaskContext"
-            ))
-        })
-    }
-
-    fn udwf(&self, name: &str) -> datafusion::common::Result<Arc<WindowUDF>> {
-        let result = self.window_functions.get(name);
-
-        result.cloned().ok_or_else(|| {
-            DataFusionError::Internal(format!(
-                "There is no UDWF named \"{name}\" in the TaskContext"
-            ))
-        })
-    }
-}
-
-impl From<&SessionState> for BallistaFunctionRegistry {
-    fn from(state: &SessionState) -> Self {
-        let scalar_functions = state.scalar_functions().clone();
-        let aggregate_functions = state.aggregate_functions().clone();
-        let window_functions = state.window_functions().clone();
-
-        Self {
-            scalar_functions,
-            aggregate_functions,
-            window_functions,
-        }
-    }
-}

ballista/executor/src/executor.rs

@@ -23,9 +23,9 @@ use crate::execution_engine::QueryStageExecutor;
 use crate::metrics::ExecutorMetricsCollector;
 use crate::metrics::LoggingMetricsCollector;
 use ballista_core::error::BallistaError;
+use ballista_core::registry::BallistaFunctionRegistry;
 use ballista_core::serde::protobuf;
 use ballista_core::serde::protobuf::ExecutorRegistration;
-use ballista_core::serde::scheduler::BallistaFunctionRegistry;
 use ballista_core::serde::scheduler::PartitionId;
 use ballista_core::ConfigProducer;
 use ballista_core::RuntimeProducer;

ballista/executor/src/executor_process.rs

@@ -25,7 +25,7 @@ use std::{env, io};
 
 use anyhow::{Context, Result};
 use arrow_flight::flight_service_server::FlightServiceServer;
-use ballista_core::serde::scheduler::BallistaFunctionRegistry;
+use ballista_core::registry::BallistaFunctionRegistry;
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use futures::stream::FuturesUnordered;

ballista/executor/src/standalone.rs

@@ -19,7 +19,7 @@ use crate::metrics::LoggingMetricsCollector;
 use crate::{execution_loop, executor::Executor, flight_service::BallistaFlightService};
 use arrow_flight::flight_service_server::FlightServiceServer;
 use ballista_core::config::BallistaConfig;
-use ballista_core::serde::scheduler::BallistaFunctionRegistry;
+use ballista_core::registry::BallistaFunctionRegistry;
 use ballista_core::utils::{default_config_producer, SessionConfigExt};
 use ballista_core::{
     error::Result,