PhysicalOptimizerRule now takes SessionConfig instead of ConfigOptions as the second argument to "optimizer()"

Author: gabotechs

datafusion-testing

@@ -19,8 +19,8 @@ use arrow::array::RecordBatch;
 use arrow::datatypes::{DataType, Field, Schema};
 use bytes::{BufMut, BytesMut};
 use criterion::{criterion_group, criterion_main, Criterion};
-use datafusion::config::ConfigOptions;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_execution::config::SessionConfig;
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
@@ -88,7 +88,7 @@ async fn create_plan() -> Arc<dyn ExecutionPlan> {
 #[derive(Clone)]
 struct BenchmarkPlan {
     plan: Arc<dyn ExecutionPlan>,
-    config: ConfigOptions,
+    config: SessionConfig,
 }
 
 impl std::fmt::Display for BenchmarkPlan {
@@ -102,8 +102,8 @@ fn bench_push_down_filter(c: &mut Criterion) {
     let plan = tokio::runtime::Runtime::new()
         .unwrap()
         .block_on(create_plan());
-    let mut config = ConfigOptions::default();
-    config.execution.parquet.pushdown_filters = true;
+    let mut config = SessionConfig::default();
+    config.options_mut().execution.parquet.pushdown_filters = true;
     let plan = BenchmarkPlan { plan, config };
     let optimizer = FilterPushdown::new();
 

datafusion/core/benches/push_down_filter.rs

@@ -2118,7 +2118,7 @@ impl DefaultPhysicalPlanner {
         for optimizer in optimizers {
             let before_schema = new_plan.schema();
             new_plan = optimizer
-                .optimize(new_plan, session_state.config_options())
+                .optimize(new_plan, session_state.config())
                 .map_err(|e| {
                     DataFusionError::Context(optimizer.name().to_string(), Box::new(e))
                 })?;
@@ -2441,7 +2441,6 @@ mod tests {
     use arrow::array::{ArrayRef, DictionaryArray, Int32Array};
     use arrow::datatypes::{DataType, Field, Int32Type};
     use arrow_schema::SchemaRef;
-    use datafusion_common::config::ConfigOptions;
     use datafusion_common::{
         assert_contains, DFSchemaRef, TableReference, ToDFSchema as _,
     };
@@ -3545,7 +3544,7 @@ digraph {
         fn optimize(
             &self,
             plan: Arc<dyn ExecutionPlan>,
-            _config: &ConfigOptions,
+            _config: &SessionConfig,
         ) -> Result<Arc<dyn ExecutionPlan>> {
             Ok(plan)
         }

datafusion/core/src/physical_planner.rs

@@ -810,8 +810,7 @@ async fn query_yields(
     task_ctx: Arc<TaskContext>,
 ) -> Result<(), Box<dyn Error>> {
     // Run plan through EnsureCooperative
-    let optimized =
-        EnsureCooperative::new().optimize(plan, task_ctx.session_config().options())?;
+    let optimized = EnsureCooperative::new().optimize(plan, task_ctx.session_config())?;
 
     // Get the stream
     let stream = physical_plan::execute_stream(optimized, task_ctx)?;

datafusion/core/tests/execution/coop.rs

@@ -38,7 +38,6 @@ use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_expr::{col, lit, Expr};
 
 use datafusion::datasource::physical_plan::FileScanConfig;
-use datafusion_common::config::ConfigOptions;
 use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::filter::FilterExec;
@@ -56,8 +55,13 @@ async fn check_stats_precision_with_filter_pushdown() {
     let table = get_listing_table(&table_path, None, &opt).await;
 
     let (_, _, state) = get_cache_runtime_state();
-    let mut options: ConfigOptions = state.config().options().as_ref().clone();
-    options.execution.parquet.pushdown_filters = true;
+    let mut session_config =
+        SessionConfig::from(state.config().options().as_ref().clone());
+    session_config
+        .options_mut()
+        .execution
+        .parquet
+        .pushdown_filters = true;
 
     // Scan without filter, stats are exact
     let exec = table.scan(&state, None, &[], None).await.unwrap();
@@ -85,7 +89,7 @@ async fn check_stats_precision_with_filter_pushdown() {
             as Arc<dyn ExecutionPlan>;
 
     let optimized_exec = FilterPushdown::new()
-        .optimize(filtered_exec, &options)
+        .optimize(filtered_exec, &session_config)
         .unwrap();
 
     assert!(

datafusion/core/tests/parquet/file_statistics.rs

@@ -27,6 +27,7 @@ use datafusion::datasource::source::DataSourceExec;
 use datafusion_common::cast::as_int64_array;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::Result;
+use datafusion_execution::config::SessionConfig;
 use datafusion_execution::TaskContext;
 use datafusion_expr::Operator;
 use datafusion_physical_expr::expressions::{self, cast};
@@ -67,7 +68,7 @@ async fn assert_count_optim_success(
     let task_ctx = Arc::new(TaskContext::default());
     let plan: Arc<dyn ExecutionPlan> = Arc::new(plan);
 
-    let config = ConfigOptions::new();
+    let config = SessionConfig::new();
     let optimized = AggregateStatistics::new().optimize(Arc::clone(&plan), &config)?;
 
     // A ProjectionExec is a sign that the count optimization was applied
@@ -264,7 +265,7 @@ async fn test_count_inexact_stat() -> Result<()> {
         Arc::clone(&schema),
     )?;
 
-    let conf = ConfigOptions::new();
+    let conf = SessionConfig::new();
     let optimized = AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;
 
     // check that the original ExecutionPlan was not replaced
@@ -308,7 +309,7 @@ async fn test_count_with_nulls_inexact_stat() -> Result<()> {
         Arc::clone(&schema),
     )?;
 
-    let conf = ConfigOptions::new();
+    let conf = SessionConfig::new();
     let optimized = AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;
 
     // check that the original ExecutionPlan was not replaced

datafusion/core/tests/physical_optimizer/aggregate_statistics.rs

@@ -27,6 +27,7 @@ use crate::physical_optimizer::test_utils::parquet_exec;
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
+use datafusion_execution::config::SessionConfig;
 use datafusion_functions_aggregate::count::count_udaf;
 use datafusion_functions_aggregate::sum::sum_udaf;
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
@@ -47,7 +48,7 @@ macro_rules! assert_optimized {
     ($PLAN: expr, @ $EXPECTED_LINES: literal $(,)?) => {
         // run optimizer
         let optimizer = CombinePartialFinalAggregate {};
-        let config = ConfigOptions::new();
+        let config = SessionConfig::new();
         let optimized = optimizer.optimize($PLAN, &config)?;
         // Now format correctly
         let plan = displayable(optimized.as_ref()).indent(true).to_string();

datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs

@@ -404,41 +404,56 @@ const SORT_DISTRIB_DISTRIB: [Run; 3] =
 
 #[derive(Clone)]
 struct TestConfig {
-    config: ConfigOptions,
+    session_config: SessionConfig,
 }
 
 impl Default for TestConfig {
     fn default() -> Self {
-        Self {
-            config: test_suite_default_config_options(),
-        }
+        let config = test_suite_default_config_options();
+        let session_config = SessionConfig::from(config.clone());
+        Self { session_config }
     }
 }
 
 impl TestConfig {
     /// If preferred, will not repartition / resort data if it is already sorted.
     fn with_prefer_existing_sort(mut self) -> Self {
-        self.config.optimizer.prefer_existing_sort = true;
+        self.session_config
+            .options_mut()
+            .optimizer
+            .prefer_existing_sort = true;
         self
     }
 
     /// If preferred, will not attempt to convert Union to Interleave.
     fn with_prefer_existing_union(mut self) -> Self {
-        self.config.optimizer.prefer_existing_union = true;
+        self.session_config
+            .options_mut()
+            .optimizer
+            .prefer_existing_union = true;
         self
     }
 
     /// If preferred, will repartition file scans.
     /// Accepts a minimum file size to repartition.
     fn with_prefer_repartition_file_scans(mut self, file_min_size: usize) -> Self {
-        self.config.optimizer.repartition_file_scans = true;
-        self.config.optimizer.repartition_file_min_size = file_min_size;
+        self.session_config
+            .options_mut()
+            .optimizer
+            .repartition_file_scans = true;
+        self.session_config
+            .options_mut()
+            .optimizer
+            .repartition_file_min_size = file_min_size;
         self
     }
 
     /// Set the preferred target partitions for query execution concurrency.
     fn with_query_execution_partitions(mut self, target_partitions: usize) -> Self {
-        self.config.execution.target_partitions = target_partitions;
+        self.session_config
+            .options_mut()
+            .execution
+            .target_partitions = target_partitions;
         self
     }
 
@@ -455,13 +470,18 @@ impl TestConfig {
 
         // Add the ancillary output requirements operator at the start:
         let optimizer = OutputRequirements::new_add_mode();
-        let mut optimized = optimizer.optimize(plan.clone(), &self.config)?;
+        let mut optimized = optimizer.optimize(plan.clone(), &self.session_config)?;
 
         // This file has 2 rules that use tree node, apply these rules to original plan consecutively
         // After these operations tree nodes should be in a consistent state.
         // This code block makes sure that these rules doesn't violate tree node integrity.
         {
-            let adjusted = if self.config.optimizer.top_down_join_key_reordering {
+            let adjusted = if self
+                .session_config
+                .options()
+                .optimizer
+                .top_down_join_key_reordering
+            {
                 // Run adjust_input_keys_ordering rule
                 let plan_requirements =
                     PlanWithKeyRequirements::new_default(plan.clone());
@@ -483,7 +503,10 @@ impl TestConfig {
             // Then run ensure_distribution rule
             DistributionContext::new_default(adjusted)
                 .transform_up(|distribution_context| {
-                    ensure_distribution(distribution_context, &self.config)
+                    ensure_distribution(
+                        distribution_context,
+                        &self.session_config.options(),
+                    )
                 })
                 .data()
                 .and_then(check_integrity)?;
@@ -494,18 +517,18 @@ impl TestConfig {
             optimized = match run {
                 Run::Distribution => {
                     let optimizer = EnforceDistribution::new();
-                    optimizer.optimize(optimized, &self.config)?
+                    optimizer.optimize(optimized, &self.session_config)?
                 }
                 Run::Sorting => {
                     let optimizer = EnforceSorting::new();
-                    optimizer.optimize(optimized, &self.config)?
+                    optimizer.optimize(optimized, &self.session_config)?
                 }
             };
         }
 
         // Remove the ancillary output requirements operator when done:
         let optimizer = OutputRequirements::new_remove_mode();
-        let optimized = optimizer.optimize(optimized, &self.config)?;
+        let optimized = optimizer.optimize(optimized, &self.session_config)?;
 
         // Now format correctly
         let actual_lines = get_plan_string(&optimized);
@@ -3340,10 +3363,13 @@ fn do_not_put_sort_when_input_is_invalid() -> Result<()> {
         "      DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
     ];
 
-    let mut config = ConfigOptions::new();
-    config.execution.target_partitions = 10;
-    config.optimizer.enable_round_robin_repartition = true;
-    config.optimizer.prefer_existing_sort = false;
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.target_partitions = 10;
+    config
+        .options_mut()
+        .optimizer
+        .enable_round_robin_repartition = true;
+    config.options_mut().optimizer.prefer_existing_sort = false;
     let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
     assert_plan_txt!(expected, dist_plan);
 
@@ -3378,10 +3404,13 @@ fn put_sort_when_input_is_valid() -> Result<()> {
         "    DataSourceExec: file_groups={10 groups: [[x:0..20], [y:0..20], [x:20..40], [y:20..40], [x:40..60], [y:40..60], [x:60..80], [y:60..80], [x:80..100], [y:80..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet",
     ];
 
-    let mut config = ConfigOptions::new();
-    config.execution.target_partitions = 10;
-    config.optimizer.enable_round_robin_repartition = true;
-    config.optimizer.prefer_existing_sort = false;
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.target_partitions = 10;
+    config
+        .options_mut()
+        .optimizer
+        .enable_round_robin_repartition = true;
+    config.options_mut().optimizer.prefer_existing_sort = false;
     let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
     assert_plan_txt!(expected, dist_plan);
 
@@ -3503,7 +3532,11 @@ async fn test_distribute_sort_parquet() -> Result<()> {
     let test_config: TestConfig =
         TestConfig::default().with_prefer_repartition_file_scans(1000);
     assert!(
-        test_config.config.optimizer.repartition_file_scans,
+        test_config
+            .session_config
+            .options()
+            .optimizer
+            .repartition_file_scans,
         "should enable scans to be repartitioned"
     );
 
@@ -3542,7 +3575,11 @@ async fn test_distribute_sort_memtable() -> Result<()> {
     let test_config: TestConfig =
         TestConfig::default().with_prefer_repartition_file_scans(1000);
     assert!(
-        test_config.config.optimizer.repartition_file_scans,
+        test_config
+            .session_config
+            .options()
+            .optimizer
+            .repartition_file_scans,
         "should enable scans to be repartitioned"
     );
 

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -34,6 +34,7 @@ use arrow::datatypes::{DataType, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{TreeNode, TransformedResult};
 use datafusion_common::{Result,  TableReference};
+use datafusion_execution::config::SessionConfig;
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr_common::operator::Operator;
@@ -110,8 +111,9 @@ impl EnforceSortingTest {
     /// Runs the enforce sorting test and returns a string with the input and
     /// optimized plan as strings for snapshot comparison using insta
     pub(crate) fn run(&self) -> String {
-        let mut config = ConfigOptions::new();
-        config.optimizer.repartition_sorts = self.repartition_sorts;
+        let mut session_config = SessionConfig::new();
+        session_config.options_mut().optimizer.repartition_sorts = self.repartition_sorts;
+        let config = session_config.options();
 
         // This file has 4 rules that use tree node, apply these rules as in the
         // EnforceSorting::optimize implementation
@@ -171,7 +173,7 @@ impl EnforceSortingTest {
 
         // Run the actual optimizer
         let optimized_physical_plan = EnforceSorting::new()
-            .optimize(Arc::clone(&self.plan), &config)
+            .optimize(Arc::clone(&self.plan), &session_config)
             .expect("enforce_sorting failed");
 
         // Get string representation of the plan
@@ -2277,7 +2279,7 @@ async fn test_commutativity() -> Result<()> {
           DataSourceExec: partitions=1, partition_sizes=[0]
     "#);
 
-    let config = ConfigOptions::new();
+    let config = SessionConfig::new();
     let rules = vec![
         Arc::new(EnforceDistribution::new()) as Arc<dyn PhysicalOptimizerRule>,
         Arc::new(EnforceSorting::new()) as Arc<dyn PhysicalOptimizerRule>,