Hive: Support DISTRIBUTE BY and SORT BY in window specs

src/ast/mod.rs

@@ -2269,12 +2269,56 @@ pub struct WindowSpec {
     pub window_name: Option<Ident>,
     /// `OVER (PARTITION BY ...)`
     pub partition_by: Vec<Expr>,
+    /// The kind of partitioning clause used in the window specification.
+    pub partition_by_kind: WindowPartitionByKind,
     /// `OVER (ORDER BY ...)`
     pub order_by: Vec<OrderByExpr>,
+    /// The kind of ordering clause used in the window specification.
+    pub order_by_kind: WindowOrderByKind,
     /// `OVER (window frame)`
     pub window_frame: Option<WindowFrame>,
 }
 
+/// The kind of partitioning clause in a window specification.
+#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum WindowPartitionByKind {
+    /// `PARTITION BY`
+    Partition,
+    /// Hive `DISTRIBUTE BY`
+    Distribute,
+}
+
+impl fmt::Display for WindowPartitionByKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            Self::Partition => "PARTITION BY",
+            Self::Distribute => "DISTRIBUTE BY",
+        })
+    }
+}
+
+/// The kind of ordering clause in a window specification.
+#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum WindowOrderByKind {
+    /// `ORDER BY`
+    Order,
+    /// Hive `SORT BY`
+    Sort,
+}
+
+impl fmt::Display for WindowOrderByKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            Self::Order => "ORDER BY",
+            Self::Sort => "SORT BY",
+        })
+    }
+}
+
 impl fmt::Display for WindowSpec {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let mut is_first = true;
@@ -2292,7 +2336,8 @@ impl fmt::Display for WindowSpec {
             is_first = false;
             write!(
                 f,
-                "PARTITION BY {}",
+                "{} {}",
+                self.partition_by_kind,
                 display_comma_separated(&self.partition_by)
             )?;
         }
@@ -2301,7 +2346,12 @@ impl fmt::Display for WindowSpec {
                 SpaceOrNewline.fmt(f)?;
             }
             is_first = false;
-            write!(f, "ORDER BY {}", display_comma_separated(&self.order_by))?;
+            write!(
+                f,
+                "{} {}",
+                self.order_by_kind,
+                display_comma_separated(&self.order_by)
+            )?;
         }
         if let Some(window_frame) = &self.window_frame {
             if !is_first {

src/parser/mod.rs

@@ -20272,15 +20272,32 @@ impl<'a> Parser<'a> {
             _ => None,
         };
 
-        let partition_by = if self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) {
-            self.parse_comma_separated(Parser::parse_expr)?
-        } else {
-            vec![]
-        };
-        let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
-            self.parse_comma_separated(Parser::parse_order_by_expr)?
+        let (partition_by_kind, partition_by) =
+            if self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) {
+                (
+                    WindowPartitionByKind::Partition,
+                    self.parse_comma_separated(Parser::parse_expr)?,
+                )
+            } else if self.parse_keywords(&[Keyword::DISTRIBUTE, Keyword::BY]) {
+                (
+                    WindowPartitionByKind::Distribute,
+                    self.parse_comma_separated(Parser::parse_expr)?,
+                )
+            } else {
+                (WindowPartitionByKind::Partition, vec![])
+            };
+        let (order_by_kind, order_by) = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
+            (
+                WindowOrderByKind::Order,
+                self.parse_comma_separated(Parser::parse_order_by_expr)?,
+            )
+        } else if self.parse_keywords(&[Keyword::SORT, Keyword::BY]) {
+            (
+                WindowOrderByKind::Sort,
+                self.parse_comma_separated(Parser::parse_order_by_expr)?,
+            )
         } else {
-            vec![]
+            (WindowOrderByKind::Order, vec![])
         };
 
         let window_frame = if !self.consume_token(&Token::RParen) {
@@ -20293,7 +20310,9 @@ impl<'a> Parser<'a> {
         Ok(WindowSpec {
             window_name,
             partition_by,
+            partition_by_kind,
             order_by,
+            order_by_kind,
             window_frame,
         })
     }

tests/sqlparser_common.rs

@@ -3095,6 +3095,7 @@ fn parse_select_qualify() {
                 over: Some(WindowType::WindowSpec(WindowSpec {
                     window_name: None,
                     partition_by: vec![Expr::Identifier(Ident::new("p"))],
+                    partition_by_kind: WindowPartitionByKind::Partition,
                     order_by: vec![OrderByExpr {
                         expr: Expr::Identifier(Ident::new("o")),
                         options: OrderByOptions {
@@ -3103,6 +3104,7 @@ fn parse_select_qualify() {
                         },
                         with_fill: None,
                     }],
+                    order_by_kind: WindowOrderByKind::Order,
                     window_frame: None,
                 })),
                 within_group: vec![]
@@ -5815,6 +5817,7 @@ fn parse_window_functions() {
             over: Some(WindowType::WindowSpec(WindowSpec {
                 window_name: None,
                 partition_by: vec![],
+                partition_by_kind: WindowPartitionByKind::Partition,
                 order_by: vec![OrderByExpr {
                     expr: Expr::Identifier(Ident::new("dt")),
                     options: OrderByOptions {
@@ -5823,6 +5826,7 @@ fn parse_window_functions() {
                     },
                     with_fill: None,
                 }],
+                order_by_kind: WindowOrderByKind::Order,
                 window_frame: None,
             })),
             within_group: vec![],
@@ -6037,6 +6041,7 @@ fn test_parse_named_window() {
                 NamedWindowExpr::WindowSpec(WindowSpec {
                     window_name: None,
                     partition_by: vec![],
+                    partition_by_kind: WindowPartitionByKind::Partition,
                     order_by: vec![OrderByExpr {
                         expr: Expr::Identifier(Ident {
                             value: "C12".to_string(),
@@ -6049,6 +6054,7 @@ fn test_parse_named_window() {
                         },
                         with_fill: None,
                     }],
+                    order_by_kind: WindowOrderByKind::Order,
                     window_frame: None,
                 }),
             ),
@@ -6065,7 +6071,9 @@ fn test_parse_named_window() {
                         quote_style: None,
                         span: Span::empty(),
                     })],
+                    partition_by_kind: WindowPartitionByKind::Partition,
                     order_by: vec![],
+                    order_by_kind: WindowOrderByKind::Order,
                     window_frame: None,
                 }),
             ),

tests/sqlparser_hive.rs

@@ -569,6 +569,13 @@ fn parse_create_table_with_map_column_comment() {
     );
 }
 
+#[test]
+fn parse_row_number_window_function() {
+    hive().verified_stmt(
+        "SELECT row_number() OVER (DISTRIBUTE BY age SORT BY update_time DESC) AS row_num FROM sdl.xxx",
+    );
+}
+
 fn hive() -> TestedDialects {
     TestedDialects::new(vec![Box::new(HiveDialect {})])
 }

tests/sqlparser_sqlite.rs

@@ -439,7 +439,9 @@ fn parse_window_function_with_filter() {
                 over: Some(WindowType::WindowSpec(WindowSpec {
                     window_name: None,
                     partition_by: vec![],
+                    partition_by_kind: WindowPartitionByKind::Partition,
                     order_by: vec![],
+                    order_by_kind: WindowOrderByKind::Order,
                     window_frame: None,
                 })),
                 filter: Some(Box::new(Expr::Identifier(Ident::new("y")))),