WIP: FFI query planner

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;
@@ -37,7 +38,7 @@ use datafusion::physical_planner::{
     DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner,
 };
 use datafusion::prelude::*;
-use datafusion_common::HashMap;
+use datafusion_common::{HashMap, exec_datafusion_err};
 use datafusion_examples::utils::{datasets::ExampleDataset, write_csv_to_parquet};
 
 /// This example demonstrates how to leverage [CacheFactory] to implement custom caching strategies for dataframes in DataFusion.
@@ -198,8 +199,16 @@ impl QueryPlanner for CacheNodeQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session: &dyn Session,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let session_state =
+            session
+                .as_any()
+                .downcast_ref::<SessionState>()
+                .ok_or_else(|| {
+                    exec_datafusion_err!("Failed to downcast Session to SessionState")
+                })?;
+
         let physical_planner =
             DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
                 CacheNodePlanner {

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -101,7 +101,9 @@ use futures::{
 use rand::{Rng, SeedableRng, rngs::StdRng};
 use tonic::async_trait;
 
-use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
+use datafusion::{
+    catalog::Session, optimizer::simplify_expressions::simplify_literal::parse_literal,
+};
 use datafusion::{
     execution::{
         RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
@@ -116,8 +118,8 @@ use datafusion::{
     prelude::*,
 };
 use datafusion_common::{
-    DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
-    plan_datafusion_err, plan_err,
+    DFSchemaRef, DataFusionError, Result, Statistics, exec_datafusion_err, internal_err,
+    not_impl_err, plan_datafusion_err, plan_err,
 };
 use datafusion_expr::{
     UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
@@ -564,8 +566,16 @@ impl QueryPlanner for TableSampleQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session: &dyn Session,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let session_state =
+            session
+                .as_any()
+                .downcast_ref::<SessionState>()
+                .ok_or_else(|| {
+                    exec_datafusion_err!("Failed to downcast Session to SessionState")
+                })?;
+
         let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
             TableSampleExtensionPlanner,
         )]);

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

@@ -56,6 +56,7 @@ use crate::{
 
 // backwards compatibility
 pub use crate::execution::session_state::SessionState;
+pub use datafusion_session::QueryPlanner;
 
 use arrow::datatypes::{Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
@@ -1970,17 +1971,6 @@ impl From<SessionContext> for SessionStateBuilder {
     }
 }
 
-/// A planner used to add extensions to DataFusion logical and physical plans.
-#[async_trait]
-pub trait QueryPlanner: 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>>;
-}
-
 /// Interface for handling `CREATE FUNCTION` statements and interacting with
 /// [SessionState] to create and register functions ([`ScalarUDF`],
 /// [`AggregateUDF`], [`WindowUDF`], and [`TableFunctionImpl`]) dynamically.
@@ -2162,6 +2152,7 @@ mod tests {
     use crate::test_util::{plan_and_collect, populate_csv_partitions};
     use arrow::datatypes::{DataType, TimeUnit};
     use datafusion_common::DataFusionError;
+    use datafusion_session::Session;
     use std::error::Error;
     use std::path::PathBuf;
 
@@ -2643,8 +2634,14 @@ mod tests {
         async fn create_physical_plan(
             &self,
             logical_plan: &LogicalPlan,
-            session_state: &SessionState,
+            session: &dyn Session,
         ) -> Result<Arc<dyn ExecutionPlan>> {
+            let session_state = session
+                .as_any()
+                .downcast_ref::<SessionState>()
+                .ok_or_else(|| {
+                    exec_datafusion_err!("Failed to downcast Session to SessionState")
+                })?;
             let physical_planner = MyPhysicalPlanner {};
             physical_planner
                 .create_physical_plan(logical_plan, session_state)

datafusion/core/src/execution/session_state.rs

@@ -28,7 +28,7 @@ use crate::datasource::file_format::FileFormatFactory;
 #[cfg(feature = "sql")]
 use crate::datasource::provider_as_source;
 use crate::execution::SessionStateDefaults;
-use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner};
+use crate::execution::context::{EmptySerializerRegistry, FunctionFactory};
 use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
 use arrow_schema::{DataType, FieldRef};
 use datafusion_catalog::MemoryCatalogProviderList;
@@ -44,7 +44,7 @@ use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan};
 use datafusion_common::tree_node::TreeNode;
 use datafusion_common::{
     DFSchema, DataFusionError, ResolvedTableReference, TableReference, config_err,
-    exec_err, plan_datafusion_err,
+    exec_datafusion_err, exec_err, plan_datafusion_err,
 };
 use datafusion_execution::TaskContext;
 use datafusion_execution::config::SessionConfig;
@@ -68,7 +68,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion_physical_plan::ExecutionPlan;
-use datafusion_session::Session;
+use datafusion_session::{QueryPlanner, Session};
 #[cfg(feature = "sql")]
 use datafusion_sql::{
     parser::{DFParserBuilder, Statement},
@@ -2115,8 +2115,15 @@ impl QueryPlanner for DefaultQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session: &dyn Session,
     ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
+        let session_state =
+            session
+                .as_any()
+                .downcast_ref::<SessionState>()
+                .ok_or_else(|| {
+                    exec_datafusion_err!("Failed to downcast Session to SessionState")
+                })?;
         let planner = DefaultPhysicalPlanner::default();
         planner
             .create_physical_plan(logical_plan, session_state)

datafusion/core/tests/user_defined/user_defined_plan.rs

@@ -91,14 +91,17 @@ use datafusion::{
 };
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
-use datafusion_common::{ScalarValue, assert_eq_or_internal_err, assert_or_internal_err};
+use datafusion_common::{
+    ScalarValue, assert_eq_or_internal_err, assert_or_internal_err, exec_datafusion_err,
+};
 use datafusion_expr::{FetchType, InvariantLevel, Projection, SortExpr};
 use datafusion_optimizer::AnalyzerRule;
 use datafusion_optimizer::optimizer::ApplyOrder;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
 use async_trait::async_trait;
 use datafusion_common::cast::as_string_view_array;
+use datafusion_session::Session;
 use futures::{Stream, StreamExt};
 
 /// Execute the specified sql and return the resulting record batches
@@ -466,8 +469,16 @@ impl QueryPlanner for TopKQueryPlanner {
     async fn create_physical_plan(
         &self,
         logical_plan: &LogicalPlan,
-        session_state: &SessionState,
+        session: &dyn Session,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let session_state =
+            session
+                .as_any()
+                .downcast_ref::<SessionState>()
+                .ok_or_else(|| {
+                    exec_datafusion_err!("Failed to downcast Session to SessionState")
+                })?;
+
         // Teach the default physical planner how to plan TopK nodes.
         let physical_planner =
             DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(

datafusion/execution/src/task.rs

@@ -212,7 +212,7 @@ impl FunctionRegistry for TaskContext {
 }
 
 /// Produce the [`TaskContext`].
-pub trait TaskContextProvider {
+pub trait TaskContextProvider: Sync + Send {
     fn task_ctx(&self) -> Arc<TaskContext>;
 }
 

datafusion/ffi/src/session/mod.rs

@@ -38,7 +38,9 @@ use datafusion_expr::{
 };
 use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_plan::ExecutionPlan;
-use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes};
+use datafusion_proto::bytes::{
+    logical_plan_from_bytes, logical_plan_to_bytes_with_extension_codec,
+};
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::logical_plan::from_proto::parse_expr;
 use datafusion_proto::logical_plan::to_proto::serialize_expr;
@@ -60,6 +62,7 @@ use crate::util::FFIResult;
 use crate::{df_result, rresult, rresult_return};
 
 pub mod config;
+pub mod planner;
 
 /// A stable struct for sharing [`Session`] across FFI boundaries.
 ///
@@ -467,8 +470,10 @@ impl Session for ForeignSession {
         &self,
         logical_plan: &LogicalPlan,
     ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
+        let codec: Arc<dyn LogicalExtensionCodec> = (&self.session.logical_codec).into();
         unsafe {
-            let logical_plan = logical_plan_to_bytes(logical_plan)?;
+            let logical_plan =
+                logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?;
             let physical_plan = df_result!(
                 (self.session.create_physical_plan)(
                     &self.session,