feat: support sample for table scan

src/query/ast/src/ast/expr.rs

@@ -789,6 +789,23 @@ pub enum Literal {
     Null,
 }
 
+impl Literal {
+    pub fn as_double(&self) -> Result<f64> {
+        match self {
+            Literal::UInt64(val) => Ok(*val as f64),
+            Literal::Float64(val) => Ok(*val),
+            Literal::Decimal256 { value, scale, .. } => {
+                let div = 10_f64.powi(*scale as i32);
+                Ok(value.as_f64() / div)
+            }
+            _ => Err(ParseError(
+                None,
+                format!("Cannot convert {:?} to double", self),
+            )),
+        }
+    }
+}
+
 impl Display for Literal {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         match self {

src/query/ast/src/ast/format/ast_format.rs

@@ -3327,6 +3327,7 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 consume,
                 pivot,
                 unpivot,
+                sample,
             } => {
                 let mut name = String::new();
                 name.push_str("TableIdentifier ");
@@ -3354,6 +3355,11 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                     name.push_str(&unpivot.to_string());
                 }
 
+                if let Some(sample) = sample {
+                    name.push(' ');
+                    name.push_str(&sample.to_string());
+                }
+
                 let mut children = Vec::new();
 
                 if let Some(temporal) = temporal {

src/query/ast/src/ast/format/syntax/query.rs

@@ -322,6 +322,7 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
             consume,
             pivot,
             unpivot,
+            sample,
         } => if let Some(catalog) = catalog {
             RcDoc::text(catalog.to_string()).append(RcDoc::text("."))
         } else {
@@ -353,6 +354,11 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
         } else {
             RcDoc::nil()
         })
+        .append(if let Some(sample) = sample {
+            RcDoc::text(format!(" {sample}"))
+        } else {
+            RcDoc::nil()
+        })
         .append(if let Some(alias) = alias {
             RcDoc::text(format!(" AS {alias}"))
         } else {

src/query/ast/src/ast/query.rs

@@ -19,6 +19,7 @@ use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
 use super::Lambda;
+use super::Literal;
 use crate::ast::write_comma_separated_list;
 use crate::ast::write_dot_separated_list;
 use crate::ast::Expr;
@@ -608,6 +609,39 @@ impl Display for TemporalClause {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum SampleLevel {
+    ROW,
+    BLOCK,
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum SampleConfig {
+    Probability(Literal),
+    RowsNum(Literal),
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct Sample {
+    pub sample_level: SampleLevel,
+    pub sample_conf: SampleConfig,
+}
+
+impl Display for Sample {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "SAMPLE ")?;
+        match self.sample_level {
+            SampleLevel::ROW => write!(f, "ROW ")?,
+            SampleLevel::BLOCK => write!(f, "BLOCK ")?,
+        }
+        match &self.sample_conf {
+            SampleConfig::Probability(prob) => write!(f, "({})", prob)?,
+            SampleConfig::RowsNum(rows) => write!(f, "({} ROWS)", rows)?,
+        }
+        Ok(())
+    }
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum TableReference {
@@ -623,6 +657,7 @@ pub enum TableReference {
         consume: bool,
         pivot: Option<Box<Pivot>>,
         unpivot: Option<Box<Unpivot>>,
+        sample: Option<Sample>,
     },
     // `TABLE(expr)[ AS alias ]`
     TableFunction {
@@ -697,6 +732,7 @@ impl Display for TableReference {
                 consume,
                 pivot,
                 unpivot,
+                sample,
             } => {
                 write_dot_separated_list(
                     f,
@@ -721,6 +757,10 @@ impl Display for TableReference {
                 if let Some(unpivot) = unpivot {
                     write!(f, " {unpivot}")?;
                 }
+
+                if let Some(sample) = sample {
+                    write!(f, " {sample}")?;
+                }
             }
             TableReference::TableFunction {
                 span: _,

src/query/ast/src/ast/statements/merge_into.rs

@@ -218,6 +218,7 @@ impl MergeSource {
                 consume: false,
                 pivot: None,
                 unpivot: None,
+                sample: None,
             },
         }
     }

src/query/ast/src/parser/query.rs

@@ -685,6 +685,7 @@ pub enum TableReferenceElement {
         consume: bool,
         pivot: Option<Box<Pivot>>,
         unpivot: Option<Box<Unpivot>>,
+        sample: Option<Sample>,
     },
     // `TABLE(expr)[ AS alias ]`
     TableFunction {
@@ -741,9 +742,43 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
     );
     let aliased_table = map(
         rule! {
-            #dot_separated_idents_1_to_3 ~ #temporal_clause? ~ (WITH ~ CONSUME)? ~ #table_alias? ~ #pivot? ~ #unpivot?
+            #dot_separated_idents_1_to_3 ~ #temporal_clause? ~ (WITH ~ CONSUME)? ~ #table_alias? ~ #pivot? ~ #unpivot? ~ SAMPLE? ~ (ROW | BLOCK)? ~ ("(" ~ #expr ~ ROWS? ~ ")")?
         },
-        |((catalog, database, table), temporal, opt_consume, alias, pivot, unpivot)| {
+        |(
+            (catalog, database, table),
+            temporal,
+            opt_consume,
+            alias,
+            pivot,
+            unpivot,
+            sample,
+            level,
+            sample_conf,
+        )| {
+            let mut table_sample = None;
+            if sample.is_some() {
+                let sample_level = match level {
+                    // If the sample level is not specified, it defaults to ROW
+                    Some(level) => match level.kind {
+                        ROW => SampleLevel::ROW,
+                        BLOCK => SampleLevel::BLOCK,
+                        _ => unreachable!(),
+                    },
+                    None => SampleLevel::ROW,
+                };
+                let mut default_sample_conf = SampleConfig::Probability(Literal::Float64(100.0));
+                if let Some((_, Expr::Literal { value, .. }, rows, _)) = sample_conf {
+                    default_sample_conf = if rows.is_some() {
+                        SampleConfig::RowsNum(value)
+                    } else {
+                        SampleConfig::Probability(value)
+                    };
+                }
+                table_sample = Some(Sample {
+                    sample_level,
+                    sample_conf: default_sample_conf,
+                })
+            };
             TableReferenceElement::Table {
                 catalog,
                 database,
@@ -753,6 +788,7 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
                 consume: opt_consume.is_some(),
                 pivot: pivot.map(Box::new),
                 unpivot: unpivot.map(Box::new),
+                sample: table_sample,
             }
         },
     );
@@ -864,6 +900,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
                 consume,
                 pivot,
                 unpivot,
+                sample,
             } => TableReference::Table {
                 span: transform_span(input.span.tokens),
                 catalog,
@@ -874,6 +911,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
                 consume,
                 pivot,
                 unpivot,
+                sample,
             },
             TableReferenceElement::TableFunction {
                 lateral,

src/query/ast/src/parser/statement.rs

@@ -4140,6 +4140,7 @@ pub fn table_reference_with_alias(i: Input) -> IResult<TableReference> {
             consume: false,
             pivot: None,
             unpivot: None,
+            sample: None,
         },
     )(i)
 }
@@ -4159,6 +4160,7 @@ pub fn table_reference_only(i: Input) -> IResult<TableReference> {
             consume: false,
             pivot: None,
             unpivot: None,
+            sample: None,
         },
     )(i)
 }

src/query/ast/src/parser/token.rs

@@ -405,6 +405,8 @@ pub enum TokenKind {
     BROTLI,
     #[token("BZ2", ignore(ascii_case))]
     BZ2,
+    #[token("BLOCK", ignore(ascii_case))]
+    BLOCK,
     #[token("CALL", ignore(ascii_case))]
     CALL,
     #[token("CASE", ignore(ascii_case))]
@@ -924,6 +926,8 @@ pub enum TokenKind {
     RETURN_FAILED_ONLY,
     #[token("REVERSE", ignore(ascii_case))]
     REVERSE,
+    #[token("SAMPLE", ignore(ascii_case))]
+    SAMPLE,
     #[token("MERGE", ignore(ascii_case))]
     MERGE,
     #[token("MATCHED", ignore(ascii_case))]
@@ -1567,6 +1571,7 @@ impl TokenKind {
             // | TokenKind::AUTHORIZATION
             // | TokenKind::BINARY
             | TokenKind::BOTH
+            | TokenKind::BLOCK
             | TokenKind::CASE
             | TokenKind::CAST
             // | TokenKind::CHECK
@@ -1624,6 +1629,7 @@ impl TokenKind {
             // | TokenKind::SIMILAR
             | TokenKind::SOME
             | TokenKind::SEMI
+            | TokenKind::SAMPLE
             // | TokenKind::SYMMETRIC
             // | TokenKind::TABLESAMPLE
             | TokenKind::THEN

src/query/ast/tests/it/parser.rs

@@ -229,6 +229,10 @@ fn test_statement() {
         r#"select * FROM t where ((a));"#,
         r#"select * FROM t where ((select 1) > 1);"#,
         r#"select ((t1.a)>=(((((t2.b)<=(t3.c))) IS NOT NULL)::INTEGER));"#,
+        r#"select * from t sample row (99);"#,
+        r#"select * from t sample block (99);"#,
+        r#"select * from t sample row (10 rows);"#,
+        r#"select * from t sample block (10 rows);"#,
         r#"insert into t (c1, c2) values (1, 2), (3, 4);"#,
         r#"insert into t (c1, c2) values (1, 2);"#,
         r#"insert into table t select * from t2;"#,