Merge pull request #7530 from andylokandy/no_start

mergify[bot] · web-flow · commit 69f82ef68f69 · 2022-09-08T14:06:09.000Z
refactor(query): replace recursion for fast-path insert with loop
diff --git a/src/query/ast/src/ast/format/syntax/dml.rs b/src/query/ast/src/ast/format/syntax/dml.rs
@@ -69,17 +69,15 @@ pub(crate) fn pretty_insert(insert_stmt: InsertStmt) -> RcDoc {
 
 fn pretty_source(source: InsertSource) -> RcDoc {
     RcDoc::line().append(match source {
-        InsertSource::Streaming {
-            format, rest_str, ..
-        } => RcDoc::text("FORMAT")
+        InsertSource::Streaming { format, rest_str } => RcDoc::text("FORMAT")
             .append(RcDoc::space())
             .append(RcDoc::text(format))
             .append(
                 RcDoc::line()
                     .nest(NEST_FACTOR)
                     .append(RcDoc::text(rest_str.to_string())),
             ),
-        InsertSource::Values { rest_str, .. } => RcDoc::text("VALUES").append(
+        InsertSource::Values { rest_str } => RcDoc::text("VALUES").append(
             RcDoc::line()
                 .nest(NEST_FACTOR)
                 .append(RcDoc::text(rest_str.to_string())),
diff --git a/src/query/ast/src/ast/statements/insert.rs b/src/query/ast/src/ast/statements/insert.rs
@@ -56,29 +56,18 @@ impl Display for InsertStmt<'_> {
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum InsertSource<'a> {
-    Streaming {
-        format: String,
-        start: usize,
-        rest_str: &'a str,
-    },
-    Values {
-        start: usize,
-        rest_str: &'a str,
-    },
-    Select {
-        query: Box<Query<'a>>,
-    },
+    Streaming { format: String, rest_str: &'a str },
+    Values { rest_str: &'a str },
+    Select { query: Box<Query<'a>> },
 }
 
 impl Display for InsertSource<'_> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            InsertSource::Streaming {
-                format, rest_str, ..
-            } => {
+            InsertSource::Streaming { format, rest_str } => {
                 write!(f, "FORMAT {format} {rest_str}")
             }
-            InsertSource::Values { rest_str, .. } => write!(f, "VALUES {rest_str}"),
+            InsertSource::Values { rest_str } => write!(f, "VALUES {rest_str}"),
             InsertSource::Select { query } => write!(f, "{query}"),
         }
     }
diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -978,32 +978,18 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
 pub fn insert_source(i: Input) -> IResult<InsertSource> {
     let streaming = map(
         rule! {
-            FORMAT ~ #ident ~ #rest_tokens
+            FORMAT ~ #ident ~ #rest_str
         },
-        |(_, format, rest_tokens)| {
-            let rest_str = &rest_tokens[0].source
-                [rest_tokens.first().unwrap().span.start..rest_tokens.last().unwrap().span.end];
-
-            InsertSource::Streaming {
-                format: format.name,
-                start: rest_tokens.first().unwrap().span.start,
-                rest_str,
-            }
+        |(_, format, rest_str)| InsertSource::Streaming {
+            format: format.name,
+            rest_str,
         },
     );
     let values = map(
         rule! {
-            VALUES ~ #rest_tokens
-        },
-        |(_, rest_tokens)| {
-            let rest_str = &rest_tokens[0].source
-                [rest_tokens.first().unwrap().span.start..rest_tokens.last().unwrap().span.end];
-
-            InsertSource::Values {
-                rest_str,
-                start: rest_tokens.first().unwrap().span.start,
-            }
+            VALUES ~ #rest_str
         },
+        |(_, rest_str)| InsertSource::Values { rest_str },
     );
     let query = map(query, |query| InsertSource::Select {
         query: Box::new(query),
@@ -1016,12 +1002,14 @@ pub fn insert_source(i: Input) -> IResult<InsertSource> {
     )(i)
 }
 
-pub fn rest_tokens<'a>(i: Input<'a>) -> IResult<&'a [Token]> {
-    if i.last().map(|token| token.kind) == Some(EOI) {
-        Ok((i.slice(i.len() - 1..), i.slice(..i.len() - 1).0))
-    } else {
-        Ok((i.slice(i.len()..), i.0))
-    }
+pub fn rest_str<'a>(i: Input<'a>) -> IResult<&'a str> {
+    // It's safe to unwrap because input must contain EOI.
+    let first_token = i.0.first().unwrap();
+    let last_token = i.0.last().unwrap();
+    Ok((
+        i.slice((i.len() - 1)..),
+        &first_token.source[first_token.span.start..last_token.span.end],
+    ))
 }
 
 pub fn column_def(i: Input) -> IResult<ColumnDefinition> {
diff --git a/src/query/ast/src/util.rs b/src/query/ast/src/util.rs
@@ -344,7 +344,7 @@ where
             let span = iter_cloned
                 .nth(len - iter.len() - 1)
                 .map(|elem| elem.span)
-                // It's safe to slice one more token because EOI is always added.
+                // It's safe to slice one more token because input must contain EOI.
                 .unwrap_or_else(|| rest.slice(..1));
 
             nom::Err::Error(Error::from_error_kind(span, err_kind))
diff --git a/src/query/ast/tests/it/testdata/statement.txt b/src/query/ast/tests/it/testdata/statement.txt
@@ -4335,7 +4335,6 @@ Insert(
             },
         ],
         source: Values {
-            start: 30,
             rest_str: "(1, 2), (3, 4);",
         },
         overwrite: false,
@@ -4360,7 +4359,6 @@ Insert(
         columns: [],
         source: Streaming {
             format: "json",
-            start: 31,
             rest_str: ";",
         },
         overwrite: false,
diff --git a/src/query/service/src/sql/planner/binder/insert.rs b/src/query/service/src/sql/planner/binder/insert.rs
@@ -108,9 +108,7 @@ impl<'a> Binder {
         };
 
         let input_source: Result<InsertInputSource> = match source.clone() {
-            InsertSource::Streaming {
-                format, rest_str, ..
-            } => {
+            InsertSource::Streaming { format, rest_str } => {
                 self.analyze_stream_format(bind_context, rest_str, Some(format), schema.clone())
                     .await
             }
diff --git a/src/query/service/src/sql/planner/mod.rs b/src/query/service/src/sql/planner/mod.rs
@@ -14,12 +14,10 @@
 
 use std::sync::Arc;
 
-use common_ast::ast::InsertSource;
-use common_ast::ast::Statement;
 use common_ast::parser::parse_sql;
+use common_ast::parser::token::Token;
 use common_ast::parser::token::TokenKind;
 use common_ast::parser::token::Tokenizer;
-use common_ast::parser::tokenize_sql;
 use common_ast::Backtrace;
 use common_exception::Result;
 use parking_lot::RwLock;
@@ -54,7 +52,8 @@ use super::optimizer::OptimizerConfig;
 use super::optimizer::OptimizerContext;
 use crate::sessions::TableContext;
 
-static MAX_TOKEN_FOR_INSERT: usize = 128;
+const PROBE_INSERT_INITIAL_TOKENS: usize = 128;
+const PROBE_INSERT_MAX_TOKENS: usize = 128 * 8;
 
 pub struct Planner {
     ctx: Arc<QueryContext>,
@@ -66,86 +65,73 @@ impl Planner {
     }
 
     pub async fn plan_sql(&mut self, sql: &str) -> Result<(Plan, MetadataRef, Option<String>)> {
-        let mut tokenizer = Tokenizer::new(sql);
-        if let Some(Ok(t)) = tokenizer.next() {
-            if t.kind == TokenKind::INSERT {
-                if let Some(Ok(last_token)) = tokenizer.take(MAX_TOKEN_FOR_INSERT).last() {
-                    match self.plan_sql_inner(sql, Some(last_token.span.start)).await {
-                        Ok(v) => return Ok(v),
-                        Err(err) => {
-                            tracing::warn!(
-                                "Faster parser path for insert failed: {}, start to fallback",
-                                err
-                            );
-                        }
-                    }
-                }
-            }
-        }
-        // fallback to normal plan_sql_inner
-        self.plan_sql_inner(sql, None).await
-    }
-
-    #[async_recursion::async_recursion]
-    async fn plan_sql_inner(
-        &mut self,
-        sql: &str,
-        last_token_index: Option<usize>,
-    ) -> Result<(Plan, MetadataRef, Option<String>)> {
         let settings = self.ctx.get_settings();
         let sql_dialect = settings.get_sql_dialect()?;
 
-        // Step 1: parse SQL text into AST
-        // If last_token_index is set, we will try to parse from the short_sql.
-        // If any error happens, it will fallback to default parser logic.
-        let short_sql = last_token_index.map(|index| &sql[..index]).unwrap_or(sql);
-        let tokens = tokenize_sql(short_sql)?;
-        let backtrace = Backtrace::new();
-
-        let (mut stmt, format) = parse_sql(&tokens, sql_dialect, &backtrace)?;
-        if last_token_index.is_some() {
-            let mut should_fallback = true;
-            if let Statement::Insert(ref mut insert) = stmt {
-                match &mut insert.source {
-                    InsertSource::Streaming {
-                        start, rest_str, ..
-                    } => {
-                        *rest_str = &sql[*start..];
-                        should_fallback = false;
-                    }
-
-                    InsertSource::Values {
-                        start, rest_str, ..
-                    } => {
-                        *rest_str = &sql[*start..];
-                        should_fallback = false;
-                    }
-                    _ => {}
-                }
+        // Step 1: Tokenize the SQL.
+        let mut tokenizer = Tokenizer::new(sql).peekable();
+        let is_insert_stmt = tokenizer
+            .peek()
+            .and_then(|token| Some(token.as_ref().ok()?.kind))
+            == Some(TokenKind::INSERT);
+        // Only tokenize the beginning tokens for `INSERT INTO` statement because it's unnecessary to tokenize tokens for values.
+        //
+        // Stop the tokenizer on unrecognized token because some values inputs (e.g. CSV) may not be recognized by the tokenizer.
+        // See also: https://github.com/datafuselabs/databend/issues/6669
+        let mut tokens: Vec<Token> = if is_insert_stmt {
+            (&mut tokenizer)
+                .take(PROBE_INSERT_INITIAL_TOKENS)
+                .take_while(|token| token.is_ok())
+                .collect::<Result<_>>()
+                .unwrap()
+        } else {
+            (&mut tokenizer).collect::<Result<_>>()?
+        };
+
+        loop {
+            let res = async {
+                // Step 2: Parse the SQL.
+                let backtrace = Backtrace::new();
+                let (stmt, format) = parse_sql(&tokens, sql_dialect, &backtrace)?;
+
+                // Step 3: Bind AST with catalog, and generate a pure logical SExpr
+                let metadata = Arc::new(RwLock::new(Metadata::create()));
+                let name_resolution_ctx = NameResolutionContext::try_from(settings.as_ref())?;
+                let binder = Binder::new(
+                    self.ctx.clone(),
+                    self.ctx.get_catalog_manager()?,
+                    name_resolution_ctx,
+                    metadata.clone(),
+                );
+                let plan = binder.bind(&stmt).await?;
+
+                // Step 4: Optimize the SExpr with optimizers, and generate optimized physical SExpr
+                let opt_ctx = Arc::new(OptimizerContext::new(OptimizerConfig {
+                    enable_distributed_optimization: !self.ctx.get_cluster().is_empty(),
+                }));
+                let optimized_plan = optimize(self.ctx.clone(), opt_ctx, plan)?;
+
+                Ok((optimized_plan, metadata.clone(), format))
             }
-            // fallback to normal plan_sql_inner
-            if should_fallback {
-                return self.plan_sql_inner(sql, None).await;
+            .await;
+
+            if res.is_err() && matches!(tokenizer.peek(), Some(Ok(_))) {
+                // Tokenize more and try again.
+                if tokens.len() < PROBE_INSERT_MAX_TOKENS {
+                    let iter = (&mut tokenizer)
+                        .take(tokens.len() * 2)
+                        .take_while(|token| token.is_ok())
+                        .map(|token| token.unwrap());
+                    tokens.extend(iter);
+                } else {
+                    let iter = (&mut tokenizer)
+                        .take_while(|token| token.is_ok())
+                        .map(|token| token.unwrap());
+                    tokens.extend(iter);
+                };
+            } else {
+                return res;
             }
         }
-
-        // Step 2: bind AST with catalog, and generate a pure logical SExpr
-        let metadata = Arc::new(RwLock::new(Metadata::create()));
-        let name_resolution_ctx = NameResolutionContext::try_from(settings.as_ref())?;
-        let binder = Binder::new(
-            self.ctx.clone(),
-            self.ctx.get_catalog_manager()?,
-            name_resolution_ctx,
-            metadata.clone(),
-        );
-        let plan = binder.bind(&stmt).await?;
-
-        // Step 3: optimize the SExpr with optimizers, and generate optimized physical SExpr
-        let opt_ctx = Arc::new(OptimizerContext::new(OptimizerConfig {
-            enable_distributed_optimization: !self.ctx.get_cluster().is_empty(),
-        }));
-        let optimized_plan = optimize(self.ctx.clone(), opt_ctx, plan)?;
-
-        Ok((optimized_plan, metadata.clone(), format))
     }
 }
diff --git a/tests/suites/0_stateless/14_clickhouse_http_handler/14_0008_tsv_input_format.result b/tests/suites/0_stateless/14_clickhouse_http_handler/14_0008_tsv_input_format.result
@@ -1,5 +1,6 @@
 ---bigint
 5632622125792883430
 ---escape
+"ac	"bc
 1a\nb	c\td
 2a b	cd
diff --git a/tests/suites/0_stateless/14_clickhouse_http_handler/14_0008_tsv_input_format.sh b/tests/suites/0_stateless/14_clickhouse_http_handler/14_0008_tsv_input_format.sh
@@ -8,7 +8,8 @@ insert into t1(a) format TSV  5632622125792883430
 EOF
 
 cat << EOF > /tmp/databend_test_tsv_escape.txt
-insert into t2(a, b) format TSV  1a\nb	c\td
+insert into t2(a, b) format TSV  "ac	"bc
+1a\nb	c\td
 2a\x20b	c\Nd
 EOF
 

-Original file line number
+Diff line change
@@ @@ -1,5 +1,6 @@ @@
 ---bigint
 5632622125792883430
 ---escape
 +"ac	"bc
 a\nb	c\td
 a b	cd