Merge pull request #9123 from youngsofun/stage_table

feat(stage): support `select from @a_stage`

Author: mergify[bot]

Cargo.lock

@@ -2196,6 +2196,31 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 let node = FormatTreeNode::with_children(format_ctx, vec![child]);
                 self.children.push(node);
             }
+            TableReference::Stage {
+                span: _,
+                location,
+                files,
+                alias,
+            } => {
+                let mut children = Vec::new();
+                if !files.is_empty() {
+                    let files = files.join(",");
+                    let files = format!("files = {}", files);
+                    children.push(FormatTreeNode::new(AstFormatContext::new(files)))
+                }
+                let stage_name = format!("Stage {:?}", location);
+                let format_ctx = if let Some(alias) = alias {
+                    AstFormatContext::with_children_alias(
+                        stage_name,
+                        children.len(),
+                        Some(format!("{}", alias)),
+                    )
+                } else {
+                    AstFormatContext::with_children(stage_name, children.len())
+                };
+                let node = FormatTreeNode::with_children(format_ctx, children);
+                self.children.push(node)
+            }
         }
     }
 

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

@@ -314,6 +314,24 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc {
                     .append(RcDoc::text(")")),
                 _ => RcDoc::nil(),
             }),
+        TableReference::Stage {
+            span: _,
+            location,
+            files,
+            alias,
+        } => RcDoc::text(location.to_string())
+            .append(if files.is_empty() {
+                RcDoc::nil()
+            } else {
+                let files = files.join(",");
+                let files = format!("FILES {}", files);
+                RcDoc::text(files)
+            })
+            .append(if let Some(a) = alias {
+                RcDoc::text(format!(" AS {a}"))
+            } else {
+                RcDoc::nil()
+            }),
     }
 }
 

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

@@ -19,6 +19,7 @@ use crate::ast::write_comma_separated_list;
 use crate::ast::write_period_separated_list;
 use crate::ast::Expr;
 use crate::ast::Identifier;
+use crate::ast::StageLocation;
 use crate::parser::token::Token;
 
 /// Root node of a query tree
@@ -175,6 +176,12 @@ pub enum TableReference<'a> {
         span: &'a [Token<'a>],
         join: Join<'a>,
     },
+    Stage {
+        span: &'a [Token<'a>],
+        location: StageLocation,
+        files: Vec<String>,
+        alias: Option<TableAlias<'a>>,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -345,6 +352,21 @@ impl<'a> Display for TableReference<'a> {
                     _ => {}
                 }
             }
+            TableReference::Stage {
+                span: _,
+                location,
+                files,
+                alias,
+            } => {
+                write!(f, "({location})")?;
+                if !files.is_empty() {
+                    let files = files.join(",");
+                    write!(f, " FILES {files}")?;
+                }
+                if let Some(alias) = alias {
+                    write!(f, " AS {alias}")?;
+                }
+            }
         }
         Ok(())
     }

src/query/ast/src/ast/query.rs

@@ -21,6 +21,7 @@ use pratt::Associativity;
 use pratt::PrattParser;
 use pratt::Precedence;
 
+use super::statement::stage_location;
 use crate::ast::*;
 use crate::input::Input;
 use crate::input::WithSpan;
@@ -257,6 +258,11 @@ pub enum TableReferenceElement<'a> {
     // ON expr | USING (ident, ...)
     JoinCondition(JoinCondition<'a>),
     Group(TableReference<'a>),
+    Stage {
+        location: StageLocation,
+        files: Vec<String>,
+        alias: Option<TableAlias<'a>>,
+    },
 }
 
 pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceElement>> {
@@ -319,8 +325,20 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
         |(_, table_ref, _)| TableReferenceElement::Group(table_ref),
     );
 
+    let aliased_stage = map(
+        rule! {
+            #stage_location ~ ( FILES ~ "=" ~ "(" ~ #comma_separated_list0(literal_string) ~ ")")? ~ #table_alias?
+        },
+        |(location, files, alias)| TableReferenceElement::Stage {
+            location,
+            alias,
+            files: files.map(|v| v.3).unwrap_or_default(),
+        },
+    );
+
     let (rest, (span, elem)) = consumed(rule! {
         #subquery
+        | #aliased_stage
         | #table_function
         | #aliased_table
         | #group
@@ -381,6 +399,16 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement<'a>>>> PrattParse
                 subquery,
                 alias,
             },
+            TableReferenceElement::Stage {
+                location,
+                files,
+                alias,
+            } => TableReference::Stage {
+                span: input.span.0,
+                location,
+                files,
+                alias,
+            },
             _ => unreachable!(),
         };
         Ok(table_ref)

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

@@ -262,6 +262,7 @@ pub fn walk_table_reference<'a, V: Visitor<'a>>(
         TableReference::Join { join, .. } => {
             visitor.visit_join(join);
         }
+        TableReference::Stage { .. } => {}
     }
 }
 

src/query/ast/src/visitors/walk.rs

@@ -262,6 +262,7 @@ pub fn walk_table_reference_mut<'a, V: VisitorMut>(
         TableReference::Join { join, .. } => {
             visitor.visit_join(join);
         }
+        TableReference::Stage { .. } => {}
     }
 }
 

src/query/ast/src/visitors/walk_mut.rs

@@ -28,13 +28,16 @@ use crate::plan::PartStatistics;
 use crate::plan::Partitions;
 use crate::plan::Projection;
 use crate::plan::PushDownInfo;
+use crate::plan::StageFileInfo;
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, Eq)]
 pub struct StageTableInfo {
     pub schema: DataSchemaRef,
     pub path: String,
     pub files: Vec<String>,
+    pub pattern: String,
     pub user_stage_info: UserStageInfo,
+    pub files_to_copy: Option<Vec<StageFileInfo>>,
 }
 
 impl StageTableInfo {

src/query/catalog/src/plan/datasource.rs

@@ -19,6 +19,7 @@ mod partition;
 mod partition_statistics;
 mod projection;
 mod pushdown;
+mod stage_file_info;
 
 pub use datasource::*;
 pub use expression::Expression;
@@ -27,3 +28,5 @@ pub use partition::*;
 pub use partition_statistics::PartStatistics;
 pub use projection::Projection;
 pub use pushdown::*;
+pub use stage_file_info::StageFileInfo;
+pub use stage_file_info::StageFileStatus;

src/query/catalog/src/plan/mod.rs

@@ -14,8 +14,6 @@
 
 use std::fmt::Debug;
 
-use common_meta_app::schema::TableInfo;
-
 #[derive(serde::Serialize, serde::Deserialize, PartialEq, Eq, Clone, Debug, Default)]
 pub struct PartStatistics {
     /// Total rows of the query read.
@@ -72,11 +70,11 @@ impl PartStatistics {
         *self = Self::default();
     }
 
-    pub fn get_description(&self, table_info: &TableInfo) -> String {
+    pub fn get_description(&self, table_desc: &str) -> String {
         if self.read_rows > 0 {
             format!(
                 "(Read from {} table, {} Read Rows:{}, Read Bytes:{}, Partitions Scanned:{}, Partitions Total:{})",
-                table_info.desc,
+                table_desc,
                 if self.is_exact {
                     "Exactly"
                 } else {
@@ -88,7 +86,7 @@ impl PartStatistics {
                 self.partitions_total,
             )
         } else {
-            format!("(Read from {} table)", table_info.desc)
+            format!("(Read from {} table)", table_desc)
         }
     }
 }