Rust optimizer improvements

dask_planner/src/lib.rs

@@ -30,6 +30,7 @@ fn rust(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<sql::function::DaskFunction>()?;
     m.add_class::<sql::table::DaskStatistics>()?;
     m.add_class::<sql::logical::PyLogicalPlan>()?;
+    m.add_class::<sql::DaskSQLOptimizerConfig>()?;
 
     // Exceptions
     m.add(

dask_planner/src/sql.rs

@@ -99,7 +99,37 @@ pub struct DaskSQLContext {
     current_schema: String,
     schemas: HashMap<String, schema::DaskSchema>,
     options: ConfigOptions,
+    optimizer_config: DaskSQLOptimizerConfig,
+}
+
+#[pyclass(name = "DaskSQLOptimizerConfig", module = "dask_planner", subclass)]
+#[derive(Debug, Clone)]
+pub struct DaskSQLOptimizerConfig {
     dynamic_partition_pruning: bool,
+    fact_dimension_ratio: Option<f64>,
+    max_fact_tables: Option<usize>,
+    preserve_user_order: Option<bool>,
+    filter_selectivity: Option<f64>,
+}
+
+#[pymethods]
+impl DaskSQLOptimizerConfig {
+    #[new]
+    pub fn new(
+        dynamic_partition_pruning: bool,
+        fact_dimension_ratio: Option<f64>,
+        max_fact_tables: Option<usize>,
+        preserve_user_order: Option<bool>,
+        filter_selectivity: Option<f64>,
+    ) -> Self {
+        Self {
+            dynamic_partition_pruning,
+            fact_dimension_ratio,
+            max_fact_tables,
+            preserve_user_order,
+            filter_selectivity,
+        }
+    }
 }
 
 impl ContextProvider for DaskSQLContext {
@@ -483,18 +513,22 @@ impl ContextProvider for DaskSQLContext {
 #[pymethods]
 impl DaskSQLContext {
     #[new]
-    pub fn new(default_catalog_name: &str, default_schema_name: &str) -> Self {
+    pub fn new(
+        default_catalog_name: &str,
+        default_schema_name: &str,
+        optimizer_config: DaskSQLOptimizerConfig,
+    ) -> Self {
         Self {
             current_catalog: default_catalog_name.to_owned(),
             current_schema: default_schema_name.to_owned(),
             schemas: HashMap::new(),
             options: ConfigOptions::new(),
-            dynamic_partition_pruning: false,
+            optimizer_config,
         }
     }
 
-    pub fn apply_dynamic_partition_pruning(&mut self, config: bool) -> PyResult<()> {
-        self.dynamic_partition_pruning = config;
+    pub fn set_optimizer_config(&mut self, config: DaskSQLOptimizerConfig) -> PyResult<()> {
+        self.optimizer_config = config;
         Ok(())
     }
 
@@ -585,21 +619,28 @@ impl DaskSQLContext {
                         Ok(existing_plan)
                     }
                     _ => {
-                        let optimized_plan = optimizer::DaskSqlOptimizer::new()
-                            .optimize(existing_plan.original_plan)
+                        let optimized_plan = optimizer::DaskSqlOptimizer::new(
+                            self.optimizer_config.fact_dimension_ratio,
+                            self.optimizer_config.max_fact_tables,
+                            self.optimizer_config.preserve_user_order,
+                            self.optimizer_config.filter_selectivity,
+                        )
+                        .optimize(existing_plan.original_plan)
+                        .map(|k| PyLogicalPlan {
+                            original_plan: k,
+                            current_node: None,
+                        })
+                        .map_err(py_optimization_exp);
+                        if self.optimizer_config.dynamic_partition_pruning {
+                            optimizer::DaskSqlOptimizer::dynamic_partition_pruner(
+                                self.optimizer_config.fact_dimension_ratio,
+                            )
+                            .optimize_once(optimized_plan.unwrap().original_plan)
                             .map(|k| PyLogicalPlan {
                                 original_plan: k,
                                 current_node: None,
                             })
-                            .map_err(py_optimization_exp);
-                        if self.dynamic_partition_pruning {
-                            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)
+                            .map_err(py_optimization_exp)
                         } else {
                             optimized_plan
                         }

dask_planner/src/sql/optimizer.rs

@@ -39,7 +39,12 @@ pub struct DaskSqlOptimizer {
 impl DaskSqlOptimizer {
     /// Creates a new instance of the DaskSqlOptimizer with all the DataFusion desired
     /// optimizers as well as any custom `OptimizerRule` trait impls that might be desired.
-    pub fn new() -> Self {
+    pub fn new(
+        fact_dimension_ratio: Option<f64>,
+        max_fact_tables: Option<usize>,
+        preserve_user_order: Option<bool>,
+        filter_selectivity: Option<f64>,
+    ) -> Self {
         debug!("Creating new instance of DaskSqlOptimizer");
 
         let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
@@ -72,7 +77,12 @@ impl DaskSqlOptimizer {
             Arc::new(PushDownFilter::new()),
             // Arc::new(SingleDistinctToGroupBy::new()),
             // Dask-SQL specific optimizations
-            Arc::new(JoinReorder::default()),
+            Arc::new(JoinReorder::new(
+                fact_dimension_ratio,
+                max_fact_tables,
+                preserve_user_order,
+                filter_selectivity,
+            )),
             // The previous optimizations added expressions and projections,
             // that might benefit from the following rules
             Arc::new(SimplifyExpressions::new()),
@@ -91,9 +101,13 @@ 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())];
+    pub fn dynamic_partition_pruner(fact_dimension_ratio: Option<f64>) -> Self {
+        let rule: Vec<Arc<dyn OptimizerRule + Sync + Send>>;
+        if let Some(f) = fact_dimension_ratio {
+            rule = vec![Arc::new(DynamicPartitionPruning::new(f))];
+        } else {
+            rule = vec![Arc::new(DynamicPartitionPruning::default())];
+        }
 
         Self {
             optimizer: Optimizer::with_rules(rule),
@@ -177,7 +191,7 @@ mod tests {
         let plan = sql_to_rel.sql_statement_to_plan(statement.clone()).unwrap();
 
         // optimize the logical plan
-        let optimizer = DaskSqlOptimizer::new();
+        let optimizer = DaskSqlOptimizer::new(None, None, None, None);
         optimizer.optimize(plan)
     }
 

dask_planner/src/sql/optimizer/dynamic_partition_pruning.rs

@@ -36,11 +36,24 @@ use log::warn;
 use crate::sql::table::DaskTableSource;
 
 // Optimizer rule for dynamic partition pruning
-pub struct DynamicPartitionPruning {}
+pub struct DynamicPartitionPruning {
+    /// Ratio of the size of the dimension tables to fact tables
+    fact_dimension_ratio: f64,
+}
 
 impl DynamicPartitionPruning {
-    pub fn new() -> Self {
-        Self {}
+    pub fn new(fact_dimension_ratio: f64) -> Self {
+        Self {
+            fact_dimension_ratio,
+        }
+    }
+}
+
+impl Default for DynamicPartitionPruning {
+    fn default() -> Self {
+        Self {
+            fact_dimension_ratio: 0.3,
+        }
     }
 }
 
@@ -106,9 +119,6 @@ impl OptimizerRule for DynamicPartitionPruning {
                         (left_table.unwrap(), right_table.unwrap());
                     let (left_field, right_field) = (left_field.unwrap(), right_field.unwrap());
 
-                    // TODO: Consider allowing the fact_dimension_ratio to be configured by the
-                    // user. See issue: https://github.com/dask-contrib/dask-sql/issues/1121
-                    let fact_dimension_ratio = 0.3;
                     let (mut left_filtered_table, mut right_filtered_table) = (None, None);
 
                     // Check if join uses an alias instead of the table name itself. Need to use
@@ -136,7 +146,7 @@ impl OptimizerRule for DynamicPartitionPruning {
                         .size
                         .unwrap_or(largest_size as usize) as f64
                         / largest_size
-                        < fact_dimension_ratio
+                        < self.fact_dimension_ratio
                     {
                         left_filtered_table =
                             read_table(left_table.clone(), left_field.clone(), tables.clone());
@@ -149,7 +159,7 @@ impl OptimizerRule for DynamicPartitionPruning {
                         .size
                         .unwrap_or(largest_size as usize) as f64
                         / largest_size
-                        < fact_dimension_ratio
+                        < self.fact_dimension_ratio
                     {
                         right_filtered_table =
                             read_table(right_table.clone(), right_field.clone(), tables.clone());

dask_planner/src/sql/optimizer/join_reorder.rs

@@ -13,25 +13,30 @@ use log::warn;
 use crate::sql::table::DaskTableSource;
 
 pub struct JoinReorder {
-    /// Maximum number of fact tables to allow in a join
-    max_fact_tables: usize,
     /// Ratio of the size of the dimension tables to fact tables
     fact_dimension_ratio: f64,
+    /// Maximum number of fact tables to allow in a join
+    max_fact_tables: usize,
     /// Whether to preserve user-defined order of unfiltered dimensions
     preserve_user_order: bool,
     /// Constant to use when determining the number of rows produced by a
     /// filtered relation
     filter_selectivity: f64,
 }
 
-impl Default for JoinReorder {
-    fn default() -> Self {
+impl JoinReorder {
+    pub fn new(
+        fact_dimension_ratio: Option<f64>,
+        max_fact_tables: Option<usize>,
+        preserve_user_order: Option<bool>,
+        filter_selectivity: Option<f64>,
+    ) -> Self {
         Self {
-            max_fact_tables: 2,
-            // FIXME: fact_dimension_ratio should be 0.3
-            fact_dimension_ratio: 0.7,
-            preserve_user_order: true,
-            filter_selectivity: 1.0,
+            // FIXME: Default value for fact_dimension_ratio should be 0.3, not 0.7
+            fact_dimension_ratio: fact_dimension_ratio.unwrap_or(0.7),
+            max_fact_tables: max_fact_tables.unwrap_or(2),
+            preserve_user_order: preserve_user_order.unwrap_or(true),
+            filter_selectivity: filter_selectivity.unwrap_or(1.0),
         }
     }
 }

dask_sql/context.py

@@ -13,6 +13,7 @@
 from dask_planner.rust import (
     DaskSchema,
     DaskSQLContext,
+    DaskSQLOptimizerConfig,
     DaskTable,
     DFOptimizationException,
     DFParsingException,
@@ -98,13 +99,20 @@ def __init__(self, logging_level=logging.INFO):
         # A started SQL server (useful for jupyter notebooks)
         self.sql_server = None
 
-        # Create the `DaskSQLContext` Rust context
-        self.context = DaskSQLContext(self.catalog_name, self.schema_name)
-        self.context.register_schema(self.schema_name, DaskSchema(self.schema_name))
+        # Create the `DaskSQLOptimizerConfig` Rust context
+        optimizer_config = DaskSQLOptimizerConfig(
+            dask_config.get("sql.dynamic_partition_pruning"),
+            dask_config.get("sql.fact_dimension_ratio"),
+            dask_config.get("sql.max_fact_tables"),
+            dask_config.get("sql.preserve_user_order"),
+            dask_config.get("sql.filter_selectivity"),
+        )
 
-        self.context.apply_dynamic_partition_pruning(
-            dask_config.get("sql.dynamic_partition_pruning")
+        # Create the `DaskSQLContext` Rust context
+        self.context = DaskSQLContext(
+            self.catalog_name, self.schema_name, optimizer_config
         )
+        self.context.register_schema(self.schema_name, DaskSchema(self.schema_name))
 
         # # Register any default plugins, if nothing was registered before.
         RelConverter.add_plugin_class(logical.DaskAggregatePlugin, replace=False)
@@ -542,11 +550,16 @@ def explain(
             :obj:`str`: a description of the created relational algebra.
 
         """
+        dynamic_partition_pruning = dask_config.get("sql.dynamic_partition_pruning")
+        if not dask_config.get("sql.verbose_optimizer"):
+            dask_config.set({"sql.dynamic_partition_pruning": False})
+
         if dataframes is not None:
             for df_name, df in dataframes.items():
                 self.create_table(df_name, df, gpu=gpu)
 
         _, rel_string = self._get_ral(sql)
+        dask_config.set({"sql.dynamic_partition_pruning": dynamic_partition_pruning})
         return rel_string
 
     def visualize(self, sql: str, filename="mydask.png") -> None:  # pragma: no cover
@@ -799,9 +812,15 @@ def _get_ral(self, sql):
         """Helper function to turn the sql query into a relational algebra and resulting column names"""
 
         logger.debug(f"Entering _get_ral('{sql}')")
-        self.context.apply_dynamic_partition_pruning(
-            dask_config.get("sql.dynamic_partition_pruning")
+
+        optimizer_config = DaskSQLOptimizerConfig(
+            dask_config.get("sql.dynamic_partition_pruning"),
+            dask_config.get("sql.fact_dimension_ratio"),
+            dask_config.get("sql.max_fact_tables"),
+            dask_config.get("sql.preserve_user_order"),
+            dask_config.get("sql.filter_selectivity"),
         )
+        self.context.set_optimizer_config(optimizer_config)
 
         # get the schema of what we currently have registered
         schemas = self._prepare_schemas()

dask_sql/sql-schema.yaml

@@ -69,6 +69,38 @@ properties:
         description: |
           Whether to apply the dynamic partition pruning optimizer rule.
 
+      verbose_optimizer:
+        type: boolean
+        description: |
+          The dynamic partition pruning optimizer rule can sometimes result in extremely long
+          c.explain() outputs which are not helpful to the user. Setting this option to true allows
+          the user to see the entire output, while setting it to false truncates the output.
+          Default is false.
+
+      fact_dimension_ratio:
+        type: [number, "null"]
+        description: |
+          Ratio of the size of the dimension tables to fact tables. Parameter for dynamic partition
+          pruning and join reorder optimizer rules.
+
+      max_fact_tables:
+        type: [integer, "null"]
+        description: |
+          Maximum number of fact tables to allow in a join. Parameter for join reorder optimizer
+          rule.
+
+      preserve_user_order:
+        type: [boolean, "null"]
+        description: |
+          Whether to preserve user-defined order of unfiltered dimensions. Parameter for join
+          reorder optimizer rule.
+
+      filter_selectivity:
+        type: [number, "null"]
+        description: |
+          Constant to use when determining the number of rows produced by a filtered relation.
+          Parameter for join reorder optimizer rule.
+
       sort:
         type: object
         properties:

dask_sql/sql.yaml

@@ -18,6 +18,16 @@ sql:
 
   dynamic_partition_pruning: True
 
+  verbose_optimizer: False
+
+  fact_dimension_ratio: null
+
+  max_fact_tables: null
+
+  preserve_user_order: null
+
+  filter_selectivity: null
+
   sort:
     topk-nelem-limit: 1000000