Spark: Parse typed numeric literal suffixes

fuzz/fuzz_seeds/spark_2912_numeric_suffix.sql

@@ -0,0 +1 @@
+SELECT 1Y, 1.5bd, 1e5D, 1Dx

src/ast/mod.rs

@@ -130,7 +130,7 @@ pub use self::trigger::{
 
 pub use self::value::{
     escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,
-    NormalizationForm, QuoteDelimitedString, TrimWhereField, Value, ValueWithSpan,
+    NormalizationForm, NumberSuffix, QuoteDelimitedString, TrimWhereField, Value, ValueWithSpan,
 };
 
 use crate::ast::helpers::key_value_options::KeyValueOptions;

src/ast/value.rs

@@ -146,6 +146,13 @@ pub enum Value {
     /// but fail on CI with the `--all-features` flag enabled
     /// Numeric literal (uses `BigDecimal` when the `bigdecimal` feature is enabled).
     Number(BigDecimal, bool),
+    /// Numeric literal with a type suffix, e.g. `1Y` or `1.5BD`. The string is the
+    /// number as written without the suffix, also with the `bigdecimal` feature.
+    ///
+    /// See [`Dialect::supports_numeric_literal_suffixes`].
+    ///
+    /// [`Dialect::supports_numeric_literal_suffixes`]: crate::dialect::Dialect::supports_numeric_literal_suffixes
+    SuffixedNumber(String, NumberSuffix),
     /// 'string value'
     SingleQuotedString(String),
     /// Dollar-quoted string literal, e.g. `$$...$$` or `$tag$...$tag$` (Postgres syntax).
@@ -262,6 +269,7 @@ impl fmt::Display for Value {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Value::Number(v, l) => write!(f, "{}{long}", v, long = if *l { "L" } else { "" }),
+            Value::SuffixedNumber(v, suffix) => write!(f, "{v}{suffix}"),
             Value::DoubleQuotedString(v) => write!(f, "\"{}\"", escape_double_quote_string(v)),
             Value::SingleQuotedString(v) => write!(f, "'{}'", escape_single_quote_string(v)),
             Value::TripleSingleQuotedString(v) => {
@@ -538,6 +546,38 @@ impl fmt::Display for NormalizationForm {
     }
 }
 
+/// The type suffix of a [`Value::SuffixedNumber`]. An `L` suffix makes a
+/// [`Value::Number`] long instead.
+///
+/// See [Spark SQL](https://spark.apache.org/docs/latest/sql-ref-literals.html#numeric-literal).
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Ord, PartialOrd, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum NumberSuffix {
+    /// `Y`, a `TINYINT`
+    TinyInt,
+    /// `S`, a `SMALLINT`
+    SmallInt,
+    /// `F`, a `FLOAT`
+    Float,
+    /// `D`, a `DOUBLE`
+    Double,
+    /// `BD`, a `DECIMAL`
+    Decimal,
+}
+
+impl fmt::Display for NumberSuffix {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            NumberSuffix::TinyInt => "Y",
+            NumberSuffix::SmallInt => "S",
+            NumberSuffix::Float => "F",
+            NumberSuffix::Double => "D",
+            NumberSuffix::Decimal => "BD",
+        })
+    }
+}
+
 pub struct EscapeQuotedString<'a> {
     string: &'a str,
     quote: char,

src/dialect/mod.rs

@@ -615,6 +615,15 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if a case-insensitive suffix gives a numeric literal its type:
+    /// `L`, `S` or `Y` after an integer, as in `1Y`, and `F`, `D` or `BD` after
+    /// any number, as in `1.5BD`. See [`NumberSuffix`].
+    ///
+    /// [`NumberSuffix`]: crate::ast::NumberSuffix
+    fn supports_numeric_literal_suffixes(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports integer literals in hexadecimal,
     /// octal and binary, e.g. `0x1F`, `0o17` and `0b101`, which then parse as
     /// numbers instead of `0x` hexadecimal strings.

src/dialect/spark.rs

@@ -112,6 +112,11 @@ impl Dialect for SparkSqlDialect {
         true
     }
 
+    /// See <https://spark.apache.org/docs/latest/sql-ref-literals.html#numeric-literal>
+    fn supports_numeric_literal_suffixes(&self) -> bool {
+        true
+    }
+
     fn require_interval_qualifier(&self) -> bool {
         true
     }

src/parser/mod.rs

@@ -2083,6 +2083,7 @@ impl<'a> Parser<'a> {
                 Ok(Expr::Value(self.parse_value()?))
             }
             Token::Number(_, _)
+            | Token::SuffixedNumber(_, _)
             | Token::SingleQuotedString(_)
             | Token::DoubleQuotedString(_)
             | Token::TripleSingleQuotedString(_)
@@ -13327,6 +13328,7 @@ impl<'a> Parser<'a> {
             Token::Number(n, l) => {
                 ok_value(Value::Number(Self::parse_number_literal(n, span.start)?, l))
             }
+            Token::SuffixedNumber(n, suffix) => ok_value(Value::SuffixedNumber(n, suffix)),
             Token::SingleQuotedString(ref s) => ok_value(Value::SingleQuotedString(
                 self.maybe_concat_string_literal(s.to_string()),
             )),

src/tokenizer.rs

@@ -40,7 +40,7 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
-use crate::ast::{DollarQuotedString, QuoteDelimitedString};
+use crate::ast::{DollarQuotedString, NumberSuffix, QuoteDelimitedString};
 use crate::dialect::Dialect;
 use crate::dialect::{
     BigQueryDialect, DuckDbDialect, GenericDialect, MySqlDialect, PostgreSqlDialect,
@@ -62,6 +62,8 @@ pub enum Token {
     Word(Word),
     /// An unsigned numeric literal
     Number(String, bool),
+    /// An unsigned numeric literal with a type suffix, e.g. `1Y`
+    SuffixedNumber(String, NumberSuffix),
     /// A character that could not be tokenized
     Char(char),
     /// Single quoted string: i.e: 'string'
@@ -303,6 +305,7 @@ impl fmt::Display for Token {
             Token::EOF => f.write_str("EOF"),
             Token::Word(ref w) => write!(f, "{w}"),
             Token::Number(ref n, l) => write!(f, "{}{long}", n, long = if *l { "L" } else { "" }),
+            Token::SuffixedNumber(ref n, suffix) => write!(f, "{n}{suffix}"),
             Token::Char(ref c) => write!(f, "{c}"),
             Token::SingleQuotedString(ref s) => write!(f, "'{s}'"),
             Token::TripleSingleQuotedString(ref s) => write!(f, "'''{s}'''"),
@@ -1707,6 +1710,10 @@ impl<'a> Tokenizer<'a> {
                         }
                     }
 
+                    if self.dialect.supports_numeric_literal_suffixes() {
+                        return Ok(Some(self.tokenize_number_suffix(chars, s)));
+                    }
+
                     let long = if chars.peek() == Some(&'L') {
                         chars.next();
                         true
@@ -2342,6 +2349,39 @@ impl<'a> Tokenizer<'a> {
         Ok(Token::Number(s, false))
     }
 
+    /// Tokenizes the number `s` with its type suffix if one follows, as in `1Y` or
+    /// `1.5BD`. A suffix directly followed by an identifier character is not one.
+    fn tokenize_number_suffix(&self, chars: &mut State, s: String) -> Token {
+        let is_integer = !s.contains(['.', 'e', 'E']);
+        let mut ahead = chars.peekable.clone();
+        // `None` is the `L` of a long number.
+        let (suffix, len) = match ahead.next().map(|ch| ch.to_ascii_uppercase()) {
+            Some('L') if is_integer => (None, 1),
+            Some('S') if is_integer => (Some(NumberSuffix::SmallInt), 1),
+            Some('Y') if is_integer => (Some(NumberSuffix::TinyInt), 1),
+            Some('F') => (Some(NumberSuffix::Float), 1),
+            Some('D') => (Some(NumberSuffix::Double), 1),
+            Some('B') if ahead.peek().is_some_and(|ch| ch.eq_ignore_ascii_case(&'D')) => {
+                ahead.next();
+                (Some(NumberSuffix::Decimal), 2)
+            }
+            _ => return Token::Number(s, false),
+        };
+        if ahead
+            .next()
+            .is_some_and(|ch| self.dialect.is_identifier_part(ch))
+        {
+            return Token::Number(s, false);
+        }
+        for _ in 0..len {
+            chars.next();
+        }
+        match suffix {
+            Some(suffix) => Token::SuffixedNumber(s, suffix),
+            None => Token::Number(s, true),
+        }
+    }
+
     // Consume characters until newline
     fn tokenize_single_line_comment(&self, chars: &mut State) -> String {
         peeking_take_while(chars, |ch| match ch {

tests/sqlparser_spark.rs

@@ -19,7 +19,7 @@
 //! Test SQL syntax specific to Apache Spark SQL.
 
 use sqlparser::ast::*;
-use sqlparser::dialect::SparkSqlDialect;
+use sqlparser::dialect::{GenericDialect, SparkSqlDialect};
 use test_utils::*;
 
 #[macro_use]
@@ -381,3 +381,53 @@ fn test_substring() {
 fn test_pipe_operator() {
     spark().verified_stmt("SELECT * FROM t |> WHERE x > 1 |> SELECT x AS y |> ORDER BY y");
 }
+
+#[test]
+fn test_numeric_literal_suffixes() {
+    let suffixed = |n: &str, suffix| Expr::value(Value::SuffixedNumber(n.into(), suffix));
+    assert_eq!(
+        spark().verified_expr("10Y"),
+        suffixed("10", NumberSuffix::TinyInt)
+    );
+    assert_eq!(
+        spark().verified_expr("10S"),
+        suffixed("10", NumberSuffix::SmallInt)
+    );
+    assert_eq!(
+        spark().verified_expr("1e5F"),
+        suffixed("1e5", NumberSuffix::Float)
+    );
+    assert_eq!(
+        spark().verified_expr("1.5D"),
+        suffixed("1.5", NumberSuffix::Double)
+    );
+    assert_eq!(
+        spark().verified_expr(".5BD"),
+        suffixed(".5", NumberSuffix::Decimal)
+    );
+    assert_eq!(
+        spark().verified_expr("10L"),
+        Expr::value(Value::Number("10".parse().unwrap(), true))
+    );
+
+    spark().verified_stmt("SELECT 1Y, 1S, 1L, 1F, 1D, 1BD");
+    spark().verified_stmt("SELECT 1.5F, 1.5D, 1.5BD, 5.F, 1.5E-3BD, -32Y");
+    spark().verified_stmt("SELECT 2 * 1.5D - 1Y FROM t WHERE x > 1BD");
+    spark().one_statement_parses_to(
+        "SELECT 1y, 1s, 1l, 1f, 1d, 1bd, 1.5bD, 9223372036854775807l",
+        "SELECT 1Y, 1S, 1L, 1F, 1D, 1BD, 1.5BD, 9223372036854775807L",
+    );
+
+    // `L`, `S` and `Y` only follow an integer, and no suffix runs into an identifier.
+    spark().one_statement_parses_to(
+        "SELECT 1.5S, 1.5Y, 1.5L, 1Dx, 1D5, 1B, 1BDx, 1La",
+        "SELECT 1.5 AS S, 1.5 AS Y, 1.5 AS L, 1 AS Dx, 1 AS D5, 1 AS B, 1 AS BDx, 1 AS La",
+    );
+    assert_eq!(
+        spark().parse_sql_statements("SELECT 1e5L").unwrap(),
+        spark().parse_sql_statements("SELECT 1e5 AS L").unwrap()
+    );
+
+    TestedDialects::new(vec![Box::new(GenericDialect {})])
+        .one_statement_parses_to("SELECT 1D, 1Y, 1l, 1L", "SELECT 1 AS D, 1 AS Y, 1 AS l, 1L");
+}