[GitHub] [arrow] alamb commented on a change in pull request #8031: ARROW-9815: [Rust][DataFusion] Add a trait for looking up scalar functions by name

[GitBox <gi...@apache.org>]

alamb commented on a change in pull request #8031:
URL: https://github.com/apache/arrow/pull/8031#discussion_r475224655



##########
File path: rust/datafusion/src/execution/context.rs
##########
@@ -189,16 +190,7 @@ impl ExecutionContext {
     /// Register a scalar UDF
     pub fn register_udf(&mut self, f: ScalarFunction) {
         let mut state = self.state.lock().expect("failed to lock mutex");
-        state.scalar_functions.insert(f.name.clone(), Box::new(f));
-    }
-
-    /// Get a reference to the registered scalar functions
-    pub fn scalar_functions(&self) -> Box<HashMap<String, Box<ScalarFunction>>> {

Review comment:
       This function actually deep-copied the HashMap. After this PR it does not

##########
File path: rust/datafusion/src/execution/context.rs
##########
@@ -317,22 +309,9 @@ impl ExecutionContext {
 
     /// Optimize the logical plan by applying optimizer rules
     pub fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
-        let rules: Vec<Box<dyn OptimizerRule>> = vec![
-            Box::new(ProjectionPushDown::new()),
-            Box::new(FilterPushDown::new()),
-            Box::new(TypeCoercionRule::new(
-                self.state
-                    .lock()
-                    .expect("failed to lock mutex")
-                    .scalar_functions
-                    .as_ref(),
-            )),
-        ];
-        let mut plan = plan.clone();
-
-        for mut rule in rules {
-            plan = rule.optimize(&plan)?;
-        }
+        let plan = ProjectionPushDown::new().optimize(&plan)?;
+        let plan = FilterPushDown::new().optimize(&plan)?;
+        let plan = TypeCoercionRule::new(self).optimize(&plan)?;

Review comment:
       Now, this code makes it clear that the TypeCoercisonRule optimizer pass is a short lived object (and can't possibly outlive the `ExecutionContext`)

##########
File path: rust/datafusion/src/execution/context.rs
##########
@@ -427,6 +406,16 @@ impl ExecutionContext {
     }
 }
 
+impl ScalarFunctionProvider for ExecutionContext {
+    fn lookup(&self, name: &str) -> Option<Arc<ScalarFunction>> {
+        self.state

Review comment:
       My long term plan is to remove the Mutex entirely from `ExecutionContext` and thus this can be made simpler

##########
File path: rust/datafusion/src/optimizer/type_coercion.rs
##########
@@ -38,17 +35,21 @@ use utils::optimize_explain;
 /// Optimizer that applies coercion rules to expressions in the logical plan.
 ///
 /// This optimizer does not alter the structure of the plan, it only changes expressions on it.
-pub struct TypeCoercionRule {
-    scalar_functions: Arc<HashMap<String, Box<ScalarFunction>>>,
+pub struct TypeCoercionRule<'a, P>
+where
+    P: ScalarFunctionProvider,
+{
+    scalar_functions: &'a P,
 }
 
-impl TypeCoercionRule {
+impl<'a, P> TypeCoercionRule<'a, P>
+where
+    P: ScalarFunctionProvider,
+{
     /// Create a new type coercion optimizer rule using meta-data about registered
     /// scalar functions
-    pub fn new(scalar_functions: &HashMap<String, Box<ScalarFunction>>) -> Self {
-        Self {
-            scalar_functions: Arc::new(scalar_functions.clone()),

Review comment:
       This used to copy the entire HashMap




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org