feat: add insert overwrite from spark sql

DerekStride · May 22, 2023 · 65b6f18 · 65b6f18
1 parent ef3df9c
commit 65b6f18
Show file tree

Hide file tree

Showing 2 changed files with 157 additions and 20 deletions.
diff --git a/grammar.js b/grammar.js
@@ -44,6 +44,7 @@ module.exports = grammar({
     keyword_replace: _ => make_keyword("replace"),
     keyword_update: _ => make_keyword("update"),
     keyword_into: _ => make_keyword("into"),
+    keyword_overwrite: _ => make_keyword("overwrite"),
     keyword_values: _ => make_keyword("values"),
     keyword_set: _ => make_keyword("set"),
     keyword_from: _ => make_keyword("from"),
@@ -1270,8 +1271,14 @@ module.exports = grammar({
         ),
       ),
       optional($.keyword_ignore),
-      optional($.keyword_into),
+      optional(
+        choice(
+          $.keyword_into,
+          $.keyword_overwrite, // Spark SQL
+        ),
+      ),
       $.table_reference,
+      optional($.table_partition), // Spark SQL
       optional(
         seq(
           $.keyword_as,
@@ -1418,26 +1425,35 @@ module.exports = grammar({
     ),
 
     table_partition: $ => seq(
-        choice(
-            // Postgres/MySQL style
-            seq(
-                $.keyword_partition,
-                $.keyword_by,
-                choice(
-                    $.keyword_range,
-                    $.keyword_hash,
-                )
-            ),
-            // Hive style
-            seq(
-                $.keyword_partitioned,
-                $.keyword_by,
-            ),
+      choice(
+        // Postgres/MySQL style
+        seq(
+          $.keyword_partition,
+          $.keyword_by,
+          choice(
+            $.keyword_range,
+            $.keyword_hash,
+          )
         ),
-        choice(
-            seq('(', $.identifier, ')'), // postgres
-            $.column_definitions // impala/hive
-        )
+        // Hive style
+        seq(
+          $.keyword_partitioned,
+          $.keyword_by,
+        ),
+        // Spark SQL
+        $.keyword_partition,
+      ),
+      choice(
+        seq('(', $.identifier, ')'), // postgres
+        $.column_definitions, // impala/hive
+        seq('(', $._key_value_pair, repeat(seq(',', $._key_value_pair)), ')',), // Spark SQL
+      )
+    ),
+
+    _key_value_pair: $ => seq(
+      field('key',$.identifier),
+      '=',
+      field('value', alias($._literal_string, $.literal)),
     ),
 
     stored_as: $ => seq(

diff --git a/test/corpus/insert.txt b/test/corpus/insert.txt
@@ -404,3 +404,124 @@ VALUES('foo','bar', 3);
         (literal)
         (literal)
         (literal)))))
+
+================================================================================
+Insert overwrite
+================================================================================
+
+INSERT OVERWRITE tab1
+SELECT
+  col1,
+  col2
+FROM
+(
+  SELECT
+    *
+  FROM
+    tab2
+  WHERE
+    key1 >= 'val'
+) a1;
+
+--------------------------------------------------------------------------------
+
+(program
+  (statement
+    (insert
+      (keyword_insert)
+      (keyword_overwrite)
+      (table_reference
+        (identifier))
+      (select
+        (keyword_select)
+        (select_expression
+          (term
+            (field
+              (identifier)))
+          (term
+            (field
+              (identifier)))))
+      (from
+        (keyword_from)
+        (relation
+          (subquery
+            (select
+              (keyword_select)
+              (select_expression
+                (all_fields)))
+            (from
+              (keyword_from)
+              (relation
+                (table_reference
+                  (identifier)))
+              (where
+                (keyword_where)
+                (binary_expression
+                  (field
+                    (identifier))
+                  (literal)))))
+          (identifier))))))
+
+================================================================================
+Insert overwrite with partition
+================================================================================
+
+INSERT OVERWRITE tab1
+PARTITION (key1 = 'val1', key2 = 'val2')
+SELECT
+  col1,
+  col2
+FROM
+(
+  SELECT
+    *
+  FROM
+    tab2
+  WHERE
+    key1 >= 'val'
+) a1;
+
+--------------------------------------------------------------------------------
+
+(program
+  (statement
+    (insert
+      (keyword_insert)
+      (keyword_overwrite)
+      (table_reference
+        (identifier))
+      (table_partition
+        (keyword_partition)
+        (identifier)
+        (literal)
+        (identifier)
+        (literal))
+      (select
+        (keyword_select)
+        (select_expression
+          (term
+            (field
+              (identifier)))
+          (term
+            (field
+              (identifier)))))
+      (from
+        (keyword_from)
+        (relation
+          (subquery
+            (select
+              (keyword_select)
+              (select_expression
+                (all_fields)))
+            (from
+              (keyword_from)
+              (relation
+                (table_reference
+                  (identifier)))
+              (where
+                (keyword_where)
+                (binary_expression
+                  (field
+                    (identifier))
+                  (literal)))))
+          (identifier))))))