feat(clickhouse): CREATE TABLE computed columns, column compression, …

…index (tobymao#3252) * feat(clickhouse): CREATE TABLE alias, col compression, index * Remove unnecessary generated options expr
1fanwang · Apr 1, 2024 · f88640b · f88640b
1 parent 28c5ee7
commit f88640b
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 3 deletions.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -261,6 +261,11 @@ class Parser(parser.Parser):
             "ArgMax",
         ]
 
+        FUNC_TOKENS = {
+            *parser.Parser.FUNC_TOKENS,
+            TokenType.SET,
+        }
+
         AGG_FUNC_MAPPING = (
             lambda functions, suffixes: {
                 f"{f}{sfx}": (f, sfx) for sfx in (suffixes + [""]) for f in functions
@@ -321,6 +326,17 @@ class Parser(parser.Parser):
             TokenType.FORMAT: lambda self: ("format", self._advance() or self._parse_id_var()),
         }
 
+        CONSTRAINT_PARSERS = {
+            **parser.Parser.CONSTRAINT_PARSERS,
+            "INDEX": lambda self: self._parse_index_constraint(),
+            "CODEC": lambda self: self._parse_compress(),
+        }
+
+        SCHEMA_UNNAMED_CONSTRAINTS = {
+            *parser.Parser.SCHEMA_UNNAMED_CONSTRAINTS,
+            "INDEX",
+        }
+
         def _parse_conjunction(self) -> t.Optional[exp.Expression]:
             this = super()._parse_conjunction()
 
@@ -512,6 +528,27 @@ def _parse_on_property(self) -> t.Optional[exp.Expression]:
                     self._retreat(index)
             return None
 
+        def _parse_index_constraint(
+            self, kind: t.Optional[str] = None
+        ) -> exp.IndexColumnConstraint:
+            # INDEX name1 expr TYPE type1(args) GRANULARITY value
+            this = self._parse_id_var()
+            expression = self._parse_conjunction()
+
+            index_type = self._match_text_seq("TYPE") and (
+                self._parse_function() or self._parse_var()
+            )
+
+            granularity = self._match_text_seq("GRANULARITY") and self._parse_term()
+
+            return self.expression(
+                exp.IndexColumnConstraint,
+                this=this,
+                expression=expression,
+                index_type=index_type,
+                granularity=granularity,
+            )
+
     class Generator(generator.Generator):
         QUERY_HINTS = False
         STRUCT_DELIMITER = ("(", ")")
@@ -590,6 +627,9 @@ class Generator(generator.Generator):
             exp.Array: inline_array_sql,
             exp.CastToStrType: rename_func("CAST"),
             exp.CountIf: rename_func("countIf"),
+            exp.CompressColumnConstraint: lambda self,
+            e: f"CODEC({self.expressions(e, key='this', flat=True)})",
+            exp.ComputedColumnConstraint: lambda self, e: f"ALIAS {self.sql(e, 'this')}",
             exp.CurrentDate: lambda self, e: self.func("CURRENT_DATE"),
             exp.DateAdd: date_delta_sql("DATE_ADD"),
             exp.DateDiff: date_delta_sql("DATE_DIFF"),
@@ -742,3 +782,15 @@ def createable_sql(self, expression: exp.Create, locations: t.DefaultDict) -> st
         def prewhere_sql(self, expression: exp.PreWhere) -> str:
             this = self.indent(self.sql(expression, "this"))
             return f"{self.seg('PREWHERE')}{self.sep()}{this}"
+
+        def indexcolumnconstraint_sql(self, expression: exp.IndexColumnConstraint) -> str:
+            this = self.sql(expression, "this")
+            this = f" {this}" if this else ""
+            expr = self.sql(expression, "expression")
+            expr = f" {expr}" if expr else ""
+            index_type = self.sql(expression, "index_type")
+            index_type = f" TYPE {index_type}" if index_type else ""
+            granularity = self.sql(expression, "granularity")
+            granularity = f" GRANULARITY {granularity}" if granularity else ""
+
+            return f"INDEX{this}{expr}{index_type}{granularity}"
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1665,13 +1665,16 @@ class GeneratedAsRowColumnConstraint(ColumnConstraintKind):
 
 
 # https://dev.mysql.com/doc/refman/8.0/en/create-table.html
+# https://github.com/ClickHouse/ClickHouse/blob/master/src/Parsers/ParserCreateQuery.h#L646
 class IndexColumnConstraint(ColumnConstraintKind):
     arg_types = {
         "this": False,
-        "schema": True,
+        "schema": False,
         "kind": False,
         "index_type": False,
         "options": False,
+        "expression": False,  # Clickhouse
+        "granularity": False,
     }
 
 

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -4504,7 +4504,7 @@ def _parse_column_def(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.
 
         constraints: t.List[exp.Expression] = []
 
-        if not kind and self._match(TokenType.ALIAS):
+        if (not kind and self._match(TokenType.ALIAS)) or self._match_text_seq("ALIAS"):
             constraints.append(
                 self.expression(
                     exp.ComputedColumnConstraint,

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -154,7 +154,9 @@ def test_clickhouse(self):
         self.validate_identity("TRUNCATE TABLE t1 ON CLUSTER test_cluster")
         self.validate_identity("TRUNCATE DATABASE db")
         self.validate_identity("TRUNCATE DATABASE db ON CLUSTER test_cluster")
-
+        self.validate_identity(
+            "CREATE TABLE t (foo String CODEC(LZ4HC(9), ZSTD, DELTA), size String ALIAS formatReadableSize(size_bytes), INDEX idx1 a TYPE bloom_filter(0.001) GRANULARITY 1, INDEX idx2 a TYPE set(100) GRANULARITY 2, INDEX idx3 a TYPE minmax GRANULARITY 3)"
+        )
         self.validate_all(
             "SELECT arrayJoin([1,2,3])",
             write={