refactor: move planning APIs to session crate

datafusion-examples/examples/dataframe/cache_factory.rs

@@ -23,6 +23,7 @@ use std::sync::{Arc, RwLock};
 
 use arrow::array::RecordBatch;
 use async_trait::async_trait;
+use datafusion::catalog::Session;
 use datafusion::catalog::memory::MemorySourceConfig;
 use datafusion::common::DFSchemaRef;
 use datafusion::error::Result;
@@ -146,7 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner {
         node: &dyn UserDefinedLogicalNode,
         logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
-        session_state: &SessionState,
+        session_state: &dyn Session,
         _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         if let Some(cache_node) = node.as_any().downcast_ref::<CacheNode>() {
@@ -200,7 +201,7 @@ impl QueryPlanner for CacheNodeQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session_state: &dyn Session,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let physical_planner =
             DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -102,9 +102,10 @@ use tonic::async_trait;
 
 use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
 use datafusion::{
+    catalog::Session,
     execution::{
-        RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
-        TaskContext, context::QueryPlanner,
+        RecordBatchStream, SendableRecordBatchStream, SessionStateBuilder, TaskContext,
+        context::QueryPlanner,
     },
     physical_expr::EquivalenceProperties,
     physical_plan::{
@@ -565,7 +566,7 @@ impl QueryPlanner for TableSampleQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session_state: &dyn Session,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
             TableSampleExtensionPlanner,
@@ -587,7 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner {
         node: &dyn UserDefinedLogicalNode,
         _logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
-        _session_state: &SessionState,
+        _session_state: &dyn Session,
         _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()

datafusion/core/src/datasource/listing_table_factory.rs

@@ -838,7 +838,10 @@ mod tests {
         use datafusion_execution::config::SessionConfig;
         use datafusion_physical_expr::PhysicalExpr;
         use datafusion_physical_plan::ExecutionPlan;
-        use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList};
+        use datafusion_session::{
+            CatalogProviderList, EmptyCatalogProviderList, QueryPlanner,
+            UnsupportedQueryPlanner,
+        };
         use std::any::Any;
         use std::collections::HashMap;
 
@@ -857,6 +860,9 @@ mod tests {
             fn catalog_list(&self) -> Arc<dyn CatalogProviderList> {
                 Arc::new(EmptyCatalogProviderList)
             }
+            fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+                Arc::new(UnsupportedQueryPlanner)
+            }
             async fn create_physical_plan(
                 &self,
                 _logical_plan: &datafusion_expr::LogicalPlan,

datafusion/core/src/execution/context/mod.rs

@@ -17,7 +17,6 @@
 
 //! [`SessionContext`] API for registering data sources and executing queries
 
-use std::any::Any;
 use std::collections::HashSet;
 use std::fmt::Debug;
 use std::sync::{Arc, Weak};
@@ -2187,16 +2186,9 @@ impl From<SessionContext> for SessionStateBuilder {
     }
 }
 
+// Re-export from this module for backwards compatibility.
 /// A planner used to add extensions to DataFusion logical and physical plans.
-#[async_trait]
-pub trait QueryPlanner: Any + Debug {
-    /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution
-    async fn create_physical_plan(
-        &self,
-        logical_plan: &LogicalPlan,
-        session_state: &SessionState,
-    ) -> Result<Arc<dyn ExecutionPlan>>;
-}
+pub use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner};
 
 /// Interface for handling `CREATE FUNCTION` statements and interacting with
 /// [SessionState] to create and register functions ([`ScalarUDF`],
@@ -2396,6 +2388,7 @@ mod tests {
     use crate::physical_planner::PhysicalPlanner;
     use async_trait::async_trait;
     use datafusion_expr::planner::TypePlanner;
+    use datafusion_session::Session;
     use sqlparser::ast;
     use tempfile::TempDir;
 
@@ -2842,7 +2835,7 @@ mod tests {
         async fn create_physical_plan(
             &self,
             _logical_plan: &LogicalPlan,
-            _session_state: &SessionState,
+            _session_state: &dyn Session,
         ) -> Result<Arc<dyn ExecutionPlan>> {
             not_impl_err!("query not supported")
         }
@@ -2851,7 +2844,7 @@ mod tests {
             &self,
             _expr: &Expr,
             _input_dfschema: &DFSchema,
-            _session_state: &SessionState,
+            _session_state: &dyn Session,
             _planning_ctx: &PhysicalPlanningContext,
         ) -> Result<Arc<dyn PhysicalExpr>> {
             unimplemented!()
@@ -2866,7 +2859,7 @@ mod tests {
         async fn create_physical_plan(
             &self,
             logical_plan: &LogicalPlan,
-            session_state: &SessionState,
+            session_state: &dyn Session,
         ) -> Result<Arc<dyn ExecutionPlan>> {
             let physical_planner = MyPhysicalPlanner {};
             physical_planner

datafusion/core/src/execution/session_state.rs

@@ -73,12 +73,10 @@ use datafusion_optimizer::{
 };
 use datafusion_physical_expr::create_physical_expr;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
-use datafusion_physical_optimizer::PhysicalOptimizerContext;
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
-use datafusion_session::Session;
+use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session};
 #[cfg(feature = "sql")]
 use datafusion_sql::{
     parser::{DFParserBuilder, Statement},
@@ -274,6 +272,22 @@ impl Session for SessionState {
         Arc::clone(self.catalog_list())
     }
 
+    fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+        Arc::clone(SessionState::query_planner(self))
+    }
+
+    fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result<LogicalPlan> {
+        SessionState::optimize(self, plan)
+    }
+
+    fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
+        SessionState::physical_optimizers(self)
+    }
+
+    fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+        SessionState::statistics_registry(self)
+    }
+
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
@@ -2322,7 +2336,7 @@ impl QueryPlanner for DefaultQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session_state: &dyn Session,
     ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
         let planner = DefaultPhysicalPlanner::default();
         planner