tmp1

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

s3select/api/src/lib.rs

@@ -0,0 +1,52 @@
+use std::fmt::Display;
+
+use datafusion::common::DataFusionError;
+use snafu::{Backtrace, Location, Snafu};
+
+pub mod query;
+pub mod server;
+
+pub type QueryResult<T> = Result<T, QueryError>;
+
+#[derive(Debug, Snafu)]
+#[snafu(visibility(pub))]
+pub enum QueryError {
+    Datafusion {
+        source: DataFusionError,
+        location: Location,
+        backtrace: Backtrace,
+    },
+}
+
+impl From<DataFusionError> for QueryError {
+    fn from(value: DataFusionError) -> Self {
+        match value {
+            DataFusionError::External(e) if e.downcast_ref::<QueryError>().is_some() => *e.downcast::<QueryError>().unwrap(),
+
+            v => Self::Datafusion {
+                source: v,
+                location: Default::default(),
+                backtrace: Backtrace::capture(),
+            },
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ResolvedTable {
+    // path
+    table: String,
+}
+
+impl ResolvedTable {
+    pub fn table(&self) -> &str {
+        &self.table
+    }
+}
+
+impl Display for ResolvedTable {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Self { table } = self;
+        write!(f, "{table}")
+    }
+}

s3select/api/src/query/ast.rs

@@ -0,0 +1,8 @@
+use datafusion::sql::sqlparser::ast::Statement;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ExtStatement {
+    /// ANSI SQL AST node
+    SqlStatement(Box<Statement>),
+    // we can expand command
+}

s3select/api/src/query/datasource/mod.rs

@@ -0,0 +1 @@
+

s3select/api/src/query/dispatcher.rs

@@ -0,0 +1,36 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use crate::QueryResult;
+
+use super::{
+    execution::{Output, QueryStateMachine},
+    logical_planner::Plan,
+    Query,
+};
+
+#[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);
+
+    async fn execute_query(&self, query: &Query) -> QueryResult<Output>;
+
+    async fn build_logical_plan(&self, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>>;
+
+    async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output>;
+
+    async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>>;
+
+    // fn running_query_infos(&self) -> Vec<QueryInfo>;
+
+    // fn running_query_status(&self) -> Vec<QueryStatus>;
+
+    // fn cancel_query(&self, id: &QueryId);
+}

s3select/api/src/query/execution.rs

@@ -0,0 +1,253 @@
+use std::fmt::Display;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicPtr, Ordering};
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::time::{Duration, Instant};
+
+use async_trait::async_trait;
+use datafusion::arrow::datatypes::{Schema, SchemaRef};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::physical_plan::SendableRecordBatchStream;
+use futures::{Stream, StreamExt, TryStreamExt};
+
+use crate::{QueryError, QueryResult};
+
+use super::logical_planner::Plan;
+use super::session::SessionCtx;
+use super::Query;
+
+pub type QueryExecutionRef = Arc<dyn QueryExecution>;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum QueryType {
+    Batch,
+    Stream,
+}
+
+impl Display for QueryType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Batch => write!(f, "batch"),
+            Self::Stream => write!(f, "stream"),
+        }
+    }
+}
+
+#[async_trait]
+pub trait QueryExecution: Send + Sync {
+    fn query_type(&self) -> QueryType {
+        QueryType::Batch
+    }
+    // 开始
+    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 {
+    StreamData(SendableRecordBatchStream),
+    Nil(()),
+}
+
+impl Output {
+    pub fn schema(&self) -> SchemaRef {
+        match self {
+            Self::StreamData(stream) => stream.schema(),
+            Self::Nil(_) => Arc::new(Schema::empty()),
+        }
+    }
+
+    pub async fn chunk_result(self) -> QueryResult<Vec<RecordBatch>> {
+        match self {
+            Self::Nil(_) => Ok(vec![]),
+            Self::StreamData(stream) => {
+                let schema = stream.schema();
+                let mut res: Vec<RecordBatch> = stream.try_collect::<Vec<RecordBatch>>().await?;
+                if res.is_empty() {
+                    res.push(RecordBatch::new_empty(schema));
+                }
+                Ok(res)
+            }
+        }
+    }
+
+    pub async fn num_rows(self) -> usize {
+        match self.chunk_result().await {
+            Ok(rb) => rb.iter().map(|e| e.num_rows()).sum(),
+            Err(_) => 0,
+        }
+    }
+
+    /// Returns the number of records affected by the query operation
+    ///
+    /// If it is a select statement, returns the number of rows in the result set
+    ///
+    /// -1 means unknown
+    ///
+    /// panic! when StreamData's number of records greater than i64::Max
+    pub async fn affected_rows(self) -> i64 {
+        self.num_rows().await as i64
+    }
+}
+
+impl Stream for Output {
+    type Item = std::result::Result<RecordBatch, QueryError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        match this {
+            Output::StreamData(stream) => stream.poll_next_unpin(cx).map_err(|e| e.into()),
+            Output::Nil(_) => Poll::Ready(None),
+        }
+    }
+}
+
+#[async_trait]
+pub trait QueryExecutionFactory {
+    async fn create_query_execution(
+        &self,
+        plan: Plan,
+        query_state_machine: QueryStateMachineRef,
+    ) -> QueryResult<QueryExecutionRef>;
+}
+
+pub type QueryStateMachineRef = Arc<QueryStateMachine>;
+
+pub struct QueryStateMachine {
+    pub session: SessionCtx,
+    pub query: Query,
+
+    state: AtomicPtr<QueryState>,
+    start: Instant,
+}
+
+impl QueryStateMachine {
+    pub fn begin(query: Query, session: SessionCtx) -> Self {
+        Self {
+            session,
+            query,
+            state: AtomicPtr::new(Box::into_raw(Box::new(QueryState::ACCEPTING))),
+            start: Instant::now(),
+        }
+    }
+
+    pub fn begin_analyze(&self) {
+        // TODO record time
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::ANALYZING)));
+    }
+
+    pub fn end_analyze(&self) {
+        // TODO record time
+    }
+
+    pub fn begin_optimize(&self) {
+        // TODO record time
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::OPTMIZING)));
+    }
+
+    pub fn end_optimize(&self) {
+        // TODO
+    }
+
+    pub fn begin_schedule(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::RUNNING(RUNNING::SCHEDULING)));
+    }
+
+    pub fn end_schedule(&self) {
+        // TODO
+    }
+
+    pub fn finish(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::FINISHED)));
+    }
+
+    pub fn cancel(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::CANCELLED)));
+    }
+
+    pub fn fail(&self) {
+        // TODO
+        self.translate_to(Box::new(QueryState::DONE(DONE::FAILED)));
+    }
+
+    pub fn state(&self) -> &QueryState {
+        unsafe { &*self.state.load(Ordering::Relaxed) }
+    }
+
+    pub fn duration(&self) -> Duration {
+        self.start.elapsed()
+    }
+
+    fn translate_to(&self, state: Box<QueryState>) {
+        self.state.store(Box::into_raw(state), Ordering::Relaxed);
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum QueryState {
+    ACCEPTING,
+    RUNNING(RUNNING),
+    DONE(DONE),
+}
+
+impl AsRef<str> for QueryState {
+    fn as_ref(&self) -> &str {
+        match self {
+            QueryState::ACCEPTING => "ACCEPTING",
+            QueryState::RUNNING(e) => e.as_ref(),
+            QueryState::DONE(e) => e.as_ref(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum RUNNING {
+    DISPATCHING,
+    ANALYZING,
+    OPTMIZING,
+    SCHEDULING,
+}
+
+impl AsRef<str> for RUNNING {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::DISPATCHING => "DISPATCHING",
+            Self::ANALYZING => "ANALYZING",
+            Self::OPTMIZING => "OPTMIZING",
+            Self::SCHEDULING => "SCHEDULING",
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum DONE {
+    FINISHED,
+    FAILED,
+    CANCELLED,
+}
+
+impl AsRef<str> for DONE {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::FINISHED => "FINISHED",
+            Self::FAILED => "FAILED",
+            Self::CANCELLED => "CANCELLED",
+        }
+    }
+}

s3select/api/src/query/function.rs

@@ -0,0 +1,23 @@
+use std::collections::HashSet;
+use std::sync::Arc;
+
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+
+use crate::QueryResult;
+
+pub type FuncMetaManagerRef = Arc<dyn FunctionMetadataManager + Send + Sync>;
+pub trait FunctionMetadataManager {
+    fn register_udf(&mut self, udf: ScalarUDF) -> QueryResult<()>;
+
+    fn register_udaf(&mut self, udaf: AggregateUDF) -> QueryResult<()>;
+
+    fn register_udwf(&mut self, udwf: WindowUDF) -> QueryResult<()>;
+
+    fn udf(&self, name: &str) -> QueryResult<Arc<ScalarUDF>>;
+
+    fn udaf(&self, name: &str) -> QueryResult<Arc<AggregateUDF>>;
+
+    fn udwf(&self, name: &str) -> QueryResult<Arc<WindowUDF>>;
+
+    fn udfs(&self) -> HashSet<String>;
+}

s3select/api/src/query/logical_planner.rs

@@ -0,0 +1,29 @@
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::logical_expr::LogicalPlan as DFPlan;
+
+#[derive(Clone)]
+pub enum Plan {
+    // only support query sql
+    /// Query plan
+    Query(QueryPlan),
+}
+
+impl Plan {
+    pub fn schema(&self) -> SchemaRef {
+        match self {
+            Self::Query(p) => SchemaRef::from(p.df_plan.schema().as_ref().to_owned()),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct QueryPlan {
+    pub df_plan: DFPlan,
+    pub is_tag_scan: bool,
+}
+
+impl QueryPlan {
+    pub fn is_explain(&self) -> bool {
+        matches!(self.df_plan, DFPlan::Explain(_) | DFPlan::Analyze(_))
+    }
+}

s3select/api/src/query/mod.rs

@@ -0,0 +1,19 @@
+pub mod ast;
+pub mod datasource;
+pub mod dispatcher;
+pub mod execution;
+pub mod function;
+pub mod logical_planner;
+pub mod parser;
+pub mod session;
+
+#[derive(Clone)]
+pub struct Context {
+    // maybe we need transfer some info?
+}
+
+#[derive(Clone)]
+pub struct Query {
+    context: Context,
+    content: String,
+}

s3select/api/src/query/parser.rs

@@ -0,0 +1,8 @@
+use std::collections::VecDeque;
+
+use super::ast::ExtStatement;
+use crate::QueryResult;
+
+pub trait Parser {
+    fn parse(&self, sql: &str) -> QueryResult<VecDeque<ExtStatement>>;
+}

s3select/api/src/query/session.rs

@@ -0,0 +1,20 @@
+use std::sync::Arc;
+
+use datafusion::execution::context::SessionState;
+
+#[derive(Clone)]
+pub struct SessionCtx {
+    desc: Arc<SessionCtxDesc>,
+    inner: SessionState,
+}
+
+impl SessionCtx {
+    pub fn inner(&self) -> &SessionState {
+        &self.inner
+    }
+}
+
+#[derive(Clone)]
+pub struct SessionCtxDesc {
+    // maybe we need some info
+}

s3select/api/src/server/dbms.rs

@@ -0,0 +1,43 @@
+use async_trait::async_trait;
+
+use crate::{
+    query::{
+        execution::{Output, QueryStateMachineRef},
+        logical_planner::Plan,
+        Query,
+    },
+    QueryResult,
+};
+
+pub struct QueryHandle {
+    query: Query,
+    result: Output,
+}
+
+impl QueryHandle {
+    pub fn new(query: Query, result: Output) -> Self {
+        Self { query, result }
+    }
+
+    pub fn query(&self) -> &Query {
+        &self.query
+    }
+
+    pub fn result(self) -> Output {
+        self.result
+    }
+}
+
+#[async_trait]
+pub trait DatabaseManagerSystem {
+    async fn start(&self) -> QueryResult<()>;
+    async fn execute(&self, query: &Query) -> QueryResult<QueryHandle>;
+    async fn build_query_state_machine(&self, query: Query) -> QueryResult<QueryStateMachineRef>;
+    async fn build_logical_plan(&self, query_state_machine: QueryStateMachineRef) -> QueryResult<Option<Plan>>;
+    async fn execute_logical_plan(
+        &self,
+        logical_plan: Plan,
+        query_state_machine: QueryStateMachineRef,
+    ) -> QueryResult<QueryHandle>;
+    fn metrics(&self) -> String;
+}

s3select/api/src/server/mod.rs

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