Implement tablesample

Author: hendrikmakait

src/ast/query.rs

@@ -104,7 +104,7 @@ impl fmt::Display for Query {
             format.fmt(f)?;
         }
         for pipe_operator in &self.pipe_operators {
-            f.write_str(" |> ")?;
+            f.write_str(" |>")?;
             pipe_operator.fmt(f)?;
         }
         Ok(())
@@ -2690,22 +2690,22 @@ impl fmt::Display for PipeOperator {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             PipeOperator::Select { exprs } => {
-                write!(f, "SELECT {}", display_comma_separated(exprs.as_slice()))
+                write!(f, " SELECT {}", display_comma_separated(exprs.as_slice()))
             }
             PipeOperator::Extend { exprs } => {
-                write!(f, "EXTEND {}", display_comma_separated(exprs.as_slice()))
+                write!(f, " EXTEND {}", display_comma_separated(exprs.as_slice()))
             }
             PipeOperator::Set { assignments } => {
-                write!(f, "SET {}", display_comma_separated(assignments.as_slice()))
+                write!(f, " SET {}", display_comma_separated(assignments.as_slice()))
             }
             PipeOperator::Drop { columns } => {
-                write!(f, "DROP {}", display_comma_separated(columns.as_slice()))
+                write!(f, " DROP {}", display_comma_separated(columns.as_slice()))
             }
             PipeOperator::As { alias } => {
-                write!(f, "AS {}", alias)
+                write!(f, " AS {}", alias)
             }
             PipeOperator::Limit { expr, offset } => {
-                write!(f, "LIMIT {}", expr)?;
+                write!(f, " LIMIT {}", expr)?;
                 if let Some(offset) = offset {
                     write!(f, " OFFSET {}", offset)?;
                 }
@@ -2715,7 +2715,7 @@ impl fmt::Display for PipeOperator {
                 full_table_exprs,
                 group_by_expr,
             } => {
-                write!(f, "AGGREGATE")?;
+                write!(f, " AGGREGATE")?;
                 if !full_table_exprs.is_empty() {
                     write!(
                         f,
@@ -2730,14 +2730,14 @@ impl fmt::Display for PipeOperator {
             }
 
             PipeOperator::Where { expr } => {
-                write!(f, "WHERE {}", expr)
+                write!(f, " WHERE {}", expr)
             }
             PipeOperator::OrderBy { exprs } => {
-                write!(f, "ORDER BY {}", display_comma_separated(exprs.as_slice()))
+                write!(f, " ORDER BY {}", display_comma_separated(exprs.as_slice()))
             }
 
             PipeOperator::TableSample { sample } => {
-                write!(f, " {}", sample)
+                write!(f, "{}", sample)
             }
         }
     }

src/parser/mod.rs

@@ -11111,10 +11111,8 @@ impl<'a> Parser<'a> {
                     pipe_operators.push(PipeOperator::OrderBy { exprs })
                 }
                 Keyword::TABLESAMPLE => {
-                    if let Some(sample) = self.maybe_parse_table_sample()? {
-                        pipe_operators.push(PipeOperator::TableSample { sample });
-                    };
-
+                    let sample = self.parse_table_sample(TableSampleModifier::TableSample)?;
+                    pipe_operators.push(PipeOperator::TableSample { sample });
                 }
                 unhandled => {
                     return Err(ParserError::ParserError(format!(
@@ -12760,8 +12758,11 @@ impl<'a> Parser<'a> {
         } else {
             return Ok(None);
         };
+        self.parse_table_sample(modifier).map(|sample| Some(sample))       
+    }
 
-        let name = match self.parse_one_of_keywords(&[
+    fn parse_table_sample(&mut self, modifier: TableSampleModifier ) -> Result<Box<TableSample>, ParserError> {
+let name = match self.parse_one_of_keywords(&[
             Keyword::BERNOULLI,
             Keyword::ROW,
             Keyword::SYSTEM,
@@ -12842,14 +12843,14 @@ impl<'a> Parser<'a> {
             None
         };
 
-        Ok(Some(Box::new(TableSample {
+        Ok(Box::new(TableSample {
             modifier,
             name,
             quantity,
             seed,
             bucket,
             offset,
-        })))
+        }))
     }
 
     fn parse_table_sample_seed(

tests/sqlparser_common.rs

@@ -15156,10 +15156,10 @@ fn parse_pipeline_operator() {
     dialects.verified_stmt("SELECT * FROM users |> ORDER BY id DESC, name ASC");
 
     // tablesample pipe operator
-    dialects.verified_stmt("SELECT * FROM tbl AS t TABLESAMPLE BERNOULLI (50)");
-    dialects.verified_stmt("SELECT * FROM tbl AS t TABLESAMPLE SYSTEM (50)");
+    dialects.verified_stmt("SELECT * FROM tbl |> TABLESAMPLE BERNOULLI (50)");
+    dialects.verified_stmt("SELECT * FROM tbl |> TABLESAMPLE SYSTEM (50)");
     // TODO: Technically, REPEATABLE is not available in BigQuery, but it is used with TABLESAMPLE in other dialects
-    dialects.verified_stmt("SELECT * FROM tbl AS t TABLESAMPLE SYSTEM (50) REPEATABLE (10)");
+    dialects.verified_stmt("SELECT * FROM tbl |> TABLESAMPLE SYSTEM (50) REPEATABLE (10)");
 
     // many pipes
     dialects.verified_stmt(