Feat(duckdb): add support for SUMMARIZE (#3840)

georgesittas · web-flow · commit 45f45eaaac5a · 2024-07-26T19:35:42.000+03:00
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
@@ -238,14 +238,15 @@ class Tokenizer(tokens.Tokenizer):
             "POSITIONAL": TokenType.POSITIONAL,
             "SIGNED": TokenType.INT,
             "STRING": TokenType.TEXT,
-            "UBIGINT": TokenType.UBIGINT,
-            "UINTEGER": TokenType.UINT,
-            "USMALLINT": TokenType.USMALLINT,
-            "UTINYINT": TokenType.UTINYINT,
+            "SUMMARIZE": TokenType.SUMMARIZE,
             "TIMESTAMP_S": TokenType.TIMESTAMP_S,
             "TIMESTAMP_MS": TokenType.TIMESTAMP_MS,
             "TIMESTAMP_NS": TokenType.TIMESTAMP_NS,
             "TIMESTAMP_US": TokenType.TIMESTAMP,
+            "UBIGINT": TokenType.UBIGINT,
+            "UINTEGER": TokenType.UINT,
+            "USMALLINT": TokenType.USMALLINT,
+            "UTINYINT": TokenType.UTINYINT,
             "VARCHAR": TokenType.TEXT,
         }
         KEYWORDS.pop("/*+")
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1439,6 +1439,11 @@ class Describe(Expression):
     arg_types = {"this": True, "style": False, "kind": False, "expressions": False}
 
 
+# https://duckdb.org/docs/guides/meta/summarize.html
+class Summarize(Expression):
+    arg_types = {"this": True, "table": False}
+
+
 class Kill(Expression):
     arg_types = {"this": True, "kind": False}
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -4067,3 +4067,7 @@ def pad_sql(self, expression: exp.Pad) -> str:
             fill_pattern = "' '"
 
         return self.func(f"{prefix}PAD", expression.this, expression.expression, fill_pattern)
+
+    def summarize_sql(self, expression: exp.Summarize) -> str:
+        table = " TABLE" if expression.args.get("table") else ""
+        return f"SUMMARIZE{table} {self.sql(expression.this)}"
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -2825,12 +2825,14 @@ def _parse_select(
             this = self._parse_derived_table_values()
         elif from_:
             this = exp.select("*").from_(from_.this, copy=False)
+        elif self._match(TokenType.SUMMARIZE):
+            table = self._match(TokenType.TABLE)
+            this = self._parse_select() or self._parse_string() or self._parse_table()
+            return self.expression(exp.Summarize, this=this, table=table)
         else:
             this = None
 
-        if parse_set_operation:
-            return self._parse_set_operations(this)
-        return this
+        return self._parse_set_operations(this) if parse_set_operation else this
 
     def _parse_with(self, skip_with_token: bool = False) -> t.Optional[exp.With]:
         if not skip_with_token and not self._match(TokenType.WITH):
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -364,6 +364,7 @@ class TokenType(AutoName):
     STORAGE_INTEGRATION = auto()
     STRAIGHT_JOIN = auto()
     STRUCT = auto()
+    SUMMARIZE = auto()
     TABLE_SAMPLE = auto()
     TAG = auto()
     TEMPORARY = auto()
diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
@@ -293,6 +293,12 @@ def test_duckdb(self):
         self.validate_identity("x -> '$.family'")
         self.validate_identity("CREATE TABLE color (name ENUM('RED', 'GREEN', 'BLUE'))")
         self.validate_identity("SELECT * FROM foo WHERE bar > $baz AND bla = $bob")
+        self.validate_identity("SUMMARIZE tbl").assert_is(exp.Summarize)
+        self.validate_identity("SUMMARIZE SELECT * FROM tbl").assert_is(exp.Summarize)
+        self.validate_identity("CREATE TABLE tbl_summary AS SELECT * FROM (SUMMARIZE tbl)")
+        self.validate_identity(
+            "SUMMARIZE TABLE 'https://blobs.duckdb.org/data/Star_Trek-Season_1.csv'"
+        ).assert_is(exp.Summarize)
         self.validate_identity(
             "SELECT * FROM x LEFT JOIN UNNEST(y)", "SELECT * FROM x LEFT JOIN UNNEST(y) ON TRUE"
         )
