Extend support for INDEX parsing (#1707)

Co-authored-by: Ifeanyi Ubah <[email protected]>

Author: LucaCappelletti94

Diff for: src/ast/ddl.rs

@@ -1174,20 +1174,33 @@ impl fmt::Display for KeyOrIndexDisplay {
 /// [1]: https://dev.mysql.com/doc/refman/8.0/en/create-table.html
 /// [2]: https://dev.mysql.com/doc/refman/8.0/en/create-index.html
 /// [3]: https://www.postgresql.org/docs/14/sql-createindex.html
-#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum IndexType {
     BTree,
     Hash,
-    // TODO add Postgresql's possible indexes
+    GIN,
+    GiST,
+    SPGiST,
+    BRIN,
+    Bloom,
+    /// Users may define their own index types, which would
+    /// not be covered by the above variants.
+    Custom(Ident),
 }
 
 impl fmt::Display for IndexType {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::BTree => write!(f, "BTREE"),
             Self::Hash => write!(f, "HASH"),
+            Self::GIN => write!(f, "GIN"),
+            Self::GiST => write!(f, "GIST"),
+            Self::SPGiST => write!(f, "SPGIST"),
+            Self::BRIN => write!(f, "BRIN"),
+            Self::Bloom => write!(f, "BLOOM"),
+            Self::Custom(name) => write!(f, "{}", name),
         }
     }
 }

Diff for: src/ast/dml.rs

@@ -34,12 +34,31 @@ pub use super::ddl::{ColumnDef, TableConstraint};
 use super::{
     display_comma_separated, display_separated, query::InputFormatClause, Assignment, ClusteredBy,
     CommentDef, Expr, FileFormat, FromTable, HiveDistributionStyle, HiveFormat, HiveIOFormat,
-    HiveRowFormat, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnCommit, OnInsert,
-    OneOrManyWithParens, OrderByExpr, Query, RowAccessPolicy, SelectItem, Setting, SqlOption,
-    SqliteOnConflict, StorageSerializationPolicy, TableEngine, TableObject, TableWithJoins, Tag,
-    WrappedCollection,
+    HiveRowFormat, Ident, IndexType, InsertAliases, MysqlInsertPriority, ObjectName, OnCommit,
+    OnInsert, OneOrManyWithParens, OrderByExpr, Query, RowAccessPolicy, SelectItem, Setting,
+    SqlOption, SqliteOnConflict, StorageSerializationPolicy, TableEngine, TableObject,
+    TableWithJoins, Tag, WrappedCollection,
 };
 
+/// Index column type.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct IndexColumn {
+    pub column: OrderByExpr,
+    pub operator_class: Option<Ident>,
+}
+
+impl Display for IndexColumn {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.column)?;
+        if let Some(operator_class) = &self.operator_class {
+            write!(f, " {}", operator_class)?;
+        }
+        Ok(())
+    }
+}
+
 /// CREATE INDEX statement.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -49,8 +68,8 @@ pub struct CreateIndex {
     pub name: Option<ObjectName>,
     #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
     pub table_name: ObjectName,
-    pub using: Option<Ident>,
-    pub columns: Vec<OrderByExpr>,
+    pub using: Option<IndexType>,
+    pub columns: Vec<IndexColumn>,
     pub unique: bool,
     pub concurrently: bool,
     pub if_not_exists: bool,

Diff for: src/ast/mod.rs

@@ -58,7 +58,7 @@ pub use self::ddl::{
     ReferentialAction, TableConstraint, TagsColumnOption, UserDefinedTypeCompositeAttributeDef,
     UserDefinedTypeRepresentation, ViewColumnDef,
 };
-pub use self::dml::{CreateIndex, CreateTable, Delete, Insert};
+pub use self::dml::{CreateIndex, CreateTable, Delete, IndexColumn, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
@@ -91,6 +91,7 @@ pub use self::value::{
 
 use crate::ast::helpers::key_value_options::KeyValueOptions;
 use crate::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageParamsObject};
+
 #[cfg(feature = "visitor")]
 pub use visitor::*;
 

Diff for: src/ast/spans.rs

@@ -704,7 +704,7 @@ impl Spanned for CreateIndex {
         let CreateIndex {
             name,
             table_name,
-            using,
+            using: _,
             columns,
             unique: _,        // bool
             concurrently: _,  // bool
@@ -719,8 +719,7 @@ impl Spanned for CreateIndex {
             name.iter()
                 .map(|i| i.span())
                 .chain(core::iter::once(table_name.span()))
-                .chain(using.iter().map(|i| i.span))
-                .chain(columns.iter().map(|i| i.span()))
+                .chain(columns.iter().map(|i| i.column.span()))
                 .chain(include.iter().map(|i| i.span))
                 .chain(with.iter().map(|i| i.span()))
                 .chain(predicate.iter().map(|i| i.span())),

Diff for: src/keywords.rs

@@ -137,11 +137,13 @@ define_keywords!(
     BIT,
     BLOB,
     BLOCK,
+    BLOOM,
     BLOOMFILTER,
     BOOL,
     BOOLEAN,
     BOTH,
     BOX,
+    BRIN,
     BROWSE,
     BTREE,
     BUCKET,
@@ -386,6 +388,8 @@ define_keywords!(
     GENERATED,
     GEOGRAPHY,
     GET,
+    GIN,
+    GIST,
     GLOBAL,
     GRANT,
     GRANTED,
@@ -805,6 +809,7 @@ define_keywords!(
     SPATIAL,
     SPECIFIC,
     SPECIFICTYPE,
+    SPGIST,
     SQL,
     SQLEXCEPTION,
     SQLSTATE,

Diff for: src/parser/mod.rs

@@ -3955,6 +3955,18 @@ impl<'a> Parser<'a> {
         true
     }
 
+    /// If the current token is one of the given `keywords`, returns the keyword
+    /// that matches, without consuming the token. Otherwise, returns [`None`].
+    #[must_use]
+    pub fn peek_one_of_keywords(&self, keywords: &[Keyword]) -> Option<Keyword> {
+        for keyword in keywords {
+            if self.peek_keyword(*keyword) {
+                return Some(*keyword);
+            }
+        }
+        None
+    }
+
     /// If the current token is one of the given `keywords`, consume the token
     /// and return the keyword that matches. Otherwise, no tokens are consumed
     /// and returns [`None`].
@@ -6406,12 +6418,13 @@ impl<'a> Parser<'a> {
         };
         let table_name = self.parse_object_name(false)?;
         let using = if self.parse_keyword(Keyword::USING) {
-            Some(self.parse_identifier()?)
+            Some(self.parse_index_type()?)
         } else {
             None
         };
+
         self.expect_token(&Token::LParen)?;
-        let columns = self.parse_comma_separated(Parser::parse_order_by_expr)?;
+        let columns = self.parse_comma_separated(Parser::parse_create_index_expr)?;
         self.expect_token(&Token::RParen)?;
 
         let include = if self.parse_keyword(Keyword::INCLUDE) {
@@ -7629,16 +7642,30 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_index_type(&mut self) -> Result<IndexType, ParserError> {
-        if self.parse_keyword(Keyword::BTREE) {
-            Ok(IndexType::BTree)
+        Ok(if self.parse_keyword(Keyword::BTREE) {
+            IndexType::BTree
         } else if self.parse_keyword(Keyword::HASH) {
-            Ok(IndexType::Hash)
-        } else {
-            self.expected("index type {BTREE | HASH}", self.peek_token())
-        }
+            IndexType::Hash
+        } else if self.parse_keyword(Keyword::GIN) {
+            IndexType::GIN
+        } else if self.parse_keyword(Keyword::GIST) {
+            IndexType::GiST
+        } else if self.parse_keyword(Keyword::SPGIST) {
+            IndexType::SPGiST
+        } else if self.parse_keyword(Keyword::BRIN) {
+            IndexType::BRIN
+        } else if self.parse_keyword(Keyword::BLOOM) {
+            IndexType::Bloom
+        } else {
+            IndexType::Custom(self.parse_identifier()?)
+        })
     }
 
-    /// Parse [USING {BTREE | HASH}]
+    /// Optionally parse the `USING` keyword, followed by an [IndexType]
+    /// Example:
+    /// ```sql
+    //// USING BTREE (name, age DESC)
+    /// ```
     pub fn parse_optional_using_then_index_type(
         &mut self,
     ) -> Result<Option<IndexType>, ParserError> {
@@ -13631,10 +13658,42 @@ impl<'a> Parser<'a> {
         }
     }
 
-    /// Parse an expression, optionally followed by ASC or DESC (used in ORDER BY)
+    /// Parse an [OrderByExpr] expression.
     pub fn parse_order_by_expr(&mut self) -> Result<OrderByExpr, ParserError> {
+        self.parse_order_by_expr_inner(false)
+            .map(|(order_by, _)| order_by)
+    }
+
+    /// Parse an [IndexColumn].
+    pub fn parse_create_index_expr(&mut self) -> Result<IndexColumn, ParserError> {
+        self.parse_order_by_expr_inner(true)
+            .map(|(column, operator_class)| IndexColumn {
+                column,
+                operator_class,
+            })
+    }
+
+    fn parse_order_by_expr_inner(
+        &mut self,
+        with_operator_class: bool,
+    ) -> Result<(OrderByExpr, Option<Ident>), ParserError> {
         let expr = self.parse_expr()?;
 
+        let operator_class: Option<Ident> = if with_operator_class {
+            // We check that if non of the following keywords are present, then we parse an
+            // identifier as operator class.
+            if self
+                .peek_one_of_keywords(&[Keyword::ASC, Keyword::DESC, Keyword::NULLS, Keyword::WITH])
+                .is_some()
+            {
+                None
+            } else {
+                self.maybe_parse(|parser| parser.parse_identifier())?
+            }
+        } else {
+            None
+        };
+
         let options = self.parse_order_by_options()?;
 
         let with_fill = if dialect_of!(self is ClickHouseDialect | GenericDialect)
@@ -13645,11 +13704,14 @@ impl<'a> Parser<'a> {
             None
         };
 
-        Ok(OrderByExpr {
-            expr,
-            options,
-            with_fill,
-        })
+        Ok((
+            OrderByExpr {
+                expr,
+                options,
+                with_fill,
+            },
+            operator_class,
+        ))
     }
 
     fn parse_order_by_options(&mut self) -> Result<OrderByOptions, ParserError> {