tmp2

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

s3select/api/src/query/analyzer.rs

@@ -0,0 +1,12 @@
+use std::sync::Arc;
+
+use datafusion::logical_expr::LogicalPlan;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type AnalyzerRef = Arc<dyn Analyzer + Send + Sync>;
+
+pub trait Analyzer {
+    fn analyze(&self, plan: &LogicalPlan, session: &SessionCtx) -> QueryResult<LogicalPlan>;
+}

s3select/api/src/query/dispatcher.rs

@@ -12,10 +12,6 @@ use super::{
 
 #[async_trait]
 pub trait QueryDispatcher: Send + Sync {
-    async fn start(&self) -> QueryResult<()>;
-
-    fn stop(&self);
-
     // fn create_query_id(&self) -> QueryId;
 
     // fn query_info(&self, id: &QueryId);

s3select/api/src/query/execution.rs

@@ -43,18 +43,6 @@ pub trait QueryExecution: Send + Sync {
     async fn start(&self) -> QueryResult<Output>;
     // 停止
     fn cancel(&self) -> QueryResult<()>;
-    // query状态
-    // 查询计划
-    // 静态信息
-    // fn info(&self) -> QueryInfo;
-    // 运行时信息
-    // fn status(&self) -> QueryStatus;
-    // sql
-    // 资源占用（cpu时间/内存/吞吐量等）
-    // 是否需要持久化query信息
-    fn need_persist(&self) -> bool {
-        false
-    }
 }
 
 pub enum Output {

s3select/api/src/query/logical_planner.rs

@@ -1,6 +1,12 @@
+use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::logical_expr::LogicalPlan as DFPlan;
 
+use crate::QueryResult;
+
+use super::ast::ExtStatement;
+use super::session::SessionCtx;
+
 #[derive(Clone)]
 pub enum Plan {
     // only support query sql
@@ -27,3 +33,8 @@ impl QueryPlan {
         matches!(self.df_plan, DFPlan::Explain(_) | DFPlan::Analyze(_))
     }
 }
+
+#[async_trait]
+pub trait LogicalPlanner {
+    async fn create_logical_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan>;
+}

s3select/api/src/query/mod.rs

@@ -1,15 +1,22 @@
+use s3s::dto::SelectObjectContentInput;
+
+pub mod analyzer;
 pub mod ast;
 pub mod datasource;
 pub mod dispatcher;
 pub mod execution;
 pub mod function;
 pub mod logical_planner;
+pub mod optimizer;
 pub mod parser;
+pub mod physical_planner;
+pub mod scheduler;
 pub mod session;
 
 #[derive(Clone)]
 pub struct Context {
     // maybe we need transfer some info?
+    pub input: SelectObjectContentInput,
 }
 
 #[derive(Clone)]
@@ -17,3 +24,18 @@ pub struct Query {
     context: Context,
     content: String,
 }
+
+impl Query {
+    #[inline(always)]
+    pub fn new(context: Context, content: String) -> Self {
+        Self { context, content }
+    }
+
+    pub fn context(&self) -> &Context {
+        &self.context
+    }
+
+    pub fn content(&self) -> &str {
+        self.content.as_str()
+    }
+}

s3select/api/src/query/optimizer.rs

@@ -0,0 +1,15 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::physical_plan::ExecutionPlan;
+
+use super::logical_planner::QueryPlan;
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type OptimizerRef = Arc<dyn Optimizer + Send + Sync>;
+
+#[async_trait]
+pub trait Optimizer {
+    async fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>>;
+}

s3select/api/src/query/physical_planner.rs

@@ -0,0 +1,21 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_planner::ExtensionPlanner;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+#[async_trait]
+pub trait PhysicalPlanner {
+    /// Given a `LogicalPlan`, create an `ExecutionPlan` suitable for execution
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionCtx,
+    ) -> QueryResult<Arc<dyn ExecutionPlan>>;
+
+    fn inject_physical_transform_rule(&mut self, rule: Arc<dyn ExtensionPlanner + Send + Sync>);
+}

s3select/api/src/query/scheduler.rs

@@ -0,0 +1,32 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::common::Result;
+use datafusion::execution::context::TaskContext;
+use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
+
+pub type SchedulerRef = Arc<dyn Scheduler + Send + Sync>;
+
+#[async_trait]
+pub trait Scheduler {
+    /// Schedule the provided [`ExecutionPlan`] on this [`Scheduler`].
+    ///
+    /// Returns a [`ExecutionResults`] that can be used to receive results as they are produced,
+    /// as a [`futures::Stream`] of [`RecordBatch`]
+    async fn schedule(&self, plan: Arc<dyn ExecutionPlan>, context: Arc<TaskContext>) -> Result<ExecutionResults>;
+}
+
+pub struct ExecutionResults {
+    stream: SendableRecordBatchStream,
+}
+
+impl ExecutionResults {
+    pub fn new(stream: SendableRecordBatchStream) -> Self {
+        Self { stream }
+    }
+
+    /// Returns a [`SendableRecordBatchStream`] of this execution
+    pub fn stream(self) -> SendableRecordBatchStream {
+        self.stream
+    }
+}

s3select/api/src/query/session.rs

@@ -1,10 +1,17 @@
 use std::sync::Arc;
 
-use datafusion::execution::context::SessionState;
+use datafusion::{
+    execution::{context::SessionState, runtime_env::RuntimeEnvBuilder, SessionStateBuilder},
+    prelude::SessionContext,
+};
+
+use crate::{object_store::EcObjectStore, QueryError, QueryResult};
+
+use super::Context;
 
 #[derive(Clone)]
 pub struct SessionCtx {
-    desc: Arc<SessionCtxDesc>,
+    _desc: Arc<SessionCtxDesc>,
     inner: SessionState,
 }
 
@@ -18,3 +25,33 @@ impl SessionCtx {
 pub struct SessionCtxDesc {
     // maybe we need some info
 }
+
+#[derive(Default)]
+pub struct SessionCtxFactory {}
+
+impl SessionCtxFactory {
+    pub fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
+        let df_session_ctx = self.build_df_session_context(context)?;
+
+        Ok(SessionCtx {
+            _desc: Arc::new(SessionCtxDesc {}),
+            inner: df_session_ctx.state(),
+        })
+    }
+
+    fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
+        let path = format!("s3://{}", context.input.bucket);
+        let store_url = url::Url::parse(&path).unwrap();
+        let store = EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
+
+        let rt = RuntimeEnvBuilder::new().build()?;
+        let df_session_state = SessionStateBuilder::new()
+            .with_runtime_env(Arc::new(rt))
+            .with_object_store(&store_url, Arc::new(store))
+            .with_default_features()
+            .build();
+        let df_session_ctx = SessionContext::new_with_state(df_session_state);
+
+        Ok(df_session_ctx)
+    }
+}