Merge pull request #7756 from youngsofun/load

refactor(query): streaming load use planner v2.

Author: mergify[bot]

src/query/service/src/interpreters/interpreter_insert_v2.rs

@@ -171,12 +171,8 @@ impl Interpreter for InsertInterpreterV2 {
                     }
                     build_res.main_pipeline.add_pipe(builder.finalize());
                 }
-                InsertInputSource::StreamingWithFormat(_) => {
-                    build_res.main_pipeline.add_pipe(
-                        ((*self.source_pipe_builder.lock()).clone())
-                            .ok_or_else(|| ErrorCode::EmptyData("empty source pipe builder"))?
-                            .finalize(),
-                    );
+                InsertInputSource::StreamingWithFormat(_, pipe) => {
+                    build_res.main_pipeline.add_pipe(pipe.clone());
                 }
                 InsertInputSource::SelectPlan(plan) => {
                     let table1 = table.clone();

src/query/service/src/servers/http/v1/load.rs

@@ -17,14 +17,10 @@ use std::sync::Arc;
 
 use common_base::base::ProgressValues;
 use common_base::base::TrySpawn;
-use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_formats::FormatFactory;
 use common_io::prelude::parse_escape_string;
-use common_io::prelude::FormatSettings;
-use common_legacy_planners::InsertInputSource;
-use common_legacy_planners::PlanNode;
 use futures::StreamExt;
 use poem::error::InternalServerError;
 use poem::error::Result as PoemResult;
@@ -37,14 +33,14 @@ use serde::Serialize;
 use tracing::error;
 
 use super::HttpQueryContext;
-use crate::interpreters::InterpreterFactory;
-use crate::pipelines::SourcePipeBuilder;
+use crate::interpreters::InterpreterFactoryV2;
 use crate::servers::http::v1::multipart_format::MultipartFormat;
-use crate::servers::http::v1::multipart_format::MultipartWorker;
 use crate::sessions::QueryContext;
 use crate::sessions::SessionType;
 use crate::sessions::TableContext;
-use crate::sql::PlanParser;
+use crate::sql::plans::InsertInputSource;
+use crate::sql::plans::Plan;
+use crate::sql::Planner;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct LoadResponse {
@@ -55,32 +51,15 @@ pub struct LoadResponse {
     pub files: Vec<String>,
 }
 
-fn get_input_format(node: &PlanNode) -> Result<&str> {
-    match node {
-        PlanNode::Insert(insert) => match &insert.source {
-            InsertInputSource::StreamingWithFormat(format) => Ok(format),
-            _ => Err(ErrorCode::UnknownFormat("Not found format name in plan")),
-        },
-        _ => Err(ErrorCode::UnknownFormat("Not found format name in plan")),
-    }
-}
-
 #[allow(clippy::manual_async_fn)]
-fn execute_query(
-    context: Arc<QueryContext>,
-    node: PlanNode,
-    source_builder: SourcePipeBuilder,
-) -> impl Future<Output = Result<()>> {
+fn execute_query(context: Arc<QueryContext>, plan: Plan) -> impl Future<Output = Result<()>> {
     async move {
-        let interpreter = InterpreterFactory::get(context.clone(), node).await?;
+        let interpreter = InterpreterFactoryV2::get(context.clone(), &plan).await?;
 
         if let Err(cause) = interpreter.start().await {
             error!("interpreter.start error: {:?}", cause);
         }
 
-        // TODO(Winter): very hack code. need remove it.
-        interpreter.set_source_pipe_builder(Option::from(source_builder))?;
-
         let mut data_stream = interpreter.execute(context).await?;
 
         while let Some(_block) = data_stream.next().await {}
@@ -94,37 +73,6 @@ fn execute_query(
     }
 }
 
-async fn new_processor_format(
-    ctx: &Arc<QueryContext>,
-    node: &PlanNode,
-    multipart: Multipart,
-) -> Result<Json<LoadResponse>> {
-    let format = get_input_format(node)?;
-    let format_settings = ctx.get_format_settings()?;
-
-    let (mut worker, builder) =
-        format_source_pipe_builder(format, ctx, node.schema(), multipart, &format_settings)?;
-
-    let handler = ctx.spawn(execute_query(ctx.clone(), node.clone(), builder));
-
-    worker.work().await;
-    let files = worker.get_files();
-
-    match handler.await {
-        Ok(Ok(_)) => Ok(()),
-        Ok(Err(cause)) => Err(cause),
-        Err(_) => Err(ErrorCode::TokioError("Maybe panic.")),
-    }?;
-
-    Ok(Json(LoadResponse {
-        error: None,
-        state: "SUCCESS".to_string(),
-        id: uuid::Uuid::new_v4().to_string(),
-        stats: ctx.get_scan_progress_value(),
-        files,
-    }))
-}
-
 #[poem::handler]
 pub async fn streaming_load(
     ctx: &HttpQueryContext,
@@ -155,95 +103,83 @@ pub async fn streaming_load(
         }
     }
 
-    let mut plan = PlanParser::parse(context.clone(), insert_sql)
+    let mut planner = Planner::new(context.clone());
+    let (mut plan, _, _) = planner
+        .plan_sql(insert_sql)
         .await
         .map_err(InternalServerError)?;
     context.attach_query_str(insert_sql);
 
-    // Block size.
-    let _max_block_size = settings.get_max_block_size().map_err(InternalServerError)? as usize;
-
     let format_settings = context.get_format_settings().map_err(InternalServerError)?;
-    let source_pipe_builder: SourcePipeBuilder = match &mut plan {
-        PlanNode::Insert(insert) => match &mut insert.source {
-            InsertInputSource::StreamingWithFormat(format) => {
-                if FormatFactory::instance().has_input(format.as_str()) {
+    let schema = plan.schema();
+    match &mut plan {
+        Plan::Insert(insert) => match &mut insert.source {
+            InsertInputSource::StrWithFormat((sql_rest, format)) => {
+                if !sql_rest.is_empty() {
+                    Err(poem::Error::from_string(
+                        "should NOT have data after `Format` in streaming load.",
+                        StatusCode::BAD_REQUEST,
+                    ))
+                } else if FormatFactory::instance().has_input(format.as_str()) {
                     let new_format = format!("{}WithNames", format);
                     if format_settings.skip_header > 0
                         && FormatFactory::instance().has_input(new_format.as_str())
                     {
                         *format = new_format;
                     }
 
-                    return match new_processor_format(&context, &plan, multipart).await {
-                        Ok(res) => Ok(res),
-                        Err(cause) => Err(InternalServerError(cause)),
-                    };
+                    let format_settings =
+                        context.get_format_settings().map_err(InternalServerError)?;
+
+                    let (mut worker, builder) = MultipartFormat::input_sources(
+                        format,
+                        context.clone(),
+                        multipart,
+                        schema,
+                        format_settings.clone(),
+                    )
+                    .map_err(InternalServerError)?;
+
+                    insert.source = InsertInputSource::StreamingWithFormat(
+                        format.to_string(),
+                        builder.finalize(),
+                    );
+
+                    let handler = context.spawn(execute_query(context.clone(), plan));
+
+                    worker.work().await;
+                    let files = worker.get_files();
+
+                    match handler.await {
+                        Ok(Ok(_)) => Ok(Json(LoadResponse {
+                            error: None,
+                            state: "SUCCESS".to_string(),
+                            id: uuid::Uuid::new_v4().to_string(),
+                            stats: context.get_scan_progress_value(),
+                            files,
+                        })),
+                        Ok(Err(cause)) => Err(cause),
+                        Err(_) => Err(ErrorCode::TokioError("Maybe panic.")),
+                    }
+                    .map_err(InternalServerError)
+                } else {
+                    Err(poem::Error::from_string(
+                        format!("format not supported for streaming load {}", format),
+                        StatusCode::BAD_REQUEST,
+                    ))
                 }
-
-                Err(poem::Error::from_string(
-                    format!(
-                        "Streaming load only supports csv format, but got {}",
-                        format
-                    ),
-                    StatusCode::BAD_REQUEST,
-                ))
             }
             _non_supported_source => Err(poem::Error::from_string(
-                "Only supports streaming upload. e.g. INSERT INTO $table FORMAT CSV",
+                "Only supports streaming upload. e.g. INSERT INTO $table FORMAT CSV, got insert ... select.",
                 StatusCode::BAD_REQUEST,
             )),
         },
         non_insert_plan => Err(poem::Error::from_string(
             format!(
                 "Only supports INSERT statement in streaming load, but got {}",
-                non_insert_plan.name()
+                non_insert_plan
             ),
             StatusCode::BAD_REQUEST,
         )),
-    }?;
-    let interpreter = InterpreterFactory::get(context.clone(), plan.clone())
-        .await
-        .map_err(InternalServerError)?;
-    let _ = interpreter
-        .set_source_pipe_builder(Some(source_pipe_builder))
-        .map_err(|e| error!("interpreter.set_source_pipe_builder.error: {:?}", e));
-    interpreter.start().await.map_err(InternalServerError)?;
-    let mut data_stream = interpreter
-        .execute(context.clone())
-        .await
-        .map_err(InternalServerError)?;
-    while let Some(_block) = data_stream.next().await {}
-    // Write Finish to query log table.
-    let _ = interpreter
-        .finish()
-        .await
-        .map_err(|e| error!("interpreter.finish error: {:?}", e));
-
-    // TODO generate id
-    // TODO duplicate by insert_label
-    let mut id = uuid::Uuid::new_v4().to_string();
-    Ok(Json(LoadResponse {
-        id,
-        state: "SUCCESS".to_string(),
-        stats: context.get_scan_progress_value(),
-        error: None,
-        files: vec![],
-    }))
-}
-
-fn format_source_pipe_builder(
-    format: &str,
-    context: &Arc<QueryContext>,
-    schema: DataSchemaRef,
-    multipart: Multipart,
-    format_settings: &FormatSettings,
-) -> Result<(Box<dyn MultipartWorker>, SourcePipeBuilder)> {
-    MultipartFormat::input_sources(
-        format,
-        context.clone(),
-        multipart,
-        schema,
-        format_settings.clone(),
-    )
+    }
 }

src/query/service/src/sql/planner/plans/insert.rs

@@ -15,6 +15,7 @@
 use common_datablocks::DataBlock;
 use common_datavalues::DataSchemaRef;
 use common_meta_types::MetaId;
+use common_pipeline_core::Pipe;
 
 use super::Plan;
 
@@ -23,7 +24,8 @@ pub enum InsertInputSource {
     #[serde(skip)]
     SelectPlan(Box<Plan>),
     // From outside streaming source
-    StreamingWithFormat(String),
+    #[serde(skip)]
+    StreamingWithFormat(String, Pipe),
     // From cloned String and format
     StrWithFormat((String, String)),
 }
@@ -66,7 +68,7 @@ impl Insert {
     pub fn format(&self) -> Option<&str> {
         match &self.source {
             InsertInputSource::SelectPlan(_) => None,
-            InsertInputSource::StreamingWithFormat(v) => Some(v.as_str()),
+            InsertInputSource::StreamingWithFormat(v, _) => Some(v.as_str()),
             InsertInputSource::StrWithFormat((_, v)) => Some(v.as_str()),
         }
     }