feat: impl grpc physical plan (#212)

* chore: rename "convert.rs" to "serde.rs"

* proto definition

* impl "projection"

* add mock_input_exec for test

* impl physical plan execution

src/datanode/Cargo.toml

@@ -18,10 +18,12 @@ axum-macros = "0.2"
 catalog = { path = "../catalog" }
 common-base = { path = "../common/base" }
 common-error = { path = "../common/error" }
+common-grpc = { path = "../common/grpc" }
 common-recordbatch = { path = "../common/recordbatch" }
 common-runtime = { path = "../common/runtime" }
 common-telemetry = { path = "../common/telemetry" }
 common-time = { path = "../common/time" }
+datafusion = { git = "https://github.com/apache/arrow-datafusion.git", branch = "arrow2", features = ["simd"] }
 datatypes = { path = "../datatypes" }
 hyper = { version = "0.14", features = ["full"] }
 log-store = { path = "../log-store" }

src/datanode/src/error.rs

@@ -1,6 +1,6 @@
 use std::any::Any;
 
-use api::convert::DecodeError;
+use api::serde::DecodeError;
 use common_error::ext::BoxedError;
 use common_error::prelude::*;
 use datatypes::prelude::ConcreteDataType;
@@ -11,26 +11,32 @@ use table::error::Error as TableError;
 #[derive(Debug, Snafu)]
 #[snafu(visibility(pub))]
 pub enum Error {
-    #[snafu(display("Fail to execute sql, source: {}", source))]
+    #[snafu(display("Failed to execute sql, source: {}", source))]
     ExecuteSql {
         #[snafu(backtrace)]
         source: query::error::Error,
     },
 
-    #[snafu(display("Fail to create catalog list, source: {}", source))]
+    #[snafu(display("Failed to execute physical plan, source: {}", source))]
+    ExecutePhysicalPlan {
+        #[snafu(backtrace)]
+        source: query::error::Error,
+    },
+
+    #[snafu(display("Failed to create catalog list, source: {}", source))]
     NewCatalog {
         #[snafu(backtrace)]
         source: catalog::error::Error,
     },
 
-    #[snafu(display("Fail to create table: {}, {}", table_name, source))]
+    #[snafu(display("Failed to create table: {}, source: {}", table_name, source))]
     CreateTable {
         table_name: String,
         #[snafu(backtrace)]
         source: TableError,
     },
 
-    #[snafu(display("Fail to get table: {}, {}", table_name, source))]
+    #[snafu(display("Failed to get table: {}, source: {}", table_name, source))]
     GetTable {
         table_name: String,
         #[snafu(backtrace)]
@@ -53,7 +59,7 @@ pub enum Error {
     ))]
     ColumnValuesNumberMismatch { columns: usize, values: usize },
 
-    #[snafu(display("Fail to parse value: {}, {}", msg, backtrace))]
+    #[snafu(display("Failed to parse value: {}, {}", msg, backtrace))]
     ParseSqlValue { msg: String, backtrace: Backtrace },
 
     #[snafu(display(
@@ -68,7 +74,7 @@ pub enum Error {
         actual: ConcreteDataType,
     },
 
-    #[snafu(display("Fail to insert value to table: {}, {}", table_name, source))]
+    #[snafu(display("Failed to insert value to table: {}, source: {}", table_name, source))]
     Insert {
         table_name: String,
         source: TableError,
@@ -77,7 +83,7 @@ pub enum Error {
     #[snafu(display("Illegal insert data"))]
     IllegalInsertData,
 
-    #[snafu(display("Fail to convert bytes to insert batch, {}", source))]
+    #[snafu(display("Failed to convert bytes to insert batch, source: {}", source))]
     DecodeInsert { source: DecodeError },
 
     #[snafu(display("Failed to start server, source: {}", source))]
@@ -86,19 +92,19 @@ pub enum Error {
         source: servers::error::Error,
     },
 
-    #[snafu(display("Fail to parse address {}, source: {}", addr, source))]
+    #[snafu(display("Failed to parse address {}, source: {}", addr, source))]
     ParseAddr {
         addr: String,
         source: std::net::AddrParseError,
     },
 
-    #[snafu(display("Fail to bind address {}, source: {}", addr, source))]
+    #[snafu(display("Failed to bind address {}, source: {}", addr, source))]
     TcpBind {
         addr: String,
         source: std::io::Error,
     },
 
-    #[snafu(display("Fail to start gRPC server, source: {}", source))]
+    #[snafu(display("Failed to start gRPC server, source: {}", source))]
     StartGrpc { source: tonic::transport::Error },
 
     #[snafu(display("Failed to create directory {}, source: {}", dir, source))]
@@ -132,12 +138,18 @@ pub enum Error {
     #[snafu(display("Invalid CREATE TABLE sql statement, cause: {}", msg))]
     InvalidCreateTableSql { msg: String, backtrace: Backtrace },
 
-    #[snafu(display("Failed to create schema when creating table: {}", source))]
+    #[snafu(display("Failed to create schema when creating table, source: {}", source))]
     CreateSchema {
         #[snafu(backtrace)]
         source: datatypes::error::Error,
     },
 
+    #[snafu(display("Failed to convert datafusion schema, source: {}", source))]
+    ConvertSchema {
+        #[snafu(backtrace)]
+        source: datatypes::error::Error,
+    },
+
     #[snafu(display("SQL data type not supported yet: {:?}", t))]
     SqlTypeNotSupported {
         t: sql::ast::DataType,
@@ -156,11 +168,17 @@ pub enum Error {
         backtrace: Backtrace,
     },
 
-    #[snafu(display("Failed to insert into system catalog table: {}", source))]
+    #[snafu(display("Failed to insert into system catalog table, source: {}", source))]
     InsertSystemCatalog {
         #[snafu(backtrace)]
         source: catalog::error::Error,
     },
+
+    #[snafu(display("Failed to decode as physical plan, source: {}", source))]
+    IntoPhysicalPlan {
+        #[snafu(backtrace)]
+        source: common_grpc::Error,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -169,10 +187,12 @@ impl ErrorExt for Error {
     fn status_code(&self) -> StatusCode {
         match self {
             Error::ExecuteSql { source } => source.status_code(),
+            Error::ExecutePhysicalPlan { source } => source.status_code(),
             Error::NewCatalog { source } => source.status_code(),
             Error::CreateTable { source, .. } => source.status_code(),
             Error::GetTable { source, .. } => source.status_code(),
             Error::Insert { source, .. } => source.status_code(),
+            Error::ConvertSchema { source, .. } => source.status_code(),
             Error::TableNotFound { .. } => StatusCode::TableNotFound,
             Error::ColumnNotFound { .. } => StatusCode::TableColumnNotFound,
             Error::ColumnValuesNumberMismatch { .. }
@@ -193,6 +213,7 @@ impl ErrorExt for Error {
             | Error::CreateDir { .. }
             | Error::InsertSystemCatalog { .. }
             | Error::Conversion { .. }
+            | Error::IntoPhysicalPlan { .. }
             | Error::UnsupportedExpr { .. } => StatusCode::Internal,
             Error::InitBackend { .. } => StatusCode::StorageUnavailable,
             Error::OpenLogStore { source } => source.status_code(),

src/datanode/src/instance.rs

@@ -20,10 +20,12 @@ use table_engine::engine::MitoEngine;
 use crate::datanode::{DatanodeOptions, ObjectStoreConfig};
 use crate::error::{
     self, ExecuteSqlSnafu, InsertSnafu, NewCatalogSnafu, Result, TableNotFoundSnafu,
+    UnsupportedExprSnafu,
 };
 use crate::metric;
 use crate::server::grpc::handler::{build_err_result, ObjectResultBuilder};
 use crate::server::grpc::insert::insertion_expr_to_request;
+use crate::server::grpc::plan::PhysicalPlanner;
 use crate::server::grpc::select::to_object_result;
 use crate::sql::{SqlHandler, SqlRequest};
 
@@ -36,6 +38,7 @@ pub struct Instance {
     sql_handler: SqlHandler<DefaultEngine>,
     // Catalog list
     catalog_manager: CatalogManagerRef,
+    physical_planner: PhysicalPlanner,
 }
 
 pub type InstanceRef = Arc<Instance>;
@@ -63,9 +66,10 @@ impl Instance {
         let query_engine = factory.query_engine().clone();
 
         Ok(Self {
-            query_engine,
+            query_engine: query_engine.clone(),
             sql_handler: SqlHandler::new(table_engine, catalog_manager.clone()),
             catalog_manager,
+            physical_planner: PhysicalPlanner::new(query_engine),
         })
     }
 
@@ -167,12 +171,23 @@ impl Instance {
     }
 
     async fn handle_select(&self, select_expr: SelectExpr) -> ObjectResult {
-        match select_expr.expr {
-            Some(select_expr::Expr::Sql(sql)) => {
-                let result = self.execute_sql(&sql).await;
-                to_object_result(result).await
+        let result = self.do_handle_select(select_expr).await;
+        to_object_result(result).await
+    }
+
+    async fn do_handle_select(&self, select_expr: SelectExpr) -> Result<Output> {
+        let expr = select_expr.expr;
+        match expr {
+            Some(select_expr::Expr::Sql(sql)) => self.execute_sql(&sql).await,
+            Some(select_expr::Expr::PhysicalPlan(api::v1::PhysicalPlan { original_ql, plan })) => {
+                self.physical_planner
+                    .execute(PhysicalPlanner::parse(plan)?, original_ql)
+                    .await
             }
-            None => ObjectResult::default(),
+            _ => UnsupportedExprSnafu {
+                name: format!("{:?}", expr),
+            }
+            .fail(),
         }
     }
 

src/datanode/src/server/grpc.rs

@@ -1,3 +1,4 @@
 pub(crate) mod handler;
-pub mod insert;
-pub mod select;
+pub(crate) mod insert;
+pub(crate) mod plan;
+pub(crate) mod select;

src/datanode/src/server/grpc/insert.rs

@@ -1,5 +1,6 @@
 use std::{
     collections::{hash_map::Entry, HashMap},
+    ops::Deref,
     sync::Arc,
 };
 
@@ -68,7 +69,7 @@ fn insert_batches(bytes_vec: Vec<Vec<u8>>) -> Result<Vec<InsertBatch>> {
     let mut insert_batches = Vec::with_capacity(bytes_vec.len());
 
     for bytes in bytes_vec {
-        insert_batches.push(bytes.try_into().context(DecodeInsertSnafu)?);
+        insert_batches.push(bytes.deref().try_into().context(DecodeInsertSnafu)?);
     }
     Ok(insert_batches)
 }

src/datanode/src/server/grpc/plan.rs

@@ -0,0 +1,44 @@
+use std::sync::Arc;
+
+use common_grpc::AsExcutionPlan;
+use common_grpc::DefaultAsPlanImpl;
+use datatypes::schema::Schema;
+use query::PhysicalPlanAdapter;
+use query::{plan::PhysicalPlan, Output, QueryEngineRef};
+use snafu::ResultExt;
+
+use crate::error::Result;
+use crate::error::{ConvertSchemaSnafu, ExecutePhysicalPlanSnafu, IntoPhysicalPlanSnafu};
+
+pub type PhysicalPlanRef = Arc<dyn PhysicalPlan>;
+
+pub struct PhysicalPlanner {
+    query_engine: QueryEngineRef,
+}
+
+impl PhysicalPlanner {
+    pub fn new(query_engine: QueryEngineRef) -> Self {
+        Self { query_engine }
+    }
+
+    pub fn parse(bytes: Vec<u8>) -> Result<PhysicalPlanRef> {
+        let physical_plan = DefaultAsPlanImpl { bytes }
+            .try_into_physical_plan()
+            .context(IntoPhysicalPlanSnafu)?;
+
+        let schema: Arc<Schema> = Arc::new(
+            physical_plan
+                .schema()
+                .try_into()
+                .context(ConvertSchemaSnafu)?,
+        );
+        Ok(Arc::new(PhysicalPlanAdapter::new(schema, physical_plan)))
+    }
+
+    pub async fn execute(&self, plan: PhysicalPlanRef, _original_ql: Vec<u8>) -> Result<Output> {
+        self.query_engine
+            .execute_physical(&plan)
+            .await
+            .context(ExecutePhysicalPlanSnafu)
+    }
+}

src/datanode/src/tests/grpc_test.rs

@@ -2,7 +2,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 
-use api::v1::{codec::InsertBatch, column, select_expr, Column, SelectExpr};
+use api::v1::{codec::InsertBatch, column, Column};
 use client::{Client, Database, ObjectResult};
 use servers::grpc::GrpcServer;
 use servers::server::Server;
@@ -86,10 +86,10 @@ async fn test_insert_and_select() {
     assert!(result.is_ok());
 
     // select
-    let select_expr = SelectExpr {
-        expr: Some(select_expr::Expr::Sql("select * from demo".to_string())),
-    };
-    let result = db.select(select_expr).await.unwrap();
+    let result = db
+        .select(client::Select::Sql("select * from demo".to_string()))
+        .await
+        .unwrap();
     assert!(matches!(result, ObjectResult::Select(_)));
     match result {
         ObjectResult::Select(select_result) => {