feat: stage file pattern match the whole string after prefix. (#12935)

Author: youngsofun

src/common/storage/src/stage.rs

@@ -92,7 +92,7 @@ pub struct StageFilesInfo {
 impl StageFilesInfo {
     fn get_pattern(&self) -> Result<Option<Regex>> {
         match &self.pattern {
-            Some(pattern) => match Regex::new(pattern) {
+            Some(pattern) => match Regex::new(&format!("^{pattern}$")) {
                 Ok(r) => Ok(Some(r)),
                 Err(e) => Err(ErrorCode::SyntaxException(format!(
                     "Pattern format invalid, got:{}, error:{:?}",
@@ -206,6 +206,7 @@ impl StageFilesInfo {
         first_only: bool,
         max_files: usize,
     ) -> Result<Vec<StageFileInfo>> {
+        let prefix_len = if path == "/" { 0 } else { path.len() };
         let root_meta = operator.stat(path).await;
         match root_meta {
             Ok(meta) => match meta.mode() {
@@ -233,7 +234,7 @@ impl StageFilesInfo {
         let mut limit: usize = 0;
         while let Some(obj) = list.try_next().await? {
             let meta = operator.metadata(&obj, StageFileInfo::meta_query()).await?;
-            if check_file(obj.path(), meta.mode(), &pattern) {
+            if check_file(&obj.path()[prefix_len..], meta.mode(), &pattern) {
                 files.push(StageFileInfo::new(obj.path().to_string(), &meta));
                 if first_only {
                     return Ok(files);
@@ -263,6 +264,7 @@ fn blocking_list_files_with_pattern(
     first_only: bool,
     max_files: usize,
 ) -> Result<Vec<StageFileInfo>> {
+    let prefix_len = if path == "/" { 0 } else { path.len() };
     let operator = operator.blocking();
 
     let root_meta = operator.stat(path);
@@ -293,7 +295,7 @@ fn blocking_list_files_with_pattern(
     for obj in list {
         let obj = obj?;
         let meta = operator.metadata(&obj, StageFileInfo::meta_query())?;
-        if check_file(obj.path(), meta.mode(), &pattern) {
+        if check_file(&obj.path()[prefix_len..], meta.mode(), &pattern) {
             files.push(StageFileInfo::new(obj.path().to_string(), &meta));
             if first_only {
                 return Ok(files);

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

@@ -2179,10 +2179,10 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         self.children.push(node);
     }
 
-    fn visit_list_stage(&mut self, location: &'ast str, pattern: &'ast str) {
+    fn visit_list_stage(&mut self, location: &'ast str, pattern: &'ast Option<String>) {
         let location_format_ctx = AstFormatContext::new(format!("Location {}", location));
         let location_child = FormatTreeNode::new(location_format_ctx);
-        let pattern_format_ctx = AstFormatContext::new(format!("Pattern {}", pattern));
+        let pattern_format_ctx = AstFormatContext::new(format!("Pattern {:?}", pattern));
         let pattern_child = FormatTreeNode::new(pattern_format_ctx);
 
         let name = "ListStage".to_string();

src/query/ast/src/ast/statements/statement.rs

@@ -186,7 +186,7 @@ pub enum Statement {
     },
     ListStage {
         location: String,
-        pattern: String,
+        pattern: Option<String>,
     },
 
     // UserDefinedFileFormat
@@ -452,7 +452,7 @@ impl Display for Statement {
             Statement::AlterUDF(stmt) => write!(f, "{stmt}")?,
             Statement::ListStage { location, pattern } => {
                 write!(f, "LIST @{location}")?;
-                if !pattern.is_empty() {
+                if let Some(pattern) = pattern {
                     write!(f, " PATTERN = '{pattern}'")?;
                 }
             }

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

@@ -1054,7 +1054,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
         },
         |(_, location, opt_pattern)| Statement::ListStage {
             location,
-            pattern: opt_pattern.map(|v| v.2).unwrap_or_default(),
+            pattern: opt_pattern.map(|v| v.2),
         },
     );
 

src/query/ast/src/visitors/visitor.rs

@@ -531,7 +531,7 @@ pub trait Visitor<'ast>: Sized {
 
     fn visit_remove_stage(&mut self, _location: &'ast str, _pattern: &'ast str) {}
 
-    fn visit_list_stage(&mut self, _location: &'ast str, _pattern: &'ast str) {}
+    fn visit_list_stage(&mut self, _location: &'ast str, _pattern: &'ast Option<String>) {}
 
     fn visit_create_file_format(
         &mut self,

src/query/ast/src/visitors/visitor_mut.rs

@@ -546,7 +546,7 @@ pub trait VisitorMut: Sized {
 
     fn visit_remove_stage(&mut self, _location: &mut String, _pattern: &mut String) {}
 
-    fn visit_list_stage(&mut self, _location: &mut String, _pattern: &mut String) {}
+    fn visit_list_stage(&mut self, _location: &mut String, _pattern: &mut Option<String>) {}
 
     fn visit_create_file_format(
         &mut self,

src/query/ast/tests/it/testdata/statement.txt

@@ -7120,7 +7120,7 @@ LIST @stage_a
 ---------- AST ------------
 ListStage {
     location: "stage_a",
-    pattern: "",
+    pattern: None,
 }
 
 
@@ -7131,7 +7131,7 @@ LIST @~
 ---------- AST ------------
 ListStage {
     location: "~",
-    pattern: "",
+    pattern: None,
 }
 
 

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

@@ -346,7 +346,14 @@ impl<'a> Binder {
 
             // Stages
             Statement::ShowStages => self.bind_rewrite_to_query(bind_context, "SELECT name, stage_type, number_of_files, creator, comment FROM system.stages ORDER BY name", RewriteKind::ShowStages).await?,
-            Statement::ListStage { location, pattern } => self.bind_rewrite_to_query(bind_context, format!("SELECT * FROM LIST_STAGE(location => '@{location}', pattern => '{pattern}')").as_str(), RewriteKind::ListStage).await?,
+            Statement::ListStage { location, pattern } => {
+                let pattern = if let Some(pattern) = pattern {
+                    format!(", pattern => '{pattern}'")
+                } else {
+                    "".to_string()
+                };
+                self.bind_rewrite_to_query(bind_context, format!("SELECT * FROM LIST_STAGE(location => '@{location}'{pattern})").as_str(), RewriteKind::ListStage).await?
+            },
             Statement::DescribeStage { stage_name } => self.bind_rewrite_to_query(bind_context, format!("SELECT * FROM system.stages WHERE name = '{stage_name}'").as_str(), RewriteKind::DescribeStage).await?,
             Statement::CreateStage(stmt) => self.bind_create_stage(stmt).await?,
             Statement::DropStage {

tests/sqllogictests/suites/stage/list_stage

@@ -8,12 +8,12 @@ parquet/multi_page/multi_page_3.parquet 4020 NULL
 parquet/multi_page/multi_page_4.parquet 6636 NULL
 
 query 
-select name, size, creator from list_stage(location => '@data/parquet/', pattern => 'complex[.]*')
+select name, size, creator from list_stage(location => '@data/parquet/', pattern => 'complex[.].*')
 ----
 parquet/complex.parquet 92762 NULL
 
 query 
-select name, size, creator from list_stage(location => '@data/', pattern => 'parquet/complex[.]*')
+select name, size, creator from list_stage(location => '@data/', pattern => 'parquet/complex[.].*')
 ----
 parquet/complex.parquet 92762 NULL
 

tests/sqllogictests/suites/stage/options/pattern

@@ -0,0 +1,23 @@
+# the following 2 cases show that `pattern` only matching sub path (or suffix) after the 'parquet/' prefix
+# wrong case
+query
+select name from list_stage(location => '@data/parquet/', pattern => 'parquet/.*_page_1.*') order by name
+----
+
+# right case
+query
+select name from list_stage(location => '@data/parquet/', pattern => 'multi_page/.*_page_1.*') order by name
+----
+parquet/multi_page/multi_page_1.parquet
+
+
+# the following 2 cases show that `pattern` need to matching match the whole suffix, it is in fact '%<pattern>$'
+# wrong case
+query
+select name from list_stage(location => '@data/parquet/', pattern => 'multi_page_1') order by name
+----
+
+query
+select name from list_stage(location => '@data/parquet/', pattern => '.*multi_page_1.*') order by name
+----
+parquet/multi_page/multi_page_1.parquet