join hints for StarRocks

src/ast/mod.rs

@@ -40,7 +40,7 @@ pub use self::ddl::{
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     Cte, CteAsMaterialized, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause,
-    ForJson, ForXml, GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinOperator,
+    ForJson, ForXml, GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinHint, JoinOperator,
     JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
     NamedWindowDefinition, NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem,
     ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator,

src/ast/query.rs

@@ -1023,6 +1023,7 @@ impl Display for TableVersion {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct Join {
     pub relation: TableFactor,
+    pub hint: Option<JoinHint>,
     pub join_operator: JoinOperator,
 }
 
@@ -1049,66 +1050,85 @@ impl fmt::Display for Join {
             }
             Suffix(constraint)
         }
+        let hint = if let Some(hint) = &self.hint {
+            match hint {
+                JoinHint::Broadcast => "[BROADCAST] ",
+                JoinHint::Bucket => "[BUCKET] ",
+                JoinHint::Colocate => "[COLOCATE] ",
+                JoinHint::Shuffle => "[SHUFFLE] ",
+                JoinHint::Unreorder => "[UNREORDER] ",
+            }
+        } else {
+            ""
+        };
         match &self.join_operator {
             JoinOperator::Inner(constraint) => write!(
                 f,
-                " {}JOIN {}{}",
+                " {}JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::LeftOuter(constraint) => write!(
                 f,
-                " {}LEFT JOIN {}{}",
+                " {}LEFT JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::RightOuter(constraint) => write!(
                 f,
-                " {}RIGHT JOIN {}{}",
+                " {}RIGHT JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::FullOuter(constraint) => write!(
                 f,
-                " {}FULL JOIN {}{}",
+                " {}FULL JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
-            JoinOperator::CrossJoin => write!(f, " CROSS JOIN {}", self.relation),
+            JoinOperator::CrossJoin => write!(f, " CROSS JOIN {}{}", hint, self.relation),
             JoinOperator::LeftSemi(constraint) => write!(
                 f,
-                " {}LEFT SEMI JOIN {}{}",
+                " {}LEFT SEMI JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::RightSemi(constraint) => write!(
                 f,
-                " {}RIGHT SEMI JOIN {}{}",
+                " {}RIGHT SEMI JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::LeftAnti(constraint) => write!(
                 f,
-                " {}LEFT ANTI JOIN {}{}",
+                " {}LEFT ANTI JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
             JoinOperator::RightAnti(constraint) => write!(
                 f,
-                " {}RIGHT ANTI JOIN {}{}",
+                " {}RIGHT ANTI JOIN {}{}{}",
                 prefix(constraint),
+                hint,
                 self.relation,
                 suffix(constraint)
             ),
-            JoinOperator::CrossApply => write!(f, " CROSS APPLY {}", self.relation),
-            JoinOperator::OuterApply => write!(f, " OUTER APPLY {}", self.relation),
+            JoinOperator::CrossApply => write!(f, " CROSS APPLY {}{}", hint, self.relation),
+            JoinOperator::OuterApply => write!(f, " OUTER APPLY {}{}", hint, self.relation),
         }
     }
 }
@@ -1136,6 +1156,17 @@ pub enum JoinOperator {
     OuterApply,
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum JoinHint {
+    Broadcast,
+    Bucket,
+    Colocate,
+    Shuffle,
+    Unreorder,
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/keywords.rs

@@ -120,8 +120,10 @@ define_keywords!(
     BOOL,
     BOOLEAN,
     BOTH,
+    BROADCAST,
     BROWSE,
     BTREE,
+    BUCKET,
     BY,
     BYPASSRLS,
     BYTEA,
@@ -157,6 +159,7 @@ define_keywords!(
     COLLATION,
     COLLECT,
     COLLECTION,
+    COLOCATE,
     COLUMN,
     COLUMNS,
     COMMENT,
@@ -637,6 +640,7 @@ define_keywords!(
     SETS,
     SHARE,
     SHOW,
+    SHUFFLE,
     SIMILAR,
     SKIP,
     SLOW,
@@ -731,6 +735,7 @@ define_keywords!(
     UNLOGGED,
     UNNEST,
     UNPIVOT,
+    UNREORDER,
     UNSAFE,
     UNSIGNED,
     UNTIL,

src/parser/mod.rs

@@ -8069,13 +8069,15 @@ impl<'a> Parser<'a> {
                 };
                 Join {
                     relation: self.parse_table_factor()?,
+                    hint: None,
                     join_operator,
                 }
             } else if self.parse_keyword(Keyword::OUTER) {
                 // MSSQL extension, similar to LEFT JOIN LATERAL .. ON 1=1
                 self.expect_keyword(Keyword::APPLY)?;
                 Join {
                     relation: self.parse_table_factor()?,
+                    hint: None,
                     join_operator: JoinOperator::OuterApply,
                 }
             } else {
@@ -8154,10 +8156,35 @@ impl<'a> Parser<'a> {
                     }
                     _ => break,
                 };
+
+                let hint = if self.consume_token(&Token::LBracket) {
+                    let hint = Some(
+                        self.expect_one_of_keywords(&[
+                            Keyword::BROADCAST,
+                            Keyword::BUCKET,
+                            Keyword::COLOCATE,
+                            Keyword::SHUFFLE,
+                            Keyword::UNREORDER,
+                        ])
+                        .map(|kw| match kw {
+                            Keyword::BROADCAST => JoinHint::Broadcast,
+                            Keyword::BUCKET => JoinHint::Bucket,
+                            Keyword::COLOCATE => JoinHint::Colocate,
+                            Keyword::SHUFFLE => JoinHint::Shuffle,
+                            Keyword::UNREORDER => JoinHint::Unreorder,
+                            _ => unreachable!(),
+                        })?,
+                    );
+                    self.expect_token(&Token::RBracket)?;
+                    hint
+                } else {
+                    None
+                };
                 let relation = self.parse_table_factor()?;
                 let join_constraint = self.parse_join_constraint(natural)?;
                 Join {
                     relation,
+                    hint,
                     join_operator: join_operator_type(join_constraint),
                 }
             };

src/test_utils.rs

@@ -304,6 +304,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
 pub fn join(relation: TableFactor) -> Join {
     Join {
         relation,
+        hint: None,
         join_operator: JoinOperator::Inner(JoinConstraint::Natural),
     }
 }