Merge pull request #9063 from ariesdevil/stream-load-refactor

feat(streaming_load): Support file_format syntax in streaming load insert sql

Author: BohuTANG

src/meta/types/src/user_stage.rs

@@ -127,6 +127,24 @@ impl FromStr for StageFileCompression {
     }
 }
 
+impl ToString for StageFileCompression {
+    fn to_string(&self) -> String {
+        match *self {
+            StageFileCompression::Auto => "auto".to_string(),
+            StageFileCompression::Gzip => "gzip".to_string(),
+            StageFileCompression::Bz2 => "bz2".to_string(),
+            StageFileCompression::Brotli => "brotli".to_string(),
+            StageFileCompression::Zstd => "zstd".to_string(),
+            StageFileCompression::Deflate => "deflate".to_string(),
+            StageFileCompression::RawDeflate => "raw_deflate".to_string(),
+            StageFileCompression::Lzo => "lzo".to_string(),
+            StageFileCompression::Snappy => "snappy".to_string(),
+            StageFileCompression::Xz => "xz".to_string(),
+            StageFileCompression::None => "none".to_string(),
+        }
+    }
+}
+
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq)]
 pub enum StageFileFormatType {
     Csv,
@@ -164,6 +182,12 @@ impl FromStr for StageFileFormatType {
     }
 }
 
+impl ToString for StageFileFormatType {
+    fn to_string(&self) -> String {
+        format!("{:?}", *self)
+    }
+}
+
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq)]
 #[serde(default)]
 pub struct FileFormatOptions {

src/query/ast/src/ast/format/ast_format.rs

@@ -963,6 +963,24 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 let streaming_node = FormatTreeNode::new(streaming_format_ctx);
                 self.children.push(streaming_node);
             }
+            InsertSource::StreamingV2 { settings, .. } => {
+                let mut file_formats_children = Vec::with_capacity(settings.len());
+                for (k, v) in settings.iter() {
+                    let file_format_name = format!("FileFormat {} = {:?}", k, v);
+                    let file_format_format_ctx = AstFormatContext::new(file_format_name);
+                    let file_format_node = FormatTreeNode::new(file_format_format_ctx);
+                    file_formats_children.push(file_format_node);
+                }
+                let file_formats_format_name = "StreamSourceFileFormats".to_string();
+                let files_formats_format_ctx = AstFormatContext::with_children(
+                    file_formats_format_name,
+                    file_formats_children.len(),
+                );
+                let files_formats_node =
+                    FormatTreeNode::with_children(files_formats_format_ctx, file_formats_children);
+
+                self.children.push(files_formats_node);
+            }
             InsertSource::Values { .. } => {
                 let values_name = "ValueSource".to_string();
                 let values_format_ctx = AstFormatContext::new(values_name);

src/query/ast/src/ast/format/syntax/dml.rs

@@ -84,6 +84,25 @@ fn pretty_source(source: InsertSource) -> RcDoc {
                     .append(RcDoc::text(rest_str.to_string()))
                     .append(RcDoc::text(start.to_string())),
             ),
+        InsertSource::StreamingV2 { settings, start } => RcDoc::text("FILE_FORMAT").append(
+            RcDoc::line()
+                .append(RcDoc::text("FILE_FORMAT_SETTINGS = "))
+                .append(parenthenized(
+                    interweave_comma(settings.iter().map(|(k, v)| {
+                        RcDoc::text(k.to_string())
+                            .append(RcDoc::space())
+                            .append(RcDoc::text("="))
+                            .append(RcDoc::space())
+                            .append(RcDoc::text(format!("{:?}", v)))
+                    }))
+                    .group(),
+                ))
+                .append(
+                    RcDoc::text("start:")
+                        .append(RcDoc::space())
+                        .append(RcDoc::text(start.to_string())),
+                ),
+        ),
         InsertSource::Values { rest_str } => RcDoc::text("VALUES").append(
             RcDoc::line()
                 .nest(NEST_FACTOR)

src/query/ast/src/ast/statements/insert.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::fmt::Display;
 use std::fmt::Formatter;
 
@@ -61,6 +62,10 @@ pub enum InsertSource<'a> {
         rest_str: &'a str,
         start: usize,
     },
+    StreamingV2 {
+        settings: BTreeMap<String, String>,
+        start: usize,
+    },
     Values {
         rest_str: &'a str,
     },
@@ -76,8 +81,13 @@ impl Display for InsertSource<'_> {
                 format,
                 rest_str,
                 start: _,
-            } => {
-                write!(f, "FORMAT {format} {rest_str}")
+            } => write!(f, "FORMAT {format} {rest_str}"),
+            InsertSource::StreamingV2 { settings, start: _ } => {
+                write!(f, " FILE_FORMAT = (")?;
+                for (k, v) in settings.iter() {
+                    write!(f, " {} = '{}'", k, v)?;
+                }
+                write!(f, " )")
             }
             InsertSource::Values { rest_str } => write!(f, "VALUES {rest_str}"),
             InsertSource::Select { query } => write!(f, "{query}"),

src/query/ast/src/parser/statement.rs

@@ -1091,14 +1091,23 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
 pub fn insert_source(i: Input) -> IResult<InsertSource> {
     let streaming = map(
         rule! {
-            FORMAT ~ #ident ~ #rest_str
+                 FORMAT ~ #ident ~ #rest_str
         },
         |(_, format, (rest_str, start))| InsertSource::Streaming {
             format: format.name,
             rest_str,
             start,
         },
     );
+    let streaming_v2 = map(
+        rule! {
+            FILE_FORMAT ~ "=" ~ #options ~ #rest_str
+        },
+        |(_, _, options, (_, start))| InsertSource::StreamingV2 {
+            settings: options,
+            start,
+        },
+    );
     let values = map(
         rule! {
             VALUES ~ #rest_str
@@ -1111,6 +1120,7 @@ pub fn insert_source(i: Input) -> IResult<InsertSource> {
 
     rule!(
         #streaming
+        | #streaming_v2
         | #values
         | #query
     )(i)

src/query/pipeline/sources/src/processors/sources/input_formats/input_context.rs

@@ -29,6 +29,7 @@ use common_formats::ClickhouseFormatType;
 use common_formats::FileFormatOptionsExt;
 use common_formats::FileFormatTypeExt;
 use common_formats::RecordDelimiter;
+use common_meta_types::FileFormatOptions;
 use common_meta_types::StageFileCompression;
 use common_meta_types::StageFileFormatType;
 use common_meta_types::UserStageInfo;
@@ -253,6 +254,48 @@ impl InputContext {
         })
     }
 
+    pub async fn try_create_from_insert_v2(
+        stream_receiver: Receiver<Result<StreamingReadBatch>>,
+        settings: Arc<Settings>,
+        file_format_options: FileFormatOptions,
+        schema: DataSchemaRef,
+        scan_progress: Arc<Progress>,
+        is_multi_part: bool,
+        block_compact_thresholds: BlockCompactThresholds,
+    ) -> Result<Self> {
+        let read_batch_size = settings.get_input_read_buffer_size()? as usize;
+        let format_typ = file_format_options.format.clone();
+        let file_format_options =
+            StageFileFormatType::get_ext_from_stage(file_format_options, &settings)?;
+        let file_format_options = format_typ.final_file_format_options(&file_format_options)?;
+        let format = Self::get_input_format(&format_typ)?;
+        let field_delimiter = file_format_options.get_field_delimiter();
+        let record_delimiter = file_format_options.get_record_delimiter()?;
+        let rows_to_skip = file_format_options.stage.skip_header as usize;
+        let compression = file_format_options.stage.compression;
+
+        let plan = StreamPlan {
+            is_multi_part,
+            compression,
+        };
+
+        Ok(InputContext {
+            format,
+            schema,
+            settings,
+            record_delimiter,
+            read_batch_size,
+            rows_to_skip,
+            field_delimiter,
+            scan_progress,
+            source: InputSource::Stream(Mutex::new(Some(stream_receiver))),
+            plan: InputPlan::StreamingLoad(plan),
+            splits: vec![],
+            block_compact_thresholds,
+            format_options: file_format_options,
+        })
+    }
+
     pub fn num_prefetch_splits(&self) -> Result<usize> {
         Ok(self.settings.get_max_threads()? as usize)
     }

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

@@ -149,6 +149,12 @@ impl Interpreter for InsertInterpreterV2 {
                         .format
                         .exec_stream(input_context.clone(), &mut build_res.main_pipeline)?;
                 }
+                InsertInputSource::StreamingWithFileFormat(_, _, input_context) => {
+                    let input_context = input_context.as_ref().expect("must success").clone();
+                    input_context
+                        .format
+                        .exec_stream(input_context.clone(), &mut build_res.main_pipeline)?;
+                }
                 InsertInputSource::SelectPlan(plan) => {
                     let table1 = table.clone();
                     let (mut select_plan, select_column_bindings) = match plan.as_ref() {
@@ -228,7 +234,8 @@ impl Interpreter for InsertInterpreterV2 {
         }
 
         let append_mode = match &self.plan.source {
-            InsertInputSource::StreamingWithFormat(_, _, _) => AppendMode::Copy,
+            InsertInputSource::StreamingWithFormat(..)
+            | InsertInputSource::StreamingWithFileFormat(..) => AppendMode::Copy,
             _ => AppendMode::Normal,
         };
 

src/query/service/src/servers/http/clickhouse_handler.rs

@@ -290,6 +290,47 @@ pub async fn clickhouse_handler_post(
                 "clickhouse insert with format {:?}, value {}",
                 input_context, *start
             );
+            let compression_alg = input_context.get_compression_alg("").map_err(BadRequest)?;
+            let start = *start;
+            handle = Some(ctx.spawn(async move {
+                gen_batches(
+                    sql,
+                    start,
+                    input_context.read_batch_size,
+                    tx,
+                    compression_alg,
+                )
+                .await
+            }));
+        } else if let InsertInputSource::StreamingWithFileFormat(
+            option_settings,
+            start,
+            input_context_ref,
+        ) = &mut insert.source
+        {
+            let (tx, rx) = tokio::sync::mpsc::channel(2);
+            let to_table = ctx
+                .get_table(&insert.catalog, &insert.database, &insert.table)
+                .await
+                .map_err(InternalServerError)?;
+
+            let input_context = Arc::new(
+                InputContext::try_create_from_insert_v2(
+                    rx,
+                    ctx.get_settings(),
+                    option_settings.clone(),
+                    schema,
+                    ctx.get_scan_progress(),
+                    false,
+                    to_table.get_block_compact_thresholds(),
+                )
+                .await
+                .map_err(InternalServerError)?,
+            );
+
+            *input_context_ref = Some(input_context.clone());
+            info!("clickhouse insert with file_format {:?}", input_context);
+
             let compression_alg = input_context.get_compression_alg("").map_err(BadRequest)?;
             let start = *start;
             handle = Some(ctx.spawn(async move {

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

@@ -100,6 +100,7 @@ pub async fn streaming_load(
         .unwrap_or("");
 
     let settings = context.get_settings();
+
     for (key, value) in req.headers().iter() {
         if settings.has_setting(key.as_str()) {
             let value = value.to_str().map_err(InternalServerError)?;
@@ -135,6 +136,7 @@ pub async fn streaming_load(
                     .await
                     .map_err(InternalServerError)?;
                 let (tx, rx) = tokio::sync::mpsc::channel(2);
+
                 let input_context = Arc::new(
                     InputContext::try_create_from_insert(
                         format.as_str(),
@@ -172,8 +174,63 @@ pub async fn streaming_load(
                     )),
                 }
             }
+            InsertInputSource::StreamingWithFileFormat(
+                option_settings,
+                start,
+                input_context_ref,
+            ) => {
+                let sql_rest = &insert_sql[*start..].trim();
+                if !sql_rest.is_empty() {
+                    return Err(poem::Error::from_string(
+                        "should NOT have data after `FILE_FORMAT` in streaming load.",
+                        StatusCode::BAD_REQUEST,
+                    ));
+                };
+                let to_table = context
+                    .get_table(&insert.catalog, &insert.database, &insert.table)
+                    .await
+                    .map_err(InternalServerError)?;
+                let (tx, rx) = tokio::sync::mpsc::channel(2);
+
+                let input_context = Arc::new(
+                    InputContext::try_create_from_insert_v2(
+                        rx,
+                        context.get_settings(),
+                        option_settings.clone(),
+                        schema,
+                        context.get_scan_progress(),
+                        false,
+                        to_table.get_block_compact_thresholds(),
+                    )
+                    .await
+                    .map_err(InternalServerError)?,
+                );
+                *input_context_ref = Some(input_context.clone());
+                tracing::info!("streaming load with file_format {:?}", input_context);
+
+                let handler = context.spawn(execute_query(context.clone(), plan));
+                let files = read_multi_part(multipart, tx, &input_context).await?;
+
+                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(poem::Error::from_string(
+                        format!("execute fail: {}", cause.message()),
+                        StatusCode::BAD_REQUEST,
+                    )),
+                    Err(_) => Err(poem::Error::from_string(
+                        "Maybe panic.",
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                    )),
+                }
+            }
             _non_supported_source => Err(poem::Error::from_string(
-                "Only supports streaming upload. e.g. INSERT INTO $table FORMAT CSV, got insert ... select.",
+                "Only supports streaming upload. e.g. INSERT INTO $table FILE_FORMAT = (type = 'CSV'), got insert ... select.",
                 StatusCode::BAD_REQUEST,
             )),
         },

src/query/sql/src/planner/binder/insert.rs

@@ -25,6 +25,7 @@ use crate::normalize_identifier;
 use crate::optimizer::optimize;
 use crate::optimizer::OptimizerConfig;
 use crate::optimizer::OptimizerContext;
+use crate::planner::binder::copy::parse_copy_file_format_options;
 use crate::plans::Insert;
 use crate::plans::InsertInputSource;
 use crate::plans::Plan;
@@ -89,6 +90,12 @@ impl<'a> Binder {
                     Ok(InsertInputSource::StreamingWithFormat(format, start, None))
                 }
             }
+            InsertSource::StreamingV2 { settings, start } => {
+                let opts = parse_copy_file_format_options(&settings)?;
+                Ok(InsertInputSource::StreamingWithFileFormat(
+                    opts, start, None,
+                ))
+            }
             InsertSource::Values { rest_str } => {
                 let data = rest_str.trim_end_matches(';').trim_start().to_owned();
                 Ok(InsertInputSource::Values(data))