refactor: execute insert with select in query engine (#1181)

* refactor: execute insert with select in query engine

* fix: resolve PR comments

src/query/src/datafusion.rs

@@ -33,12 +33,21 @@ use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
 use common_telemetry::timer;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::ExecutionPlan;
+use datafusion_common::ResolvedTableReference;
+use datafusion_expr::{DmlStatement, LogicalPlan as DfLogicalPlan, WriteOp};
 use datatypes::schema::Schema;
+use futures_util::StreamExt;
+use session::context::QueryContextRef;
 use snafu::{OptionExt, ResultExt};
+use table::requests::InsertRequest;
+use table::TableRef;
 
 pub use crate::datafusion::catalog_adapter::DfCatalogListAdapter;
 pub use crate::datafusion::planner::DfContextProviderAdapter;
-use crate::error::{DataFusionSnafu, QueryExecutionSnafu, Result};
+use crate::error::{
+    CatalogNotFoundSnafu, CatalogSnafu, CreateRecordBatchSnafu, DataFusionSnafu,
+    QueryExecutionSnafu, Result, SchemaNotFoundSnafu, TableNotFoundSnafu, UnsupportedExprSnafu,
+};
 use crate::executor::QueryExecutor;
 use crate::logical_optimizer::LogicalOptimizer;
 use crate::physical_optimizer::PhysicalOptimizer;
@@ -56,6 +65,83 @@ impl DatafusionQueryEngine {
     pub fn new(state: Arc<QueryEngineState>) -> Self {
         Self { state }
     }
+
+    async fn exec_query_plan(&self, plan: LogicalPlan) -> Result<Output> {
+        let mut ctx = QueryEngineContext::new(self.state.session_state());
+
+        // `create_physical_plan` will optimize logical plan internally
+        let physical_plan = self.create_physical_plan(&mut ctx, &plan).await?;
+        let physical_plan = self.optimize_physical_plan(&mut ctx, physical_plan)?;
+
+        Ok(Output::Stream(self.execute_stream(&ctx, &physical_plan)?))
+    }
+
+    async fn exec_insert_plan(
+        &self,
+        dml: &DmlStatement,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        let default_catalog = query_ctx.current_catalog();
+        let default_schema = query_ctx.current_schema();
+        let table_name = dml
+            .table_name
+            .as_table_reference()
+            .resolve(&default_catalog, &default_schema);
+        let table = self.find_table(&table_name).await?;
+
+        let output = self
+            .exec_query_plan(LogicalPlan::DfPlan((*dml.input).clone()))
+            .await?;
+        let mut stream = match output {
+            Output::RecordBatches(batches) => batches.as_stream(),
+            Output::Stream(stream) => stream,
+            _ => unreachable!(),
+        };
+
+        let mut affected_rows = 0;
+        while let Some(batch) = stream.next().await {
+            let batch = batch.context(CreateRecordBatchSnafu)?;
+            let request = InsertRequest::try_from_recordbatch(&table_name, table.schema(), batch)
+                .map_err(BoxedError::new)
+                .context(QueryExecutionSnafu)?;
+
+            let rows = table
+                .insert(request)
+                .await
+                .map_err(BoxedError::new)
+                .context(QueryExecutionSnafu)?;
+            affected_rows += rows;
+        }
+        Ok(Output::AffectedRows(affected_rows))
+    }
+
+    async fn find_table(&self, table_name: &ResolvedTableReference<'_>) -> Result<TableRef> {
+        let catalog_name = table_name.catalog.as_ref();
+        let schema_name = table_name.schema.as_ref();
+        let table_name = table_name.table.as_ref();
+
+        let catalog = self
+            .state
+            .catalog_list()
+            .catalog(catalog_name)
+            .context(CatalogSnafu)?
+            .context(CatalogNotFoundSnafu {
+                catalog: catalog_name,
+            })?;
+        let schema =
+            catalog
+                .schema(schema_name)
+                .context(CatalogSnafu)?
+                .context(SchemaNotFoundSnafu {
+                    schema: schema_name,
+                })?;
+        let table = schema
+            .table(table_name)
+            .await
+            .context(CatalogSnafu)?
+            .context(TableNotFoundSnafu { table: table_name })?;
+        Ok(table)
+    }
 }
 
 #[async_trait]
@@ -75,14 +161,17 @@ impl QueryEngine for DatafusionQueryEngine {
         optimised_plan.schema()
     }
 
-    async fn execute(&self, plan: &LogicalPlan) -> Result<Output> {
-        let logical_plan = self.optimize(plan)?;
-
-        let mut ctx = QueryEngineContext::new(self.state.session_state());
-        let physical_plan = self.create_physical_plan(&mut ctx, &logical_plan).await?;
-        let physical_plan = self.optimize_physical_plan(&mut ctx, physical_plan)?;
-
-        Ok(Output::Stream(self.execute_stream(&ctx, &physical_plan)?))
+    async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output> {
+        match plan {
+            LogicalPlan::DfPlan(DfLogicalPlan::Dml(dml)) => match dml.op {
+                WriteOp::Insert => self.exec_insert_plan(&dml, query_ctx).await,
+                _ => UnsupportedExprSnafu {
+                    name: format!("DML op {}", dml.op),
+                }
+                .fail(),
+            },
+            _ => self.exec_query_plan(plan).await,
+        }
     }
 
     fn register_udf(&self, udf: ScalarUdf) {
@@ -292,11 +381,11 @@ mod tests {
         let stmt = QueryLanguageParser::parse_sql(sql).unwrap();
         let plan = engine
             .planner()
-            .plan(stmt, Arc::new(QueryContext::new()))
+            .plan(stmt, QueryContext::arc())
             .await
             .unwrap();
 
-        let output = engine.execute(&plan).await.unwrap();
+        let output = engine.execute(plan, QueryContext::arc()).await.unwrap();
 
         match output {
             Output::Stream(recordbatch) => {

src/query/src/error.rs

@@ -72,8 +72,11 @@ pub enum Error {
     #[snafu(display("The SQL string has multiple statements, query: {}", query))]
     MultipleStatements { query: String, backtrace: Backtrace },
 
-    #[snafu(display("Failed to convert datatype: {}", source))]
-    Datatype { source: datatypes::error::Error },
+    #[snafu(display("Failed to convert Datafusion schema: {}", source))]
+    ConvertDatafusionSchema {
+        #[snafu(backtrace)]
+        source: datatypes::error::Error,
+    },
 
     #[snafu(display("Failed to parse timestamp `{}`: {}", raw, source))]
     ParseTimestamp {
@@ -123,9 +126,10 @@ impl ErrorExt for Error {
             | ParseFloat { .. } => StatusCode::InvalidArguments,
             QueryAccessDenied { .. } => StatusCode::AccessDenied,
             Catalog { source } => source.status_code(),
-            VectorComputation { source } => source.status_code(),
+            VectorComputation { source } | ConvertDatafusionSchema { source } => {
+                source.status_code()
+            }
             CreateRecordBatch { source } => source.status_code(),
-            Datatype { source } => source.status_code(),
             QueryExecution { source } | QueryPlan { source } => source.status_code(),
             DataFusion { .. } => StatusCode::Internal,
             Sql { source } => source.status_code(),

src/query/src/plan.rs

@@ -18,7 +18,7 @@ use datafusion_expr::LogicalPlan as DfLogicalPlan;
 use datatypes::schema::Schema;
 use snafu::ResultExt;
 
-use crate::error::Result;
+use crate::error::{ConvertDatafusionSchemaSnafu, Result};
 
 /// A LogicalPlan represents the different types of relational
 /// operators (such as Projection, Filter, etc) and can be created by
@@ -42,7 +42,7 @@ impl LogicalPlan {
                 df_schema
                     .clone()
                     .try_into()
-                    .context(crate::error::DatatypeSnafu)
+                    .context(ConvertDatafusionSchemaSnafu)
             }
         }
     }

src/query/src/query_engine.rs

@@ -56,7 +56,7 @@ pub trait QueryEngine: Send + Sync {
 
     async fn describe(&self, plan: LogicalPlan) -> Result<Schema>;
 
-    async fn execute(&self, plan: &LogicalPlan) -> Result<Output>;
+    async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output>;
 
     fn register_udf(&self, udf: ScalarUdf);
 

src/query/src/tests.rs

@@ -41,7 +41,7 @@ async fn exec_selection(engine: QueryEngineRef, sql: &str) -> Vec<RecordBatch> {
         .await
         .unwrap();
     let Output::Stream(stream) = engine
-        .execute(&plan)
+        .execute(plan, QueryContext::arc())
         .await
         .unwrap() else { unreachable!() };
     util::collect(stream).await.unwrap()

src/query/src/tests/query_engine_test.rs

@@ -77,7 +77,7 @@ async fn test_datafusion_query_engine() -> Result<()> {
         .unwrap(),
     );
 
-    let output = engine.execute(&plan).await?;
+    let output = engine.execute(plan, QueryContext::arc()).await?;
 
     let recordbatch = match output {
         Output::Stream(recordbatch) => recordbatch,