implement rewrite for FilterNullJoinKeys

Author: Lordworms

datafusion/optimizer/src/filter_null_join_keys.rs

@@ -19,10 +19,12 @@
 
 use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result};
 use datafusion_expr::{
     and, logical_plan::Filter, logical_plan::JoinType, Expr, ExprSchemable, LogicalPlan,
 };
+use std::mem;
 use std::sync::Arc;
 
 /// The FilterNullJoinKeys rule will identify inner joins with equi-join conditions
@@ -32,66 +34,69 @@ use std::sync::Arc;
 #[derive(Default)]
 pub struct FilterNullJoinKeys {}
 
-impl FilterNullJoinKeys {
-    pub const NAME: &'static str = "filter_null_join_keys";
-}
-
 impl OptimizerRule for FilterNullJoinKeys {
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
-        config: &dyn OptimizerConfig,
+        _plan: &LogicalPlan,
+        _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
+        internal_err!("Should have called FilterNullJoinKeys::rewrite")
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::BottomUp)
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
         if !config.options().optimizer.filter_null_join_keys {
-            return Ok(None);
+            return Ok(Transformed::no(plan));
         }
 
         match plan {
-            LogicalPlan::Join(join) if join.join_type == JoinType::Inner => {
-                let mut join = join.clone();
-
+            LogicalPlan::Join(mut join) if join.join_type == JoinType::Inner => {
                 let left_schema = join.left.schema();
                 let right_schema = join.right.schema();
 
                 let mut left_filters = vec![];
                 let mut right_filters = vec![];
 
-                for (l, r) in &join.on {
+                for (l, r) in &mut join.on {
                     if l.nullable(left_schema)? {
-                        left_filters.push(l.clone());
+                        left_filters.push(mem::take(l));
                     }
 
                     if r.nullable(right_schema)? {
-                        right_filters.push(r.clone());
+                        right_filters.push(mem::take(r));
                     }
                 }
 
                 if !left_filters.is_empty() {
                     let predicate = create_not_null_predicate(left_filters);
                     join.left = Arc::new(LogicalPlan::Filter(Filter::try_new(
-                        predicate,
-                        join.left.clone(),
+                        predicate, join.left,
                     )?));
                 }
                 if !right_filters.is_empty() {
                     let predicate = create_not_null_predicate(right_filters);
                     join.right = Arc::new(LogicalPlan::Filter(Filter::try_new(
-                        predicate,
-                        join.right.clone(),
+                        predicate, join.right,
                     )?));
                 }
-                Ok(Some(LogicalPlan::Join(join)))
+                Ok(Transformed::yes(LogicalPlan::Join(join)))
             }
-            _ => Ok(None),
+            _ => Ok(Transformed::no(plan)),
         }
     }
-
     fn name(&self) -> &str {
-        Self::NAME
-    }
-
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::BottomUp)
+        "filter_null_join_keys"
     }
 }