From f3c3515e044f8384badc43000a44326acbec53a7 Mon Sep 17 00:00:00 2001
From: Matthias Queitsch <matthias.queitsch@mailbox.org>
Date: Sun, 18 Jun 2023 00:16:08 +0200
Subject: [PATCH] feat: optimize athena tables

---
 grammar.js               | 17 +++++++++++++++++
 test/corpus/optimize.txt | 28 ++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/grammar.js b/grammar.js
index bb346ea0..92253752 100644
--- a/grammar.js
+++ b/grammar.js
@@ -223,6 +223,9 @@ module.exports = grammar({
     keyword_compute: _ => make_keyword("compute"),
     keyword_stats: _ => make_keyword("stats"),
     keyword_statistics: _ => make_keyword("statistics"),
+    keyword_optimize: _ => make_keyword("optimize"),
+    keyword_rewrite: _ => make_keyword("rewrite"),
+    keyword_bin_pack: _ => make_keyword("bin_pack"),
     keyword_incremental: _ => make_keyword("incremental"),
     keyword_location: _ => make_keyword("location"),
     keyword_partitioned: _ => make_keyword("partitioned"),
@@ -1349,6 +1352,7 @@ module.exports = grammar({
 
     // Compute stats for Impala and Hive
     _compute_stats: $ => choice(
+      // Hive
       seq(
         $.keyword_analyze,
         $.keyword_table,
@@ -1370,6 +1374,7 @@ module.exports = grammar({
         ),
         optional($.keyword_noscan),
       ),
+      // Impala
       seq(
         $.keyword_compute,
         optional(
@@ -1384,6 +1389,18 @@ module.exports = grammar({
           )
         )
       ),
+      // Athena/Iceberg
+      seq(
+        $.keyword_optimize,
+        $.table_reference,
+        $.keyword_rewrite,
+        $.keyword_data,
+        $.keyword_using,
+        $.keyword_bin_pack,
+        optional(
+          $.where,
+        )
+      ),
     ),
 
     // TODO: this does not account for partitions specs like
diff --git a/test/corpus/optimize.txt b/test/corpus/optimize.txt
index 01b4046e..4eff6f78 100644
--- a/test/corpus/optimize.txt
+++ b/test/corpus/optimize.txt
@@ -68,3 +68,31 @@ NOSCAN
     (keyword_cache)
     (keyword_metadata)
     (keyword_noscan)))
+
+================================================================================
+Athena/Iceberg: Optimize table
+================================================================================
+
+OPTIMIZE mytable REWRITE DATA USING BIN_PACK
+  WHERE col1 is not null
+--------------------------------------------------------------------------------
+
+(program
+  (statement
+    (keyword_optimize)
+    (table_reference
+      (identifier))
+    (keyword_rewrite)
+    (keyword_data)
+    (keyword_using)
+    (keyword_bin_pack)
+    (where
+      (keyword_where)
+      (binary_expression
+        (field
+          (identifier))
+        (is_not
+          (keyword_is)
+          (keyword_not))
+        (literal
+          (keyword_null))))))