Add support for table sample, initial commit

Author: yoavcloud

src/ast/mod.rs

@@ -69,8 +69,10 @@ 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, TableVersion, TableWithJoins,
-    Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill,
+    TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableSample,
+    TableSampleBernoulli, TableSampleBucket, TableSampleImplicit, TableSampleSystem,
+    TableSampleUnit, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode, Values,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 pub use self::trigger::{

src/ast/query.rs

@@ -1002,6 +1002,12 @@ pub enum TableFactor {
         partitions: Vec<Ident>,
         /// Optional PartiQL JsonPath: <https://partiql.org/dql/from.html>
         json_path: Option<JsonPath>,
+        /// Optional table sample modifier
+        /// See: <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#sample-clause>
+        sample: Option<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,
     },
     Derived {
         lateral: bool,
@@ -1146,6 +1152,121 @@ pub enum TableFactor {
     },
 }
 
+/// The table sample modifier 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 TableSample {
+    Bernoulli(TableSampleBernoulli),
+    System(TableSampleSystem),
+    Bucket(TableSampleBucket),
+    Implicit(TableSampleImplicit),
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableSampleBernoulli {
+    pub probability: Option<Expr>,
+    pub value: Option<Expr>,
+    pub unit: Option<TableSampleUnit>,
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableSampleSystem {
+    pub probability: Expr,
+    pub seed: Option<Expr>,
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableSampleUnit {
+    Rows,
+    Percent,
+}
+
+impl fmt::Display for TableSampleUnit {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            TableSampleUnit::Percent => write!(f, "PERCENT"),
+            TableSampleUnit::Rows => write!(f, "ROWS"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableSampleBucket {
+    pub bucket: Value,
+    pub total: Value,
+    pub on: Option<Expr>,
+}
+
+impl fmt::Display for TableSampleBucket {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "BUCKET {} OUT OF {}", self.bucket, self.total)?;
+        if let Some(on) = &self.on {
+            write!(f, " ON {}", on)?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct TableSampleImplicit {
+    pub value: Value,
+    pub unit: Option<TableSampleUnit>,
+}
+
+impl fmt::Display for TableSampleImplicit {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.value)?;
+        if let Some(unit) = &self.unit {
+            write!(f, " {}", unit)?;
+        }
+        Ok(())
+    }
+}
+
+impl fmt::Display for TableSample {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, " TABLESAMPLE")?;
+        match self {
+            TableSample::Bernoulli(sample) => {
+                write!(f, " BERNOULLI (")?;
+                if let Some(probability) = &sample.probability {
+                    write!(f, "{})", probability)?;
+                } else if let Some(value) = &sample.value {
+                    write!(f, "{}", value)?;
+                    if let Some(unit) = &sample.unit {
+                        write!(f, " {}", unit)?;
+                    }
+                    write!(f, ")")?;
+                }
+            }
+            TableSample::System(sample) => {
+                write!(f, " SYSTEM ({})", sample.probability)?;
+                if let Some(seed) = &sample.seed {
+                    write!(f, " SEED ({})", seed)?;
+                }
+            }
+            TableSample::Bucket(sample) => {
+                write!(f, " ({})", sample)?;
+            }
+            TableSample::Implicit(sample) => {
+                write!(f, " ({})", sample)?;
+            }
+        }
+        Ok(())
+    }
+}
+
 /// The source of values in a `PIVOT` operation.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1404,6 +1525,8 @@ impl fmt::Display for TableFactor {
                 partitions,
                 with_ordinality,
                 json_path,
+                sample,
+                sample_before_alias,
             } => {
                 write!(f, "{name}")?;
                 if let Some(json_path) = json_path {
@@ -1426,6 +1549,9 @@ impl fmt::Display for TableFactor {
                 if *with_ordinality {
                     write!(f, " WITH ORDINALITY")?;
                 }
+                if let (Some(sample), true) = (sample, sample_before_alias) {
+                    write!(f, "{sample}")?;
+                }
                 if let Some(alias) = alias {
                     write!(f, " AS {alias}")?;
                 }
@@ -1435,6 +1561,9 @@ impl fmt::Display for TableFactor {
                 if let Some(version) = version {
                     write!(f, "{version}")?;
                 }
+                if let (Some(sample), false) = (sample, sample_before_alias) {
+                    write!(f, "{sample}")?;
+                }
                 Ok(())
             }
             TableFactor::Derived {

src/ast/spans.rs

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

src/dialect/hive.rs

@@ -61,4 +61,9 @@ impl Dialect for HiveDialect {
     fn supports_load_data(&self) -> bool {
         true
     }
+
+    /// See Hive <https://cwiki.apache.org/confluence/display/hive/languagemanual+sampling>
+    fn supports_table_sample_before_alias(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -707,6 +707,13 @@ pub trait Dialect: Debug + Any {
     fn is_reserved_for_identifier(&self, kw: Keyword) -> bool {
         keywords::RESERVED_FOR_IDENTIFIER.contains(&kw)
     }
+
+    /// Returns true if the dialect supports the `TABLESAMPLE` option
+    /// before the table alias option.
+    /// <https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#_7_6_table_reference>
+    fn supports_table_sample_before_alias(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/keywords.rs

@@ -120,6 +120,7 @@ define_keywords!(
     BEGIN,
     BEGIN_FRAME,
     BEGIN_PARTITION,
+    BERNOULLI,
     BETWEEN,
     BIGDECIMAL,
     BIGINT,
@@ -128,12 +129,14 @@ define_keywords!(
     BINDING,
     BIT,
     BLOB,
+    BLOCK,
     BLOOMFILTER,
     BOOL,
     BOOLEAN,
     BOTH,
     BROWSE,
     BTREE,
+    BUCKET,
     BUCKETS,
     BY,
     BYPASSRLS,
@@ -679,6 +682,7 @@ define_keywords!(
     RUN,
     SAFE,
     SAFE_CAST,
+    SAMPLE,
     SAVEPOINT,
     SCHEMA,
     SCHEMAS,
@@ -689,6 +693,7 @@ define_keywords!(
     SECONDARY,
     SECRET,
     SECURITY,
+    SEED,
     SELECT,
     SEMI,
     SENSITIVE,
@@ -931,6 +936,9 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::CONNECT,
     // Reserved for snowflake MATCH_RECOGNIZE
     Keyword::MATCH_RECOGNIZE,
+    // Reserved for Snowflake table sample
+    Keyword::SAMPLE,
+    Keyword::TABLESAMPLE,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`