Skip to content

Commit

Permalink
Minor: Add getter for logical optimizer rules (apache#12379)
Browse files Browse the repository at this point in the history
* feat: new getter method for optimizer rules in parity with its physical counterpart

* style: use rust_lint

* chore: make struct local to the test scope
  • Loading branch information
maronavenue authored Sep 8, 2024
1 parent 5e37bb9 commit 3e1850d
Showing 1 changed file with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions datafusion/core/src/execution/session_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -801,6 +801,11 @@ impl SessionState {
&mut self.config
}

/// Return the logical optimizers
pub fn optimizers(&self) -> &[Arc<dyn OptimizerRule + Send + Sync>] {
&self.optimizer.rules
}

/// Return the physical optimizers
pub fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
&self.physical_optimizers.rules
Expand Down Expand Up @@ -1826,6 +1831,8 @@ mod tests {
use datafusion_common::Result;
use datafusion_execution::config::SessionConfig;
use datafusion_expr::Expr;
use datafusion_optimizer::optimizer::OptimizerRule;
use datafusion_optimizer::Optimizer;
use datafusion_sql::planner::{PlannerContext, SqlToRel};
use std::collections::HashMap;
use std::sync::Arc;
Expand Down Expand Up @@ -1922,4 +1929,31 @@ mod tests {
assert!(new_state.catalog_list().catalog(&default_catalog).is_none());
Ok(())
}

#[test]
fn test_session_state_with_optimizer_rules() {
struct DummyRule {}

impl OptimizerRule for DummyRule {
fn name(&self) -> &str {
"dummy_rule"
}
}
// test building sessions with fresh set of rules
let state = SessionStateBuilder::new()
.with_optimizer_rules(vec![Arc::new(DummyRule {})])
.build();

assert_eq!(state.optimizers().len(), 1);

// test adding rules to default recommendations
let state = SessionStateBuilder::new()
.with_optimizer_rule(Arc::new(DummyRule {}))
.build();

assert_eq!(
state.optimizers().len(),
Optimizer::default().rules.len() + 1
);
}
}

0 comments on commit 3e1850d

Please sign in to comment.