Skip to content

Commit d3ca9b0

Browse files
tustvoldtustvold
authored
Decouple physical optimizer from SessionConfig (#3887) (#4749)
* Decouple physical optimizer from SessionConfig (#3887) * Logical conflicts
1 parent b686b68 commit d3ca9b0

10 files changed

Lines changed: 57 additions & 52 deletions

File tree

‎datafusion/core/src/physical_optimizer/aggregate_statistics.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@
1818
//! Utilizing exact statistics from sources to avoid scanning data
1919
use std::sync::Arc;
2020

21+
use crate::config::ConfigOptions;
2122
use datafusion_expr::utils::COUNT_STAR_EXPANSION;
2223

23-
use crate::execution::context::SessionConfig;
2424
use crate::physical_plan::aggregates::{AggregateExec, AggregateMode};
2525
use crate::physical_plan::empty::EmptyExec;
2626
use crate::physical_plan::projection::ProjectionExec;
@@ -51,7 +51,7 @@ impl PhysicalOptimizerRule for AggregateStatistics {
5151
fn optimize(
5252
&self,
5353
plan: Arc<dyn ExecutionPlan>,
54-
config: &SessionConfig,
54+
config: &ConfigOptions,
5555
) -> Result<Arc<dyn ExecutionPlan>> {
5656
if let Some(partial_agg_exec) = take_optimizable(&*plan) {
5757
let partial_agg_exec = partial_agg_exec
@@ -307,9 +307,10 @@ mod tests {
307307
agg: TestAggregate,
308308
) -> Result<()> {
309309
let session_ctx = SessionContext::new();
310-
let conf = session_ctx.copied_config();
310+
let state = session_ctx.state();
311311
let plan = Arc::new(plan) as _;
312-
let optimized = AggregateStatistics::new().optimize(Arc::clone(&plan), &conf)?;
312+
let optimized = AggregateStatistics::new()
313+
.optimize(Arc::clone(&plan), state.config_options())?;
313314

314315
// A ProjectionExec is a sign that the count optimization was applied
315316
assert!(optimized.as_any().is::<ProjectionExec>());
@@ -548,7 +549,7 @@ mod tests {
548549
Arc::clone(&schema),
549550
)?;
550551

551-
let conf = SessionConfig::new();
552+
let conf = ConfigOptions::new();
552553
let optimized =
553554
AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;
554555

@@ -591,7 +592,7 @@ mod tests {
591592
Arc::clone(&schema),
592593
)?;
593594

594-
let conf = SessionConfig::new();
595+
let conf = ConfigOptions::new();
595596
let optimized =
596597
AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;
597598

‎datafusion/core/src/physical_optimizer/coalesce_batches.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
//! CoalesceBatches optimizer that groups batches together rows
1919
//! in bigger batches to avoid overhead with small batches
2020
21+
use crate::config::ConfigOptions;
2122
use crate::{
2223
error::Result,
2324
physical_optimizer::PhysicalOptimizerRule,
@@ -46,7 +47,7 @@ impl PhysicalOptimizerRule for CoalesceBatches {
4647
fn optimize(
4748
&self,
4849
plan: Arc<dyn crate::physical_plan::ExecutionPlan>,
49-
_config: &crate::execution::context::SessionConfig,
50+
_config: &ConfigOptions,
5051
) -> Result<Arc<dyn crate::physical_plan::ExecutionPlan>> {
5152
let target_batch_size = self.target_batch_size;
5253
plan.transform_up(&|plan| {

‎datafusion/core/src/physical_optimizer/enforcement.rs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@
1818
//! Enforcement optimizer rules are used to make sure the plan's Distribution and Ordering
1919
//! requirements are met by inserting necessary [[RepartitionExec]] and [[SortExec]].
2020
//!
21-
use crate::config::OPT_TOP_DOWN_JOIN_KEY_REORDERING;
21+
use crate::config::{
22+
ConfigOptions, OPT_TARGET_PARTITIONS, OPT_TOP_DOWN_JOIN_KEY_REORDERING,
23+
};
2224
use crate::error::Result;
2325
use crate::physical_optimizer::utils::{add_sort_above_child, ordering_satisfy};
2426
use crate::physical_optimizer::PhysicalOptimizerRule;
@@ -34,7 +36,6 @@ use crate::physical_plan::sorts::sort::SortOptions;
3436
use crate::physical_plan::windows::WindowAggExec;
3537
use crate::physical_plan::Partitioning;
3638
use crate::physical_plan::{with_new_children_if_necessary, Distribution, ExecutionPlan};
37-
use crate::prelude::SessionConfig;
3839
use arrow::datatypes::SchemaRef;
3940
use datafusion_expr::logical_plan::JoinType;
4041
use datafusion_physical_expr::equivalence::EquivalenceProperties;
@@ -69,11 +70,10 @@ impl PhysicalOptimizerRule for BasicEnforcement {
6970
fn optimize(
7071
&self,
7172
plan: Arc<dyn ExecutionPlan>,
72-
config: &SessionConfig,
73+
config: &ConfigOptions,
7374
) -> Result<Arc<dyn ExecutionPlan>> {
74-
let target_partitions = config.target_partitions();
75+
let target_partitions = config.get_usize(OPT_TARGET_PARTITIONS).unwrap();
7576
let top_down_join_key_reordering = config
76-
.config_options()
7777
.get_bool(OPT_TOP_DOWN_JOIN_KEY_REORDERING)
7878
.unwrap_or_default();
7979
let new_plan = if top_down_join_key_reordering {
@@ -1135,10 +1135,12 @@ mod tests {
11351135
($EXPECTED_LINES: expr, $PLAN: expr) => {
11361136
let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect();
11371137

1138+
let mut config = ConfigOptions::new();
1139+
config.set_usize(OPT_TARGET_PARTITIONS, 10);
1140+
11381141
// run optimizer
11391142
let optimizer = BasicEnforcement {};
1140-
let optimized = optimizer
1141-
.optimize($PLAN, &SessionConfig::new().with_target_partitions(10))?;
1143+
let optimized = optimizer.optimize($PLAN, &config)?;
11421144

11431145
// Now format correctly
11441146
let plan = displayable(optimized.as_ref()).indent().to_string();

‎datafusion/core/src/physical_optimizer/global_sort_selection.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,13 @@
1919
2020
use std::sync::Arc;
2121

22+
use crate::config::ConfigOptions;
2223
use crate::error::Result;
2324
use crate::physical_optimizer::PhysicalOptimizerRule;
2425
use crate::physical_plan::rewrite::TreeNodeRewritable;
2526
use crate::physical_plan::sorts::sort::SortExec;
2627
use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
2728
use crate::physical_plan::ExecutionPlan;
28-
use crate::prelude::SessionConfig;
2929

3030
/// Currently for a sort operator, if
3131
/// - there are more than one input partitions
@@ -48,7 +48,7 @@ impl PhysicalOptimizerRule for GlobalSortSelection {
4848
fn optimize(
4949
&self,
5050
plan: Arc<dyn ExecutionPlan>,
51-
_config: &SessionConfig,
51+
_config: &ConfigOptions,
5252
) -> Result<Arc<dyn ExecutionPlan>> {
5353
plan.transform_up(&|plan| {
5454
Ok(plan

‎datafusion/core/src/physical_optimizer/join_selection.rs‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@ use std::sync::Arc;
2020

2121
use arrow::datatypes::Schema;
2222

23-
use crate::config::OPT_HASH_JOIN_SINGLE_PARTITION_THRESHOLD;
24-
use crate::execution::context::SessionConfig;
23+
use crate::config::{ConfigOptions, OPT_HASH_JOIN_SINGLE_PARTITION_THRESHOLD};
2524
use crate::logical_expr::JoinType;
2625
use crate::physical_plan::expressions::Column;
2726
use crate::physical_plan::joins::{
@@ -211,10 +210,9 @@ impl PhysicalOptimizerRule for JoinSelection {
211210
fn optimize(
212211
&self,
213212
plan: Arc<dyn ExecutionPlan>,
214-
session_config: &SessionConfig,
213+
config: &ConfigOptions,
215214
) -> Result<Arc<dyn ExecutionPlan>> {
216-
let collect_left_threshold: usize = session_config
217-
.config_options()
215+
let collect_left_threshold: usize = config
218216
.get_u64(OPT_HASH_JOIN_SINGLE_PARTITION_THRESHOLD)
219217
.unwrap_or_default()
220218
.try_into()
@@ -508,7 +506,7 @@ mod tests {
508506
.unwrap();
509507

510508
let optimized_join = JoinSelection::new()
511-
.optimize(Arc::new(join), &SessionConfig::new())
509+
.optimize(Arc::new(join), &ConfigOptions::new())
512510
.unwrap();
513511

514512
let swapping_projection = optimized_join
@@ -556,7 +554,7 @@ mod tests {
556554
.unwrap();
557555

558556
let optimized_join = JoinSelection::new()
559-
.optimize(Arc::new(join), &SessionConfig::new())
557+
.optimize(Arc::new(join), &ConfigOptions::new())
560558
.unwrap();
561559

562560
let swapping_projection = optimized_join
@@ -609,7 +607,7 @@ mod tests {
609607
let original_schema = join.schema();
610608

611609
let optimized_join = JoinSelection::new()
612-
.optimize(Arc::new(join), &SessionConfig::new())
610+
.optimize(Arc::new(join), &ConfigOptions::new())
613611
.unwrap();
614612

615613
let swapped_join = optimized_join
@@ -638,7 +636,7 @@ mod tests {
638636
$EXPECTED_LINES.iter().map(|s| *s).collect::<Vec<&str>>();
639637

640638
let optimized = JoinSelection::new()
641-
.optimize(Arc::new($PLAN), &SessionConfig::new())
639+
.optimize(Arc::new($PLAN), &ConfigOptions::new())
642640
.unwrap();
643641

644642
let plan = displayable(optimized.as_ref()).indent().to_string();
@@ -725,7 +723,7 @@ mod tests {
725723
.unwrap();
726724

727725
let optimized_join = JoinSelection::new()
728-
.optimize(Arc::new(join), &SessionConfig::new())
726+
.optimize(Arc::new(join), &ConfigOptions::new())
729727
.unwrap();
730728

731729
let swapped_join = optimized_join
@@ -950,7 +948,7 @@ mod tests {
950948
.unwrap();
951949

952950
let optimized_join = JoinSelection::new()
953-
.optimize(Arc::new(join), &SessionConfig::new())
951+
.optimize(Arc::new(join), &ConfigOptions::new())
954952
.unwrap();
955953

956954
if !is_swapped {

‎datafusion/core/src/physical_optimizer/optimize_sorts.rs‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
//! " SortExec: [non_nullable_col@1 ASC]",
2626
//! in the physical plan. The first sort is unnecessary since its result is overwritten
2727
//! by another SortExec. Therefore, this rule removes it from the physical plan.
28+
use crate::config::ConfigOptions;
2829
use crate::error::Result;
2930
use crate::physical_optimizer::utils::{
3031
add_sort_above_child, ordering_satisfy, ordering_satisfy_concrete,
@@ -34,7 +35,6 @@ use crate::physical_plan::rewrite::TreeNodeRewritable;
3435
use crate::physical_plan::sorts::sort::SortExec;
3536
use crate::physical_plan::windows::WindowAggExec;
3637
use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
37-
use crate::prelude::SessionConfig;
3838
use arrow::datatypes::SchemaRef;
3939
use datafusion_common::{reverse_sort_options, DataFusionError};
4040
use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
@@ -122,7 +122,7 @@ impl PhysicalOptimizerRule for OptimizeSorts {
122122
fn optimize(
123123
&self,
124124
plan: Arc<dyn ExecutionPlan>,
125-
_config: &SessionConfig,
125+
_config: &ConfigOptions,
126126
) -> Result<Arc<dyn ExecutionPlan>> {
127127
// Execute a post-order traversal to adjust input key ordering:
128128
let plan_requirements = PlanWithCorrespondingSort::new(plan);
@@ -557,7 +557,7 @@ mod tests {
557557
#[tokio::test]
558558
async fn test_remove_unnecessary_sort() -> Result<()> {
559559
let session_ctx = SessionContext::new();
560-
let conf = session_ctx.copied_config();
560+
let state = session_ctx.state();
561561
let schema = create_test_schema()?;
562562
let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
563563
as Arc<dyn ExecutionPlan>;
@@ -589,7 +589,7 @@ mod tests {
589589
expected, actual
590590
);
591591
let optimized_physical_plan =
592-
OptimizeSorts::new().optimize(physical_plan, &conf)?;
592+
OptimizeSorts::new().optimize(physical_plan, state.config_options())?;
593593
let formatted = displayable(optimized_physical_plan.as_ref())
594594
.indent()
595595
.to_string();
@@ -608,7 +608,7 @@ mod tests {
608608
#[tokio::test]
609609
async fn test_remove_unnecessary_sort_window_multilayer() -> Result<()> {
610610
let session_ctx = SessionContext::new();
611-
let conf = session_ctx.copied_config();
611+
let state = session_ctx.state();
612612
let schema = create_test_schema()?;
613613
let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
614614
as Arc<dyn ExecutionPlan>;
@@ -690,7 +690,7 @@ mod tests {
690690
expected, actual
691691
);
692692
let optimized_physical_plan =
693-
OptimizeSorts::new().optimize(physical_plan, &conf)?;
693+
OptimizeSorts::new().optimize(physical_plan, state.config_options())?;
694694
let formatted = displayable(optimized_physical_plan.as_ref())
695695
.indent()
696696
.to_string();
@@ -715,7 +715,7 @@ mod tests {
715715
#[tokio::test]
716716
async fn test_add_required_sort() -> Result<()> {
717717
let session_ctx = SessionContext::new();
718-
let conf = session_ctx.copied_config();
718+
let state = session_ctx.state();
719719
let schema = create_test_schema()?;
720720
let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
721721
as Arc<dyn ExecutionPlan>;
@@ -736,7 +736,7 @@ mod tests {
736736
expected, actual
737737
);
738738
let optimized_physical_plan =
739-
OptimizeSorts::new().optimize(physical_plan, &conf)?;
739+
OptimizeSorts::new().optimize(physical_plan, state.config_options())?;
740740
let formatted = displayable(optimized_physical_plan.as_ref())
741741
.indent()
742742
.to_string();
@@ -760,7 +760,7 @@ mod tests {
760760
#[tokio::test]
761761
async fn test_remove_unnecessary_sort1() -> Result<()> {
762762
let session_ctx = SessionContext::new();
763-
let conf = session_ctx.copied_config();
763+
let state = session_ctx.state();
764764
let schema = create_test_schema()?;
765765
let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
766766
as Arc<dyn ExecutionPlan>;
@@ -803,7 +803,7 @@ mod tests {
803803
expected, actual
804804
);
805805
let optimized_physical_plan =
806-
OptimizeSorts::new().optimize(physical_plan, &conf)?;
806+
OptimizeSorts::new().optimize(physical_plan, state.config_options())?;
807807
let formatted = displayable(optimized_physical_plan.as_ref())
808808
.indent()
809809
.to_string();
@@ -827,7 +827,7 @@ mod tests {
827827
#[tokio::test]
828828
async fn test_change_wrong_sorting() -> Result<()> {
829829
let session_ctx = SessionContext::new();
830-
let conf = session_ctx.copied_config();
830+
let state = session_ctx.state();
831831
let schema = create_test_schema()?;
832832
let source = Arc::new(MemoryExec::try_new(&[], schema.clone(), None)?)
833833
as Arc<dyn ExecutionPlan>;
@@ -865,7 +865,7 @@ mod tests {
865865
expected, actual
866866
);
867867
let optimized_physical_plan =
868-
OptimizeSorts::new().optimize(physical_plan, &conf)?;
868+
OptimizeSorts::new().optimize(physical_plan, state.config_options())?;
869869
let formatted = displayable(optimized_physical_plan.as_ref())
870870
.indent()
871871
.to_string();

‎datafusion/core/src/physical_optimizer/optimizer.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,8 @@
1919
2020
use std::sync::Arc;
2121

22-
use crate::{
23-
error::Result, execution::context::SessionConfig, physical_plan::ExecutionPlan,
24-
};
22+
use crate::config::ConfigOptions;
23+
use crate::{error::Result, physical_plan::ExecutionPlan};
2524

2625
/// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
2726
/// computes the same results, but in a potentially more efficient
@@ -31,7 +30,7 @@ pub trait PhysicalOptimizerRule {
3130
fn optimize(
3231
&self,
3332
plan: Arc<dyn ExecutionPlan>,
34-
config: &SessionConfig,
33+
config: &ConfigOptions,
3534
) -> Result<Arc<dyn ExecutionPlan>>;
3635

3736
/// A human readable name for this optimizer rule

‎datafusion/core/src/physical_optimizer/repartition.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@
1919
use std::sync::Arc;
2020

2121
use super::optimizer::PhysicalOptimizerRule;
22+
use crate::config::{ConfigOptions, OPT_TARGET_PARTITIONS};
23+
use crate::error::Result;
2224
use crate::physical_plan::Partitioning::*;
2325
use crate::physical_plan::{
2426
repartition::RepartitionExec, with_new_children_if_necessary, ExecutionPlan,
2527
};
26-
use crate::{error::Result, execution::context::SessionConfig};
2728

2829
/// Optimizer that introduces repartition to introduce more
2930
/// parallelism in the plan
@@ -207,14 +208,15 @@ impl PhysicalOptimizerRule for Repartition {
207208
fn optimize(
208209
&self,
209210
plan: Arc<dyn ExecutionPlan>,
210-
config: &SessionConfig,
211+
config: &ConfigOptions,
211212
) -> Result<Arc<dyn ExecutionPlan>> {
213+
let target_partitions = config.get_usize(OPT_TARGET_PARTITIONS).unwrap();
212214
// Don't run optimizer if target_partitions == 1
213-
if config.target_partitions() == 1 {
215+
if target_partitions == 1 {
214216
Ok(plan)
215217
} else {
216218
optimize_partitions(
217-
config.target_partitions(),
219+
target_partitions,
218220
plan.clone(),
219221
plan.output_ordering().is_none(),
220222
false,
@@ -360,8 +362,10 @@ mod tests {
360362
($EXPECTED_LINES: expr, $PLAN: expr) => {
361363
let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect();
362364

365+
let mut config = ConfigOptions::new();
366+
config.set_usize(OPT_TARGET_PARTITIONS, 10);
367+
363368
// run optimizer
364-
let config = SessionConfig::new().with_target_partitions(10);
365369
let optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
366370
Arc::new(Repartition::new()),
367371
// The `BasicEnforcement` is an essential rule to be applied.

0 commit comments

Comments
 (0)