Dynamic partition pruning

dask_planner/src/sql.rs

@@ -546,13 +546,22 @@ impl DaskSQLContext {
                         warn!("This LogicalPlan does not support Optimization. Returning original");
                         Ok(existing_plan)
                     }
-                    _ => optimizer::DaskSqlOptimizer::new()
-                        .optimize(existing_plan.original_plan)
-                        .map(|k| PyLogicalPlan {
-                            original_plan: k,
-                            current_node: None,
-                        })
-                        .map_err(py_optimization_exp),
+                    _ => {
+                        let optimized_plan = optimizer::DaskSqlOptimizer::new()
+                            .optimize(existing_plan.original_plan)
+                            .map(|k| PyLogicalPlan {
+                                original_plan: k,
+                                current_node: None,
+                            })
+                            .map_err(py_optimization_exp);
+                        optimizer::DaskSqlOptimizer::dynamic_partition_pruner()
+                            .optimize_once(optimized_plan.unwrap().original_plan)
+                            .map(|k| PyLogicalPlan {
+                                original_plan: k,
+                                current_node: None,
+                            })
+                            .map_err(py_optimization_exp)
+                    }
                 }
             }
             Err(e) => Err(py_optimization_exp(e)),

dask_planner/src/sql/optimizer.rs

@@ -24,6 +24,9 @@ use datafusion_python::{
 };
 use log::{debug, trace};
 
+mod dynamic_partition_pruning;
+use dynamic_partition_pruning::DynamicPartitionPruning;
+
 mod join_reorder;
 use join_reorder::JoinReorder;
 
@@ -86,13 +89,32 @@ impl DaskSqlOptimizer {
         }
     }
 
+    // Create a separate instance of this optimization rule, since we want to ensure that it only
+    // runs one time
+    pub fn dynamic_partition_pruner() -> Self {
+        let rule: Vec<Arc<dyn OptimizerRule + Sync + Send>> =
+            vec![Arc::new(DynamicPartitionPruning::new())];
+
+        Self {
+            optimizer: Optimizer::with_rules(rule),
+        }
+    }
+
     /// Iterates through the configured `OptimizerRule`(s) to transform the input `LogicalPlan`
     /// to its final optimized form
     pub(crate) fn optimize(&self, plan: LogicalPlan) -> Result<LogicalPlan, DataFusionError> {
         let config = OptimizerContext::new();
         self.optimizer.optimize(&plan, &config, Self::observe)
     }
 
+    /// Iterates once through the configured `OptimizerRule`(s) to transform the input `LogicalPlan`
+    /// to its final optimized form
+    pub(crate) fn optimize_once(&self, plan: LogicalPlan) -> Result<LogicalPlan, DataFusionError> {
+        let mut config = OptimizerContext::new();
+        config = OptimizerContext::with_max_passes(config, 1);
+        self.optimizer.optimize(&plan, &config, Self::observe)
+    }
+
     fn observe(optimized_plan: &LogicalPlan, optimization: &dyn OptimizerRule) {
         trace!(
             "== AFTER APPLYING RULE {} ==\n{}\n",