refactor: pass the TaskContext when execute physical plan rather than holding it

interpreters/src/tests.rs

@@ -22,10 +22,9 @@ use catalog::{
 };
 use catalog_impls::table_based::TableBasedManager;
 use common_types::request_id::RequestId;
-use datafusion::execution::runtime_env::RuntimeEnv;
-use query_engine::{
-    datafusion_impl::physical_planner::DatafusionPhysicalPlannerImpl, executor::ExecutorImpl,
-};
+use datafusion::execution::runtime_env::RuntimeConfig;
+use df_operator::registry::{FunctionRegistry, FunctionRegistryImpl};
+use query_engine::{datafusion_impl::DatafusionQueryEngineImpl, QueryEngineRef};
 use query_frontend::{
     parser::Parser, plan::Plan, planner::Planner, provider::MetaProvider, tests::MockMetaProvider,
 };
@@ -60,6 +59,7 @@ where
     pub meta_provider: M,
     pub catalog_manager: ManagerRef,
     pub table_manipulator: TableManipulatorRef,
+    pub query_engine: QueryEngineRef,
 }
 
 impl<M> Env<M>
@@ -77,11 +77,8 @@ where
 {
     async fn build_factory(&self) -> Factory {
         Factory::new(
-            Arc::new(ExecutorImpl),
-            Arc::new(DatafusionPhysicalPlannerImpl::new(
-                query_engine::Config::default(),
-                Arc::new(RuntimeEnv::default()),
-            )),
+            self.query_engine.executor(),
+            self.query_engine.physical_planner(),
             self.catalog_manager.clone(),
             self.engine(),
             self.table_manipulator.clone(),
@@ -232,11 +229,8 @@ where
         let table_operator = TableOperator::new(catalog_manager.clone());
         let table_manipulator = Arc::new(TableManipulatorImpl::new(table_operator));
         let insert_factory = Factory::new(
-            Arc::new(ExecutorImpl),
-            Arc::new(DatafusionPhysicalPlannerImpl::new(
-                query_engine::Config::default(),
-                Arc::new(RuntimeEnv::default()),
-            )),
+            self.query_engine.executor(),
+            self.query_engine.physical_planner(),
             catalog_manager.clone(),
             self.engine(),
             table_manipulator.clone(),
@@ -255,11 +249,8 @@ where
         let select_sql =
             "SELECT key1, key2, field1, field2, field3, field4, field5 from test_missing_columns_table";
         let select_factory = Factory::new(
-            Arc::new(ExecutorImpl),
-            Arc::new(DatafusionPhysicalPlannerImpl::new(
-                query_engine::Config::default(),
-                Arc::new(RuntimeEnv::default()),
-            )),
+            self.query_engine.executor(),
+            self.query_engine.physical_planner(),
             catalog_manager,
             self.engine(),
             table_manipulator,
@@ -379,12 +370,22 @@ async fn test_interpreters<T: EngineBuildContext>(engine_context: T) {
     let catalog_manager = Arc::new(build_catalog_manager(engine.clone()).await);
     let table_operator = TableOperator::new(catalog_manager.clone());
     let table_manipulator = Arc::new(TableManipulatorImpl::new(table_operator));
+    let function_registry = Arc::new(FunctionRegistryImpl::default());
+    let query_engine = Box::new(
+        DatafusionQueryEngineImpl::new(
+            query_engine::Config::default(),
+            RuntimeConfig::default(),
+            function_registry.to_df_function_registry(),
+        )
+        .unwrap(),
+    );
 
     let env = Env {
         engine: test_ctx.clone_engine(),
         meta_provider: mock,
         catalog_manager,
         table_manipulator,
+        query_engine,
     };
 
     env.test_create_table().await;

query_engine/src/datafusion_impl/codec.rs

@@ -28,6 +28,7 @@ pub struct DataFusionPhysicalPlanEncoderImpl {
 }
 
 impl DataFusionPhysicalPlanEncoderImpl {
+    // TODO: use `SessionContext` to init it.
     pub fn new(
         runtime_env: Arc<RuntimeEnv>,
         function_registry: Arc<dyn FunctionRegistry + Send + Sync>,

query_engine/src/datafusion_impl/executor.rs

@@ -0,0 +1,78 @@
+// Copyright 2023 The CeresDB Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{sync::Arc, time::Instant};
+
+use async_trait::async_trait;
+use generic_error::BoxError;
+use log::{debug, info};
+use snafu::ResultExt;
+use table_engine::stream::SendableRecordBatchStream;
+use time_ext::InstantExt;
+
+use crate::{
+    context::Context,
+    datafusion_impl::DfContextBuilder,
+    error::*,
+    executor::Executor,
+    physical_planner::{PhysicalPlanPtr, TaskContext},
+};
+
+#[derive(Debug, Clone)]
+pub struct DatafusionExecutorImpl {
+    // Datafuison session context builder
+    df_ctx_builder: Arc<DfContextBuilder>,
+}
+
+impl DatafusionExecutorImpl {
+    pub fn new(df_ctx_builder: Arc<DfContextBuilder>) -> Self {
+        Self { df_ctx_builder }
+    }
+}
+
+#[async_trait]
+impl Executor for DatafusionExecutorImpl {
+    async fn execute(
+        &self,
+        ctx: &Context,
+        physical_plan: PhysicalPlanPtr,
+    ) -> Result<SendableRecordBatchStream> {
+        debug!(
+            "DatafusionExecutorImpl begin to execute plan, request_id:{}, physical_plan: {:?}",
+            ctx.request_id, physical_plan
+        );
+
+        let begin_instant = Instant::now();
+
+        // TODO: build the `TaskContext` directly rather than through `SessionContext`.
+        let session_ctx = self.df_ctx_builder.build(ctx);
+        let df_task_ctx = session_ctx.task_ctx();
+        let task_ctx = TaskContext::default().with_datafusion_task_ctx(df_task_ctx);
+        let stream = physical_plan
+            .execute(&task_ctx)
+            .box_err()
+            .with_context(|| ExecutorWithCause {
+                msg: Some("failed to execute physical plan".to_string()),
+            })?;
+
+        info!(
+            "DatafusionExecutorImpl finish to execute plan, request_id:{}, cost:{}ms, plan_and_metrics: {}",
+            ctx.request_id,
+            begin_instant.saturating_elapsed().as_millis(),
+            physical_plan.metrics_to_string()
+        );
+
+        Ok(stream)
+    }
+}

query_engine/src/datafusion_impl/mod.rs

@@ -12,24 +12,36 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
+use std::{fmt, sync::Arc, time::Instant};
 
-use datafusion::execution::{
-    runtime_env::{RuntimeConfig, RuntimeEnv},
-    FunctionRegistry,
+use datafusion::{
+    execution::{
+        context::{QueryPlanner, SessionState},
+        runtime_env::{RuntimeConfig, RuntimeEnv},
+        FunctionRegistry,
+    },
+    optimizer::analyzer::Analyzer,
+    physical_optimizer::PhysicalOptimizerRule,
+    prelude::{SessionConfig, SessionContext},
 };
+use table_engine::provider::CeresdbOptions;
 
 use crate::{
     codec::PhysicalPlanCodecRef,
+    context::Context,
     datafusion_impl::{
-        codec::DataFusionPhysicalPlanEncoderImpl, physical_planner::DatafusionPhysicalPlannerImpl,
+        codec::DataFusionPhysicalPlanEncoderImpl, executor::DatafusionExecutorImpl,
+        logical_optimizer::type_conversion::TypeConversion,
+        physical_planner::DatafusionPhysicalPlannerImpl,
+        physical_planner_extension::QueryPlannerAdapter,
     },
-    executor::{ExecutorImpl, ExecutorRef},
+    executor::ExecutorRef,
     physical_planner::PhysicalPlannerRef,
     Config, QueryEngine,
 };
 
 pub mod codec;
+pub mod executor;
 pub mod logical_optimizer;
 pub mod physical_optimizer;
 pub mod physical_plan;
@@ -53,12 +65,15 @@ impl DatafusionQueryEngineImpl {
         function_registry: Arc<dyn FunctionRegistry + Send + Sync>,
     ) -> Result<Self> {
         let runtime_env = Arc::new(RuntimeEnv::new(runtime_config).unwrap());
-
-        let physical_planner = Arc::new(DatafusionPhysicalPlannerImpl::new(
+        let physical_planner = Arc::new(QueryPlannerAdapter);
+        let df_ctx_builder = Arc::new(DfContextBuilder::new(
             config,
             runtime_env.clone(),
+            physical_planner,
         ));
-        let executor = Arc::new(ExecutorImpl);
+
+        let physical_planner = Arc::new(DatafusionPhysicalPlannerImpl::new(df_ctx_builder.clone()));
+        let executor = Arc::new(DatafusionExecutorImpl::new(df_ctx_builder));
         let physical_plan_codec = Arc::new(DataFusionPhysicalPlanEncoderImpl::new(
             runtime_env,
             function_registry,
@@ -85,3 +100,94 @@ impl QueryEngine for DatafusionQueryEngineImpl {
         self.physical_plan_codec.clone()
     }
 }
+
+/// Datafusion context builder
+#[derive(Clone)]
+pub struct DfContextBuilder {
+    config: Config,
+    runtime_env: Arc<RuntimeEnv>,
+    physical_planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+impl fmt::Debug for DfContextBuilder {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("DfContextBuilder")
+            .field("config", &self.config)
+            .field("runtime_env", &self.runtime_env)
+            .field("physical_planner", &"QueryPlannerAdapter")
+            .finish()
+    }
+}
+
+impl DfContextBuilder {
+    pub fn new(
+        config: Config,
+        runtime_env: Arc<RuntimeEnv>,
+        physical_planner: Arc<dyn QueryPlanner + Send + Sync>,
+    ) -> Self {
+        Self {
+            config,
+            runtime_env,
+            physical_planner,
+        }
+    }
+
+    pub fn build(&self, ctx: &Context) -> SessionContext {
+        let timeout = ctx
+            .deadline
+            .map(|deadline| deadline.duration_since(Instant::now()).as_millis() as u64);
+        let ceresdb_options = CeresdbOptions {
+            request_id: ctx.request_id.as_u64(),
+            request_timeout: timeout,
+        };
+        let mut df_session_config = SessionConfig::new()
+            .with_default_catalog_and_schema(
+                ctx.default_catalog.clone(),
+                ctx.default_schema.clone(),
+            )
+            .with_target_partitions(self.config.read_parallelism);
+
+        df_session_config
+            .options_mut()
+            .extensions
+            .insert(ceresdb_options);
+
+        // Using default logcial optimizer, if want to add more custom rule, using
+        // `add_optimizer_rule` to add.
+        let state = SessionState::with_config_rt(df_session_config, self.runtime_env.clone())
+            .with_query_planner(self.physical_planner.clone());
+
+        // Register analyzer rules
+        let state = Self::register_analyzer_rules(state);
+
+        // Register iox optimizers, used by influxql.
+        let state = influxql_query::logical_optimizer::register_iox_logical_optimizers(state);
+
+        SessionContext::with_state(state)
+    }
+
+    // TODO: this is not used now, bug of RepartitionAdapter is already fixed in
+    // datafusion itself. Remove this code in future.
+    #[allow(dead_code)]
+    fn apply_adapters_for_physical_optimize_rules(
+        default_rules: &[Arc<dyn PhysicalOptimizerRule + Send + Sync>],
+    ) -> Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>> {
+        let mut new_rules = Vec::with_capacity(default_rules.len());
+        for rule in default_rules {
+            new_rules.push(physical_optimizer::may_adapt_optimize_rule(rule.clone()))
+        }
+
+        new_rules
+    }
+
+    fn register_analyzer_rules(mut state: SessionState) -> SessionState {
+        // Our analyzer has high priority, so first add we custom rules, then add the
+        // default ones.
+        state = state.with_analyzer_rules(vec![Arc::new(TypeConversion)]);
+        for rule in Analyzer::new().rules {
+            state = state.add_analyzer_rule(rule);
+        }
+
+        state
+    }
+}

query_engine/src/datafusion_impl/physical_plan.rs

@@ -20,28 +20,25 @@ use std::{
 };
 
 use async_trait::async_trait;
-use datafusion::{
-    execution::context::TaskContext,
-    physical_plan::{
-        coalesce_partitions::CoalescePartitionsExec, display::DisplayableExecutionPlan,
-        ExecutionPlan,
-    },
-    prelude::SessionContext,
+use datafusion::physical_plan::{
+    coalesce_partitions::CoalescePartitionsExec, display::DisplayableExecutionPlan, ExecutionPlan,
 };
 use generic_error::BoxError;
-use snafu::ResultExt;
+use snafu::{OptionExt, ResultExt};
 use table_engine::stream::{FromDfStream, SendableRecordBatchStream};
 
-use crate::{error::*, physical_planner::PhysicalPlan};
+use crate::{
+    error::*,
+    physical_planner::{PhysicalPlan, TaskContext},
+};
 
 pub struct DataFusionPhysicalPlanImpl {
-    ctx: SessionContext,
     plan: Arc<dyn ExecutionPlan>,
 }
 
 impl DataFusionPhysicalPlanImpl {
-    pub fn with_plan(ctx: SessionContext, plan: Arc<dyn ExecutionPlan>) -> Self {
-        Self { ctx, plan }
+    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+        Self { plan }
     }
 
     pub fn as_df_physical_plan(&self) -> Arc<dyn ExecutionPlan> {
@@ -59,12 +56,17 @@ impl Debug for DataFusionPhysicalPlanImpl {
 
 #[async_trait]
 impl PhysicalPlan for DataFusionPhysicalPlanImpl {
-    fn execute(&self) -> Result<SendableRecordBatchStream> {
-        let task_context = Arc::new(TaskContext::from(&self.ctx));
+    fn execute(&self, task_ctx: &TaskContext) -> Result<SendableRecordBatchStream> {
+        let df_task_ctx =
+            task_ctx
+                .try_to_datafusion_task_ctx()
+                .with_context(|| PhysicalPlanNoCause {
+                    msg: Some("datafusion task ctx not found".to_string()),
+                })?;
         let partition_count = self.plan.output_partitioning().partition_count();
         let df_stream = if partition_count <= 1 {
             self.plan
-                .execute(0, task_context)
+                .execute(0, df_task_ctx)
                 .box_err()
                 .context(PhysicalPlanWithCause {
                     msg: Some(format!("partition_count:{partition_count}")),
@@ -74,7 +76,7 @@ impl PhysicalPlan for DataFusionPhysicalPlanImpl {
             let plan = CoalescePartitionsExec::new(self.plan.clone());
             // MergeExec must produce a single partition
             assert_eq!(1, plan.output_partitioning().partition_count());
-            plan.execute(0, task_context)
+            plan.execute(0, df_task_ctx)
                 .box_err()
                 .context(PhysicalPlanWithCause {
                     msg: Some(format!("partition_count:{partition_count}")),