Code review comments

Author: yoavcloud

src/ast/mod.rs

@@ -69,8 +69,8 @@ pub use self::query::{
     OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
     SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
-    TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableSample,
-    TableSampleBernoulli, TableSampleBucket, TableSampleImplicit, TableSampleSystem,
+    TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableSampleBernoulli,
+    TableSampleBucket, TableSampleImplicit, TableSampleKind, TableSampleMethod, TableSampleSystem,
     TableSampleUnit, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode, Values,
     WildcardAdditionalOptions, With, WithFill,
 };

src/ast/query.rs

@@ -1004,10 +1004,7 @@ pub enum TableFactor {
         json_path: Option<JsonPath>,
         /// Optional table sample modifier
         /// See: <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#sample-clause>
-        sample: Option<Box<TableSample>>,
-        /// Position of the table sample modifier in the table factor. Default is after the table alias
-        /// e.g. `SELECT * FROM tbl t TABLESAMPLE (10 ROWS)`. See `Dialect::supports_table_sample_before_alias`.
-        sample_before_alias: bool,
+        sample: Option<TableSampleKind>,
     },
     Derived {
         lateral: bool,
@@ -1156,7 +1153,19 @@ pub enum TableFactor {
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub enum TableSample {
+
+pub enum TableSampleKind {
+    /// Table sample located before the table alias option
+    BeforeTableAlias(Box<TableSampleMethod>),
+    /// Table sample located after the table alias option
+    AfterTableAlias(Box<TableSampleMethod>),
+}
+
+/// The table sample method options
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableSampleMethod {
     Bernoulli(TableSampleBernoulli),
     System(TableSampleSystem),
     Bucket(TableSampleBucket),
@@ -1234,11 +1243,11 @@ impl fmt::Display for TableSampleImplicit {
     }
 }
 
-impl fmt::Display for TableSample {
+impl fmt::Display for TableSampleMethod {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, " TABLESAMPLE")?;
         match self {
-            TableSample::Bernoulli(sample) => {
+            TableSampleMethod::Bernoulli(sample) => {
                 write!(f, " BERNOULLI (")?;
                 if let Some(probability) = &sample.probability {
                     write!(f, "{})", probability)?;
@@ -1250,16 +1259,16 @@ impl fmt::Display for TableSample {
                     write!(f, ")")?;
                 }
             }
-            TableSample::System(sample) => {
+            TableSampleMethod::System(sample) => {
                 write!(f, " SYSTEM ({})", sample.probability)?;
                 if let Some(repeatable) = &sample.repeatable {
                     write!(f, " REPEATABLE ({})", repeatable)?;
                 }
             }
-            TableSample::Bucket(sample) => {
+            TableSampleMethod::Bucket(sample) => {
                 write!(f, " ({})", sample)?;
             }
-            TableSample::Implicit(sample) => {
+            TableSampleMethod::Implicit(sample) => {
                 write!(f, " ({})", sample)?;
             }
         }
@@ -1526,7 +1535,6 @@ impl fmt::Display for TableFactor {
                 with_ordinality,
                 json_path,
                 sample,
-                sample_before_alias,
             } => {
                 write!(f, "{name}")?;
                 if let Some(json_path) = json_path {
@@ -1549,7 +1557,7 @@ impl fmt::Display for TableFactor {
                 if *with_ordinality {
                     write!(f, " WITH ORDINALITY")?;
                 }
-                if let (Some(sample), true) = (sample, sample_before_alias) {
+                if let Some(TableSampleKind::BeforeTableAlias(sample)) = sample {
                     write!(f, "{sample}")?;
                 }
                 if let Some(alias) = alias {
@@ -1561,7 +1569,7 @@ impl fmt::Display for TableFactor {
                 if let Some(version) = version {
                     write!(f, "{version}")?;
                 }
-                if let (Some(sample), false) = (sample, sample_before_alias) {
+                if let Some(TableSampleKind::AfterTableAlias(sample)) = sample {
                     write!(f, "{sample}")?;
                 }
                 Ok(())

src/ast/spans.rs

@@ -1700,7 +1700,6 @@ impl Spanned for TableFactor {
                 partitions: _,
                 json_path: _,
                 sample: _,
-                sample_before_alias: _,
             } => union_spans(
                 name.0
                     .iter()

src/dialect/hive.rs

@@ -68,7 +68,7 @@ impl Dialect for HiveDialect {
     }
 
     /// See Hive <https://cwiki.apache.org/confluence/display/hive/languagemanual+sampling>
-    fn supports_implicit_table_sample(&self) -> bool {
+    fn supports_implicit_table_sample_method(&self) -> bool {
         true
     }
 }

src/dialect/mod.rs

@@ -709,15 +709,23 @@ pub trait Dialect: Debug + Any {
     }
 
     /// Returns true if this dialect supports the `TABLESAMPLE` option
-    /// before the table alias option.
+    /// before the table alias option. For example:
+    ///
+    /// Table sample before alias: `SELECT * FROM tbl AS t TABLESAMPLE (10)`
+    /// Table sample after alias: `SELECT * FROM tbl TABLESAMPLE (10) AS t`
+    ///
     /// <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#_7_6_table_reference>
     fn supports_table_sample_before_alias(&self) -> bool {
         false
     }
 
-    /// Returns true if this dialect support not specifying a table sample method.
+    /// Returns true if this dialect support not specifying a table sample method. For example:
+    ///
+    /// Implicit table sample method: `SELECT * FROM tbl TABLESAMPLE (10)`
+    /// Explicit table sample method: `SELECT * FROM tbl TABLESAMPLE BERNOULLI (10)`
+    ///
     /// <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#sample-clause>
-    fn supports_implicit_table_sample(&self) -> bool {
+    fn supports_implicit_table_sample_method(&self) -> bool {
         false
     }
 }

src/parser/mod.rs

@@ -10595,13 +10595,9 @@ impl<'a> Parser<'a> {
             let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
 
             let mut sample = None;
-            let mut sample_before_alias = false;
             if self.dialect.supports_table_sample_before_alias() {
-                sample = self.maybe_parse_table_sample()?;
-                if sample.is_some() {
-                    // No need to modify the default is no sample option
-                    // exists on the statement
-                    sample_before_alias = true;
+                if let Some(parsed_sample) = self.maybe_parse_table_sample()? {
+                    sample = Some(TableSampleKind::BeforeTableAlias(parsed_sample));
                 }
             }
 
@@ -10620,8 +10616,9 @@ impl<'a> Parser<'a> {
             };
 
             if !self.dialect.supports_table_sample_before_alias() {
-                sample = self.maybe_parse_table_sample()?;
-                sample_before_alias = false;
+                if let Some(parsed_sample) = self.maybe_parse_table_sample()? {
+                    sample = Some(TableSampleKind::AfterTableAlias(parsed_sample));
+                }
             }
 
             let mut table = TableFactor::Table {
@@ -10634,7 +10631,6 @@ impl<'a> Parser<'a> {
                 with_ordinality,
                 json_path,
                 sample,
-                sample_before_alias,
             };
 
             while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {
@@ -10655,14 +10651,15 @@ impl<'a> Parser<'a> {
         }
     }
 
-    fn maybe_parse_table_sample(&mut self) -> Result<Option<Box<TableSample>>, ParserError> {
+    fn maybe_parse_table_sample(&mut self) -> Result<Option<Box<TableSampleMethod>>, ParserError> {
         if self
             .parse_one_of_keywords(&[Keyword::SAMPLE, Keyword::TABLESAMPLE])
             .is_none()
         {
             return Ok(None);
         }
 
+        // Try to parse based on an explicit table sample method keyword
         let sample = if self
             .parse_one_of_keywords(&[Keyword::BERNOULLI, Keyword::ROW])
             .is_some()
@@ -10678,7 +10675,7 @@ impl<'a> Parser<'a> {
                 (Some(expr), None, None)
             };
             self.expect_token(&Token::RParen)?;
-            TableSample::Bernoulli(TableSampleBernoulli {
+            TableSampleMethod::Bernoulli(TableSampleBernoulli {
                 probability,
                 value,
                 unit,
@@ -10701,10 +10698,11 @@ impl<'a> Parser<'a> {
             } else {
                 None
             };
-            TableSample::System(TableSampleSystem {
+            TableSampleMethod::System(TableSampleSystem {
                 probability,
                 repeatable: seed,
             })
+        // Try to parse without an explicit table sample method keyword
         } else if self.peek_token().token == Token::LParen {
             self.expect_token(&Token::LParen)?;
             if self.parse_keyword(Keyword::BUCKET) {
@@ -10717,10 +10715,10 @@ impl<'a> Parser<'a> {
                     None
                 };
                 self.expect_token(&Token::RParen)?;
-                TableSample::Bucket(TableSampleBucket { bucket, total, on })
+                TableSampleMethod::Bucket(TableSampleBucket { bucket, total, on })
             } else {
-                let value = match self.try_parse(|p| p.parse_number_value()) {
-                    Ok(num) => num,
+                let value = match self.maybe_parse(|p| p.parse_number_value()) {
+                    Ok(Some(num)) => num,
                     _ => {
                         if let Token::Word(w) = self.next_token().token {
                             Value::Placeholder(w.value)
@@ -10733,10 +10731,10 @@ impl<'a> Parser<'a> {
                     }
                 };
                 if self.peek_token().token == Token::RParen
-                    && !self.dialect.supports_implicit_table_sample()
+                    && !self.dialect.supports_implicit_table_sample_method()
                 {
                     self.expect_token(&Token::RParen)?;
-                    TableSample::Bernoulli(TableSampleBernoulli {
+                    TableSampleMethod::Bernoulli(TableSampleBernoulli {
                         probability: Some(Expr::Value(value)),
                         unit: None,
                         value: None,
@@ -10750,7 +10748,7 @@ impl<'a> Parser<'a> {
                         None
                     };
                     self.expect_token(&Token::RParen)?;
-                    TableSample::Implicit(TableSampleImplicit { value, unit })
+                    TableSampleMethod::Implicit(TableSampleImplicit { value, unit })
                 }
             }
         } else {

src/test_utils.rs

@@ -347,7 +347,6 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         with_ordinality: false,
         json_path: None,
         sample: None,
-        sample_before_alias: false,
     }
 }
 
@@ -362,7 +361,6 @@ pub fn table_from_name(name: ObjectName) -> TableFactor {
         with_ordinality: false,
         json_path: None,
         sample: None,
-        sample_before_alias: false,
     }
 }
 
@@ -380,7 +378,6 @@ pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> Ta
         with_ordinality: false,
         json_path: None,
         sample: None,
-        sample_before_alias: false,
     }
 }
 

tests/sqlparser_bigquery.rs

@@ -1545,7 +1545,6 @@ fn parse_table_time_travel() {
                 with_ordinality: false,
                 json_path: None,
                 sample: None,
-                sample_before_alias: false,
             },
             joins: vec![]
         },]
@@ -1646,7 +1645,6 @@ fn parse_merge() {
                     with_ordinality: false,
                     json_path: None,
                     sample: None,
-                    sample_before_alias: false,
                 },
                 table
             );
@@ -1664,7 +1662,6 @@ fn parse_merge() {
                     with_ordinality: false,
                     json_path: None,
                     sample: None,
-                    sample_before_alias: false,
                 },
                 source
             );