Databricks: Add support for OPTIMIZE statement

Add support for Databricks Delta Lake OPTIMIZE statement syntax:
- OPTIMIZE table_name [WHERE predicate] [ZORDER BY (col1, ...)]

This extends the existing OptimizeTable AST to support both ClickHouse
and Databricks syntax by adding:
- has_table_keyword: distinguishes OPTIMIZE TABLE (ClickHouse) from
  OPTIMIZE (Databricks)
- predicate: optional WHERE clause for partition filtering
- zorder: optional ZORDER BY clause for data colocation

Author: funcpp

src/ast/mod.rs

@@ -4654,22 +4654,34 @@ pub enum Statement {
         /// Legacy copy-style options.
         options: Vec<CopyLegacyOption>,
     },
+    /// ClickHouse:
     /// ```sql
     /// OPTIMIZE TABLE [db.]name [ON CLUSTER cluster] [PARTITION partition | PARTITION ID 'partition_id'] [FINAL] [DEDUPLICATE [BY expression]]
     /// ```
-    ///
     /// See ClickHouse <https://clickhouse.com/docs/en/sql-reference/statements/optimize>
+    ///
+    /// Databricks:
+    /// ```sql
+    /// OPTIMIZE table_name [WHERE predicate] [ZORDER BY (col_name1 [, ...])]
+    /// ```
+    /// See Databricks <https://docs.databricks.com/en/sql/language-manual/delta-optimize.html>
     OptimizeTable {
         /// Table name to optimize.
         name: ObjectName,
-        /// Optional cluster identifier.
+        /// Whether the `TABLE` keyword was present (ClickHouse uses `OPTIMIZE TABLE`, Databricks uses `OPTIMIZE`).
+        has_table_keyword: bool,
+        /// Optional cluster identifier (ClickHouse).
         on_cluster: Option<Ident>,
-        /// Optional partition spec.
+        /// Optional partition spec (ClickHouse).
         partition: Option<Partition>,
-        /// Whether `FINAL` was specified.
+        /// Whether `FINAL` was specified (ClickHouse).
         include_final: bool,
-        /// Optional deduplication settings.
+        /// Optional deduplication settings (ClickHouse).
         deduplicate: Option<Deduplicate>,
+        /// Optional WHERE predicate (Databricks).
+        predicate: Option<Expr>,
+        /// Optional ZORDER BY columns (Databricks).
+        zorder: Option<Vec<Expr>>,
     },
     /// ```sql
     /// LISTEN
@@ -6243,12 +6255,19 @@ impl fmt::Display for Statement {
             }
             Statement::OptimizeTable {
                 name,
+                has_table_keyword,
                 on_cluster,
                 partition,
                 include_final,
                 deduplicate,
+                predicate,
+                zorder,
             } => {
-                write!(f, "OPTIMIZE TABLE {name}")?;
+                write!(f, "OPTIMIZE")?;
+                if *has_table_keyword {
+                    write!(f, " TABLE")?;
+                }
+                write!(f, " {name}")?;
                 if let Some(on_cluster) = on_cluster {
                     write!(f, " ON CLUSTER {on_cluster}")?;
                 }
@@ -6261,6 +6280,12 @@ impl fmt::Display for Statement {
                 if let Some(deduplicate) = deduplicate {
                     write!(f, " {deduplicate}")?;
                 }
+                if let Some(predicate) = predicate {
+                    write!(f, " WHERE {predicate}")?;
+                }
+                if let Some(zorder) = zorder {
+                    write!(f, " ZORDER BY ({})", display_comma_separated(zorder))?;
+                }
                 Ok(())
             }
             Statement::LISTEN { channel } => {

src/parser/mod.rs

@@ -697,8 +697,10 @@ impl<'a> Parser<'a> {
                     self.parse_install()
                 }
                 Keyword::LOAD => self.parse_load(),
-                // `OPTIMIZE` is clickhouse specific https://clickhouse.tech/docs/en/sql-reference/statements/optimize/
-                Keyword::OPTIMIZE if dialect_of!(self is ClickHouseDialect | GenericDialect) => {
+                // `OPTIMIZE` is clickhouse/databricks specific
+                // ClickHouse: https://clickhouse.tech/docs/en/sql-reference/statements/optimize/
+                // Databricks: https://docs.databricks.com/en/sql/language-manual/delta-optimize.html
+                Keyword::OPTIMIZE if dialect_of!(self is ClickHouseDialect | DatabricksDialect | GenericDialect) => {
                     self.parse_optimize_table()
                 }
                 // `COMMENT` is snowflake specific https://docs.snowflake.com/en/sql-reference/sql/comment
@@ -18204,13 +18206,24 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// ClickHouse:
     /// ```sql
     /// OPTIMIZE TABLE [db.]name [ON CLUSTER cluster] [PARTITION partition | PARTITION ID 'partition_id'] [FINAL] [DEDUPLICATE [BY expression]]
     /// ```
     /// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/optimize)
+    ///
+    /// Databricks:
+    /// ```sql
+    /// OPTIMIZE table_name [WHERE predicate] [ZORDER BY (col_name1 [, ...])]
+    /// ```
+    /// [Databricks](https://docs.databricks.com/en/sql/language-manual/delta-optimize.html)
     pub fn parse_optimize_table(&mut self) -> Result<Statement, ParserError> {
-        self.expect_keyword_is(Keyword::TABLE)?;
+        // Check for TABLE keyword (ClickHouse uses it, Databricks does not)
+        let has_table_keyword = self.parse_keyword(Keyword::TABLE);
+
         let name = self.parse_object_name(false)?;
+
+        // ClickHouse-specific options
         let on_cluster = self.parse_optional_on_cluster()?;
 
         let partition = if self.parse_keyword(Keyword::PARTITION) {
@@ -18224,6 +18237,7 @@ impl<'a> Parser<'a> {
         };
 
         let include_final = self.parse_keyword(Keyword::FINAL);
+
         let deduplicate = if self.parse_keyword(Keyword::DEDUPLICATE) {
             if self.parse_keyword(Keyword::BY) {
                 Some(Deduplicate::ByExpression(self.parse_expr()?))
@@ -18234,12 +18248,31 @@ impl<'a> Parser<'a> {
             None
         };
 
+        // Databricks-specific options
+        let predicate = if self.parse_keyword(Keyword::WHERE) {
+            Some(self.parse_expr()?)
+        } else {
+            None
+        };
+
+        let zorder = if self.parse_keywords(&[Keyword::ZORDER, Keyword::BY]) {
+            self.expect_token(&Token::LParen)?;
+            let columns = self.parse_comma_separated(|p| p.parse_expr())?;
+            self.expect_token(&Token::RParen)?;
+            Some(columns)
+        } else {
+            None
+        };
+
         Ok(Statement::OptimizeTable {
             name,
+            has_table_keyword,
             on_cluster,
             partition,
             include_final,
             deduplicate,
+            predicate,
+            zorder,
         })
     }
 

tests/sqlparser_databricks.rs

@@ -389,3 +389,70 @@ fn parse_table_time_travel() {
         .parse_sql_statements("SELECT 1 FROM t1 VERSION AS OF 1 - 2",)
         .is_err())
 }
+
+#[test]
+fn parse_optimize_table() {
+    // Basic OPTIMIZE (Databricks style - no TABLE keyword)
+    databricks().verified_stmt("OPTIMIZE my_table");
+    databricks().verified_stmt("OPTIMIZE db.my_table");
+    databricks().verified_stmt("OPTIMIZE catalog.db.my_table");
+
+    // With WHERE clause
+    databricks().verified_stmt("OPTIMIZE my_table WHERE date = '2023-01-01'");
+    databricks().verified_stmt("OPTIMIZE my_table WHERE date >= '2023-01-01' AND date < '2023-02-01'");
+
+    // With ZORDER BY clause
+    databricks().verified_stmt("OPTIMIZE my_table ZORDER BY (col1)");
+    databricks().verified_stmt("OPTIMIZE my_table ZORDER BY (col1, col2)");
+    databricks().verified_stmt("OPTIMIZE my_table ZORDER BY (col1, col2, col3)");
+
+    // Combined WHERE and ZORDER BY
+    databricks().verified_stmt("OPTIMIZE my_table WHERE date = '2023-01-01' ZORDER BY (col1)");
+    databricks().verified_stmt("OPTIMIZE my_table WHERE date >= '2023-01-01' ZORDER BY (col1, col2)");
+
+    // Verify AST structure
+    match databricks()
+        .verified_stmt("OPTIMIZE my_table WHERE date = '2023-01-01' ZORDER BY (col1, col2)")
+    {
+        Statement::OptimizeTable {
+            name,
+            has_table_keyword,
+            on_cluster,
+            partition,
+            include_final,
+            deduplicate,
+            predicate,
+            zorder,
+        } => {
+            assert_eq!(name.to_string(), "my_table");
+            assert!(!has_table_keyword);
+            assert!(on_cluster.is_none());
+            assert!(partition.is_none());
+            assert!(!include_final);
+            assert!(deduplicate.is_none());
+            assert!(predicate.is_some());
+            assert_eq!(
+                zorder,
+                Some(vec![
+                    Expr::Identifier(Ident::new("col1")),
+                    Expr::Identifier(Ident::new("col2")),
+                ])
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    // Negative cases
+    assert_eq!(
+        databricks()
+            .parse_sql_statements("OPTIMIZE my_table ZORDER BY")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: (, found: EOF".to_string())
+    );
+    assert_eq!(
+        databricks()
+            .parse_sql_statements("OPTIMIZE my_table ZORDER BY ()")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: an expression, found: )".to_string())
+    );
+}