Spark, Databricks: parse the << and >> shift operators

src/dialect/databricks.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::dialect::Dialect;
+use crate::dialect::{Dialect, Precedence};
+use crate::parser::{Parser, ParserError};
+use crate::tokenizer::Token;
 
 /// A [`Dialect`] for [Databricks SQL](https://www.databricks.com/)
 ///
@@ -25,6 +27,21 @@ use crate::dialect::Dialect;
 pub struct DatabricksDialect;
 
 impl Dialect for DatabricksDialect {
+    /// `<<` and `>>`, binding tighter than `&` and looser than `+`, as Spark's
+    /// grammar orders them.
+    fn supports_bitwise_shift_operators(&self) -> bool {
+        true
+    }
+
+    fn get_next_precedence(&self, parser: &Parser) -> Option<Result<u8, ParserError>> {
+        match parser.peek_token_ref().token {
+            Token::ShiftLeft | Token::ShiftRight => {
+                Some(Ok(self.prec_value(Precedence::Ampersand) + 1))
+            }
+            _ => None,
+        }
+    }
+
     // see https://docs.databricks.com/en/sql/language-manual/sql-ref-identifiers.html
 
     fn is_delimited_identifier_start(&self, ch: char) -> bool {

src/dialect/spark.rs

@@ -19,9 +19,10 @@
 use alloc::boxed::Box;
 
 use crate::ast::{BinaryOperator, Expr};
-use crate::dialect::Dialect;
+use crate::dialect::{Dialect, Precedence};
 use crate::keywords::Keyword;
 use crate::parser::{Parser, ParserError};
+use crate::tokenizer::Token;
 
 /// A [`Dialect`] for [Apache Spark SQL](https://spark.apache.org/docs/latest/sql-ref.html).
 ///
@@ -31,6 +32,21 @@ use crate::parser::{Parser, ParserError};
 pub struct SparkSqlDialect;
 
 impl Dialect for SparkSqlDialect {
+    /// `<<` and `>>`, binding tighter than `&` and looser than `+`, as Spark's
+    /// grammar orders them.
+    fn supports_bitwise_shift_operators(&self) -> bool {
+        true
+    }
+
+    fn get_next_precedence(&self, parser: &Parser) -> Option<Result<u8, ParserError>> {
+        match parser.peek_token_ref().token {
+            Token::ShiftLeft | Token::ShiftRight => {
+                Some(Ok(self.prec_value(Precedence::Ampersand) + 1))
+            }
+            _ => None,
+        }
+    }
+
     // See https://spark.apache.org/docs/latest/sql-ref-identifier.html
     fn is_delimited_identifier_start(&self, ch: char) -> bool {
         matches!(ch, '`')

tests/sqlparser_spark.rs

@@ -485,3 +485,42 @@ fn test_describe_table_is_not_a_table_query() {
     // EXPLAIN is unaffected.
     spark().one_statement_parses_to("EXPLAIN TABLE t", "EXPLAIN SELECT * FROM t");
 }
+
+#[test]
+fn test_bitwise_shift_operators() {
+    spark().verified_stmt("SELECT 4 >> 1, 1 << 2");
+    // Tighter than `&`, looser than `+`.
+    match spark().verified_expr("1 & 3 << 1") {
+        Expr::BinaryOp {
+            op: BinaryOperator::BitwiseAnd,
+            right,
+            ..
+        } => assert!(matches!(
+            *right,
+            Expr::BinaryOp {
+                op: BinaryOperator::PGBitwiseShiftLeft,
+                ..
+            }
+        )),
+        other => panic!("Expected & over <<, got {other:?}"),
+    }
+    match spark().verified_expr("1 << 1 + 1") {
+        Expr::BinaryOp {
+            op: BinaryOperator::PGBitwiseShiftLeft,
+            right,
+            ..
+        } => assert!(matches!(
+            *right,
+            Expr::BinaryOp {
+                op: BinaryOperator::Plus,
+                ..
+            }
+        )),
+        other => panic!("Expected << over +, got {other:?}"),
+    }
+    spark().verified_stmt("SELECT (1 & 3) << 1");
+    match spark().verified_expr("a >> b") {
+        Expr::BinaryOp { op, .. } => assert_eq!(op, BinaryOperator::PGBitwiseShiftRight),
+        other => panic!("Expected a shift, got {other:?}"),
+    }
+}