feat: streaming insert from select query results

Author: killme2008

Cargo.lock

@@ -27,6 +27,7 @@ common-runtime = { path = "../common/runtime" }
 common-telemetry = { path = "../common/telemetry" }
 common-time = { path = "../common/time" }
 datafusion.workspace = true
+datafusion-expr.workspace = true
 datatypes = { path = "../datatypes" }
 futures = "0.3"
 hyper = { version = "0.14", features = ["full"] }

src/datanode/Cargo.toml

@@ -15,6 +15,7 @@
 use std::any::Any;
 
 use common_error::prelude::*;
+use datatypes::prelude::ConcreteDataType;
 use storage::error::Error as StorageError;
 use table::error::Error as TableError;
 
@@ -101,6 +102,18 @@ pub enum Error {
     ))]
     ColumnValuesNumberMismatch { columns: usize, values: usize },
 
+    #[snafu(display(
+        "Columns type mismatch, column: {}, expected type: {:?}, acutal: {:?}",
+        column,
+        expected,
+        actual,
+    ))]
+    ColumnTypeMismatch {
+        column: String,
+        expected: ConcreteDataType,
+        actual: ConcreteDataType,
+    },
+
     #[snafu(display("Failed to parse sql value, source: {}", source))]
     ParseSqlValue {
         #[snafu(backtrace)]
@@ -378,6 +391,7 @@ impl ErrorExt for Error {
             | Error::VectorComputation { source } => source.status_code(),
 
             Error::ColumnValuesNumberMismatch { .. }
+            | Error::ColumnTypeMismatch { .. }
             | Error::InvalidSql { .. }
             | Error::NotSupportSql { .. }
             | Error::KeyColumnNotFound { .. }

src/datanode/src/error.rs

@@ -21,6 +21,7 @@ use common_recordbatch::RecordBatches;
 use common_telemetry::logging::info;
 use common_telemetry::timer;
 use datatypes::schema::Schema;
+use futures::StreamExt;
 use query::parser::{QueryLanguageParser, QueryStatement};
 use servers::error as server_error;
 use servers::promql::PromqlHandler;
@@ -35,6 +36,7 @@ use table::requests::{CreateDatabaseRequest, DropTableRequest};
 use crate::error::{self, BumpTableIdSnafu, ExecuteSqlSnafu, Result, TableIdProviderNotFoundSnafu};
 use crate::instance::Instance;
 use crate::metric;
+use crate::sql::insert::InsertRequests;
 use crate::sql::SqlRequest;
 
 impl Instance {
@@ -56,11 +58,28 @@ impl Instance {
                     .context(ExecuteSqlSnafu)
             }
             QueryStatement::Sql(Statement::Insert(i)) => {
-                let request = self
+                let requests = self
                     .sql_handler
                     .insert_to_request(self.catalog_manager.clone(), *i, query_ctx.clone())
                     .await?;
-                self.sql_handler.execute(request, query_ctx).await
+
+                match requests {
+                    InsertRequests::Request(request) => {
+                        self.sql_handler.execute(request, query_ctx.clone()).await
+                    }
+
+                    InsertRequests::Stream(mut s) => {
+                        let mut output = None;
+                        while let Some(request) = s.next().await {
+                            output = Some(
+                                self.sql_handler
+                                    .execute(request?, query_ctx.clone())
+                                    .await?,
+                            )
+                        }
+                        Ok(output.unwrap())
+                    }
+                }
             }
             QueryStatement::Sql(Statement::Delete(d)) => {
                 let request = SqlRequest::Delete(*d);

src/datanode/src/instance/sql.rs

@@ -34,7 +34,7 @@ mod alter;
 mod create;
 mod delete;
 mod drop_table;
-mod insert;
+pub(crate) mod insert;
 
 #[derive(Debug)]
 pub enum SqlRequest {
@@ -160,6 +160,7 @@ mod tests {
     use tempdir::TempDir;
 
     use super::*;
+    use crate::sql::insert::InsertRequests;
 
     struct DemoTable;
 
@@ -260,7 +261,7 @@ mod tests {
             .unwrap();
 
         match request {
-            SqlRequest::Insert(req) => {
+            InsertRequests::Request(SqlRequest::Insert(req)) => {
                 assert_eq!(req.table_name, "demo");
                 let columns_values = req.columns_values;
                 assert_eq!(4, columns_values.len());