tmp1

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

s3select/query/src/execution/factory.rs

@@ -0,0 +1,109 @@
+use std::sync::Arc;
+
+use api::query::execution::QueryExecutionFactory;
+
+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,
+        }
+    }
+}
+
+#[async_trait]
+impl QueryExecutionFactory for SqlQueryExecutionFactory {
+    async fn create_query_execution(
+        &self,
+        plan: Plan,
+        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()))),
+        }
+    }
+}
+
+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

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