tmp2

Signed-off-by: junxiang Mu <1948535941@qq.com>

s3select/query/src/execution/factory.rs

@@ -1,42 +1,29 @@
 use std::sync::Arc;
 
-use api::query::execution::QueryExecutionFactory;
+use api::{
+    query::{
+        execution::{QueryExecutionFactory, QueryExecutionRef, QueryStateMachineRef},
+        logical_planner::Plan,
+        optimizer::Optimizer,
+        scheduler::SchedulerRef,
+    },
+    QueryError,
+};
+use async_trait::async_trait;
+
+use super::query::SqlQueryExecution;
 
 pub type QueryExecutionFactoryRef = Arc<dyn QueryExecutionFactory + Send + Sync>;
 
 pub struct SqlQueryExecutionFactory {
     optimizer: Arc<dyn Optimizer + Send + Sync>,
     scheduler: SchedulerRef,
-    query_tracker: Arc<QueryTracker>,
-    trigger_executor_factory: TriggerExecutorFactoryRef,
-    runtime: Arc<DedicatedExecutor>,
-    stream_checker_manager: StreamCheckerManagerRef,
 }
 
 impl SqlQueryExecutionFactory {
     #[inline(always)]
-    pub fn new(
-        optimizer: Arc<dyn Optimizer + Send + Sync>,
-        scheduler: SchedulerRef,
-        query_tracker: Arc<QueryTracker>,
-        stream_checker_manager: StreamCheckerManagerRef,
-        config: Arc<QueryOptions>,
-    ) -> Self {
-        // Only do periodic scheduling, no need for many threads
-        let trigger_executor_runtime = DedicatedExecutor::new("stream-trigger", config.stream_trigger_cpu);
-        let trigger_executor_factory = Arc::new(TriggerExecutorFactory::new(Arc::new(trigger_executor_runtime)));
-
-        // perform stream-related preparations, not actual operator execution
-        let runtime = Arc::new(DedicatedExecutor::new("stream-executor", config.stream_executor_cpu));
-
-        Self {
-            optimizer,
-            scheduler,
-            query_tracker,
-            trigger_executor_factory,
-            runtime,
-            stream_checker_manager,
-        }
+    pub fn new(optimizer: Arc<dyn Optimizer + Send + Sync>, scheduler: SchedulerRef) -> Self {
+        Self { optimizer, scheduler }
     }
 }
 
@@ -48,62 +35,12 @@ impl QueryExecutionFactory for SqlQueryExecutionFactory {
         state_machine: QueryStateMachineRef,
     ) -> Result<QueryExecutionRef, QueryError> {
         match plan {
-            Plan::Query(query_plan) => {
-                // 获取执行计划中所有涉及到的stream source
-                let stream_providers = extract_stream_providers(&query_plan);
-
-                // (含有流表, explain, dml)
-                match (!stream_providers.is_empty(), query_plan.is_explain(), is_dml(&query_plan)) {
-                    (false, _, _) | (true, true, _) => Ok(Arc::new(SqlQueryExecution::new(
-                        state_machine,
-                        query_plan,
-                        self.optimizer.clone(),
-                        self.scheduler.clone(),
-                    ))),
-                    (true, false, true) => {
-                        // 流操作
-                        // stream source + dml + !explain
-                        let options = state_machine.session.inner().config().into();
-                        let exec = MicroBatchStreamExecutionBuilder::new(MicroBatchStreamExecutionDesc {
-                            plan: Arc::new(query_plan),
-                            options,
-                        })
-                        .with_stream_providers(stream_providers)
-                        .build(
-                            state_machine,
-                            self.scheduler.clone(),
-                            self.trigger_executor_factory.clone(),
-                            self.runtime.clone(),
-                        )
-                        .await?;
-
-                        Ok(Arc::new(exec))
-                    }
-                    (true, false, false) => {
-                        // stream source + !dml + !explain
-                        Err(QueryError::NotImplemented {
-                            err: "Stream table can only be used as source table in insert select statements.".to_string(),
-                        })
-                    }
-                }
-            }
-            Plan::DDL(ddl_plan) => Ok(Arc::new(DDLExecution::new(state_machine, self.stream_checker_manager.clone(), ddl_plan))),
-            Plan::DML(dml_plan) => Ok(Arc::new(DMLExecution::new(state_machine, dml_plan))),
-            Plan::SYSTEM(sys_plan) => Ok(Arc::new(SystemExecution::new(state_machine, sys_plan, self.query_tracker.clone()))),
+            Plan::Query(query_plan) => Ok(Arc::new(SqlQueryExecution::new(
+                state_machine,
+                query_plan,
+                self.optimizer.clone(),
+                self.scheduler.clone(),
+            ))),
         }
     }
 }
-
-fn is_dml(query_plan: &QueryPlan) -> bool {
-    match &query_plan.df_plan {
-        LogicalPlan::Dml(_) => true,
-        LogicalPlan::Extension(Extension { node }) => downcast_plan_node::<TableWriterMergePlanNode>(node.as_ref()).is_some(),
-        _ => false,
-    }
-}
-
-impl Drop for SqlQueryExecutionFactory {
-    fn drop(&mut self) {
-        self.runtime.shutdown();
-    }
-}

s3select/query/src/execution/mod.rs

@@ -1 +1,3 @@
 pub mod factory;
+pub mod query;
+pub mod scheduler;

s3select/query/src/execution/query.rs

@@ -0,0 +1,92 @@
+use std::sync::Arc;
+
+use api::query::execution::{Output, QueryExecution, QueryStateMachineRef};
+use api::query::logical_planner::QueryPlan;
+use api::query::optimizer::Optimizer;
+use api::query::scheduler::SchedulerRef;
+use api::{QueryError, QueryResult};
+use async_trait::async_trait;
+use futures::stream::AbortHandle;
+use parking_lot::Mutex;
+use tracing::debug;
+
+pub struct SqlQueryExecution {
+    query_state_machine: QueryStateMachineRef,
+    plan: QueryPlan,
+    optimizer: Arc<dyn Optimizer + Send + Sync>,
+    scheduler: SchedulerRef,
+
+    abort_handle: Mutex<Option<AbortHandle>>,
+}
+
+impl SqlQueryExecution {
+    pub fn new(
+        query_state_machine: QueryStateMachineRef,
+        plan: QueryPlan,
+        optimizer: Arc<dyn Optimizer + Send + Sync>,
+        scheduler: SchedulerRef,
+    ) -> Self {
+        Self {
+            query_state_machine,
+            plan,
+            optimizer,
+            scheduler,
+            abort_handle: Mutex::new(None),
+        }
+    }
+
+    async fn start(&self) -> QueryResult<Output> {
+        // begin optimize
+        self.query_state_machine.begin_optimize();
+        let physical_plan = self.optimizer.optimize(&self.plan, &self.query_state_machine.session).await?;
+        self.query_state_machine.end_optimize();
+
+        // begin schedule
+        self.query_state_machine.begin_schedule();
+        let stream = self
+            .scheduler
+            .schedule(physical_plan.clone(), self.query_state_machine.session.inner().task_ctx())
+            .await?
+            .stream();
+
+        debug!("Success build result stream.");
+        self.query_state_machine.end_schedule();
+
+        Ok(Output::StreamData(stream))
+    }
+}
+
+#[async_trait]
+impl QueryExecution for SqlQueryExecution {
+    async fn start(&self) -> QueryResult<Output> {
+        let (task, abort_handle) = futures::future::abortable(self.start());
+
+        {
+            *self.abort_handle.lock() = Some(abort_handle);
+        }
+
+        task.await.map_err(|_| QueryError::Cancel)?
+    }
+
+    fn cancel(&self) -> QueryResult<()> {
+        debug!(
+            "cancel sql query execution: sql: {}, state: {:?}",
+            self.query_state_machine.query.content(),
+            self.query_state_machine.state()
+        );
+
+        // change state
+        self.query_state_machine.cancel();
+        // stop future task
+        if let Some(e) = self.abort_handle.lock().as_ref() {
+            e.abort()
+        };
+
+        debug!(
+            "canceled sql query execution: sql: {}, state: {:?}",
+            self.query_state_machine.query.content(),
+            self.query_state_machine.state()
+        );
+        Ok(())
+    }
+}

s3select/query/src/execution/scheduler/local.rs

@@ -0,0 +1,22 @@
+use std::sync::Arc;
+
+use api::query::scheduler::{ExecutionResults, Scheduler};
+use async_trait::async_trait;
+use datafusion::error::DataFusionError;
+use datafusion::execution::context::TaskContext;
+use datafusion::physical_plan::{execute_stream, ExecutionPlan};
+
+pub struct LocalScheduler {}
+
+#[async_trait]
+impl Scheduler for LocalScheduler {
+    async fn schedule(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        context: Arc<TaskContext>,
+    ) -> Result<ExecutionResults, DataFusionError> {
+        let stream = execute_stream(plan, context)?;
+
+        Ok(ExecutionResults::new(stream))
+    }
+}

s3select/query/src/execution/scheduler/mod.rs

@@ -0,0 +1 @@
+pub mod local;