Snowflake: Support SAMPLE clause on subqueries (#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>,
     },
     /// `TABLE(<expr>)[ AS <alias> ]`
     TableFunction {
@@ -2071,6 +2073,7 @@ impl fmt::Display for TableFactor {
                 lateral,
                 subquery,
                 alias,
+                sample,
             } => {
                 if *lateral {
                     write!(f, "LATERAL ")?;
@@ -2083,6 +2086,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::Function {

src/ast/spans.rs

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

src/parser/mod.rs

@@ -15073,6 +15073,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)
@@ -15880,13 +15881,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![]
             }])),
@@ -7863,6 +7864,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()])),
@@ -8871,6 +8873,7 @@ fn lateral_derived() {
             lateral,
             ref subquery,
             alias: Some(ref alias),
+            sample: _,
         } = join.relation
         {
             assert_eq!(lateral_in, lateral);
@@ -9950,6 +9953,7 @@ fn parse_merge() {
                         pipe_operators: vec![],
                     }),
                     alias: table_alias(true, "stg"),
+                    sample: None,
                 }
             );
             assert_eq!(source, source_no_into);

tests/sqlparser_snowflake.rs

@@ -3406,6 +3406,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 @~");