Skip to content

Commit 3e1850d

Browse files
authored
Minor: Add getter for logical optimizer rules (#12379)
* 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
1 parent 5e37bb9 commit 3e1850d

File tree

1 file changed

+34
-0
lines changed

1 file changed

+34
-0
lines changed

datafusion/core/src/execution/session_state.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -801,6 +801,11 @@ impl SessionState {
801801
&mut self.config
802802
}
803803

804+
/// Return the logical optimizers
805+
pub fn optimizers(&self) -> &[Arc<dyn OptimizerRule + Send + Sync>] {
806+
&self.optimizer.rules
807+
}
808+
804809
/// Return the physical optimizers
805810
pub fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
806811
&self.physical_optimizers.rules
@@ -1826,6 +1831,8 @@ mod tests {
18261831
use datafusion_common::Result;
18271832
use datafusion_execution::config::SessionConfig;
18281833
use datafusion_expr::Expr;
1834+
use datafusion_optimizer::optimizer::OptimizerRule;
1835+
use datafusion_optimizer::Optimizer;
18291836
use datafusion_sql::planner::{PlannerContext, SqlToRel};
18301837
use std::collections::HashMap;
18311838
use std::sync::Arc;
@@ -1922,4 +1929,31 @@ mod tests {
19221929
assert!(new_state.catalog_list().catalog(&default_catalog).is_none());
19231930
Ok(())
19241931
}
1932+
1933+
#[test]
1934+
fn test_session_state_with_optimizer_rules() {
1935+
struct DummyRule {}
1936+
1937+
impl OptimizerRule for DummyRule {
1938+
fn name(&self) -> &str {
1939+
"dummy_rule"
1940+
}
1941+
}
1942+
// test building sessions with fresh set of rules
1943+
let state = SessionStateBuilder::new()
1944+
.with_optimizer_rules(vec![Arc::new(DummyRule {})])
1945+
.build();
1946+
1947+
assert_eq!(state.optimizers().len(), 1);
1948+
1949+
// test adding rules to default recommendations
1950+
let state = SessionStateBuilder::new()
1951+
.with_optimizer_rule(Arc::new(DummyRule {}))
1952+
.build();
1953+
1954+
assert_eq!(
1955+
state.optimizers().len(),
1956+
Optimizer::default().rules.len() + 1
1957+
);
1958+
}
19251959
}

0 commit comments

Comments
 (0)