Snowflake: Support SAMPLE clause on subqueries (apache#2164)

Author: finchxxia

src/ast/query.rs

@@ -1325,6 +1325,8 @@ pub enum TableFactor {
         subquery: Box<Query>,
         /// Optional alias for the derived table.
         alias: Option<TableAlias>,
+        /// Optional table sample modifier
+        sample: Option<TableSampleKind>,
     },
     /// A pass-through query string that is not parsed.
     /// This is useful while building/rewriting queries with a known valid SQL string and to avoid parsing it.
@@ -2077,6 +2079,7 @@ impl fmt::Display for TableFactor {
                 lateral,
                 subquery,
                 alias,
+                sample,
             } => {
                 if *lateral {
                     write!(f, "LATERAL ")?;
@@ -2089,6 +2092,9 @@ impl fmt::Display for TableFactor {
                 if let Some(alias) = alias {
                     write!(f, " {alias}")?;
                 }
+                if let Some(TableSampleKind::AfterTableAlias(sample)) = sample {
+                    write!(f, " {sample}")?;
+                }
                 Ok(())
             }
             TableFactor::PassThroughQuery { query, alias } => {

src/ast/spans.rs

@@ -1922,6 +1922,7 @@ impl Spanned for TableFactor {
                 lateral: _,
                 subquery,
                 alias,
+                sample: _,
             } => subquery
                 .span()
                 .union_opt(&alias.as_ref().map(|alias| alias.span())),

src/parser/mod.rs

@@ -15105,6 +15105,7 @@ impl<'a> Parser<'a> {
                     pipe_operators: vec![],
                 }),
                 alias,
+                sample: None,
             })
         } else if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect)
             && self.parse_keyword(Keyword::UNNEST)
@@ -15912,13 +15913,20 @@ impl<'a> Parser<'a> {
         let subquery = self.parse_query()?;
         self.expect_token(&Token::RParen)?;
         let alias = self.maybe_parse_table_alias()?;
+
+        // Parse optional SAMPLE clause after alias
+        let sample = self
+            .maybe_parse_table_sample()?
+            .map(TableSampleKind::AfterTableAlias);
+
         Ok(TableFactor::Derived {
             lateral: match lateral {
                 Lateral => true,
                 NotLateral => false,
             },
             subquery,
             alias,
+            sample,
         })
     }
 

tests/sqlparser_common.rs

@@ -512,7 +512,8 @@ fn parse_update_set_from() {
                         format_clause: None,
                         pipe_operators: vec![],
                     }),
-                    alias: table_alias(true, "t2")
+                    alias: table_alias(true, "t2"),
+                    sample: None,
                 },
                 joins: vec![]
             }])),
@@ -7864,6 +7865,7 @@ fn parse_derived_tables() {
                     lateral: false,
                     subquery: Box::new(verified_query("(SELECT 1) UNION (SELECT 2)")),
                     alias: table_alias(true, "t1"),
+                    sample: None,
                 },
                 joins: vec![Join {
                     relation: table_from_name(ObjectName::from(vec!["t2".into()])),
@@ -8872,6 +8874,7 @@ fn lateral_derived() {
             lateral,
             ref subquery,
             alias: Some(ref alias),
+            sample: _,
         } = join.relation
         {
             assert_eq!(lateral_in, lateral);
@@ -9951,6 +9954,7 @@ fn parse_merge() {
                         pipe_operators: vec![],
                     }),
                     alias: table_alias(true, "stg"),
+                    sample: None,
                 }
             );
             assert_eq!(source, source_no_into);

tests/sqlparser_snowflake.rs

@@ -3538,6 +3538,23 @@ fn test_table_sample() {
     snowflake_and_generic().verified_stmt("SELECT id FROM mytable TABLESAMPLE (10) SEED (1)");
 }
 
+#[test]
+fn test_subquery_sample() {
+    // Test SAMPLE clause on subqueries (derived tables)
+    snowflake_and_generic().verified_stmt("SELECT * FROM (SELECT * FROM mytable) SAMPLE (10)");
+    snowflake_and_generic()
+        .verified_stmt("SELECT * FROM (SELECT * FROM mytable) SAMPLE (10000 ROWS)");
+    snowflake_and_generic()
+        .verified_stmt("SELECT * FROM (SELECT * FROM mytable) AS t SAMPLE (50 PERCENT)");
+    // Nested subquery with SAMPLE
+    snowflake_and_generic().verified_stmt(
+        "SELECT * FROM (SELECT * FROM (SELECT report_from FROM mytable) SAMPLE (10000 ROWS)) AS anon_1",
+    );
+    // SAMPLE with SEED on subquery
+    snowflake_and_generic()
+        .verified_stmt("SELECT * FROM (SELECT * FROM mytable) SAMPLE (10) SEED (42)");
+}
+
 #[test]
 fn parse_ls_and_rm() {
     snowflake().one_statement_parses_to("LS @~", "LIST @~");