From 6db06ce508a27363155a2153168174b4e3a60882 Mon Sep 17 00:00:00 2001
From: Maron Montano <maron.montano@gmail.com>
Date: Sun, 8 Sep 2024 03:27:39 +0800
Subject: [PATCH 1/3] feat: new getter method for optimizer rules in parity
 with its physical counterpart

---
 .../core/src/execution/session_state.rs       | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index 675ac798bf4e..30b4aec18dfa 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -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
@@ -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;
@@ -1922,4 +1929,29 @@ mod tests {
         assert!(new_state.catalog_list().catalog(&default_catalog).is_none());
         Ok(())
     }
+
+    #[test]
+    fn test_session_state_with_optimizer_rules() {
+        // test building sessions with supplied rules
+
+        let state = SessionStateBuilder::new()
+            .with_optimizer_rules(vec![Arc::new(DummyRule {})])
+            .build();
+
+        assert_eq!(state.optimizers().len(), 1);
+        
+        let state = SessionStateBuilder::new()
+            .with_optimizer_rule(Arc::new(DummyRule {}))
+            .build();
+
+        assert_eq!(state.optimizers().len(), Optimizer::default().rules.len() + 1);
+    }
+
+    struct DummyRule {}
+    
+    impl OptimizerRule for DummyRule {
+        fn name(&self) -> &str {
+            "dummy_rule"
+        }
+    }
 }

From c0349a13fc8c76d3ce97dd30358884ceb77bc1e5 Mon Sep 17 00:00:00 2001
From: Maron Montano <maron.montano@gmail.com>
Date: Sun, 8 Sep 2024 03:36:47 +0800
Subject: [PATCH 2/3] style: use rust_lint

---
 datafusion/core/src/execution/session_state.rs | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index 30b4aec18dfa..286a57a8ddc6 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -1939,16 +1939,19 @@ mod tests {
             .build();
 
         assert_eq!(state.optimizers().len(), 1);
-        
+
         let state = SessionStateBuilder::new()
             .with_optimizer_rule(Arc::new(DummyRule {}))
             .build();
 
-        assert_eq!(state.optimizers().len(), Optimizer::default().rules.len() + 1);
+        assert_eq!(
+            state.optimizers().len(),
+            Optimizer::default().rules.len() + 1
+        );
     }
 
     struct DummyRule {}
-    
+
     impl OptimizerRule for DummyRule {
         fn name(&self) -> &str {
             "dummy_rule"

From bba3fe52b022c7a01e524cc138ee57d2f957f274 Mon Sep 17 00:00:00 2001
From: Maron Montano <maron.montano@gmail.com>
Date: Sun, 8 Sep 2024 11:03:16 +0800
Subject: [PATCH 3/3] chore: make struct local to the test scope

---
 datafusion/core/src/execution/session_state.rs | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index 286a57a8ddc6..5e8d22b6ccbc 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -1932,14 +1932,21 @@ mod tests {
 
     #[test]
     fn test_session_state_with_optimizer_rules() {
-        // test building sessions with supplied 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();
@@ -1949,12 +1956,4 @@ mod tests {
             Optimizer::default().rules.len() + 1
         );
     }
-
-    struct DummyRule {}
-
-    impl OptimizerRule for DummyRule {
-        fn name(&self) -> &str {
-            "dummy_rule"
-        }
-    }
 }