Refactor clickhouse partition by.

alibaba · Aug 9, 2024 · 0dbf76d · 0dbf76d
1 parent cfc1ac5
commit 0dbf76d
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 26 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/ast/CKCreateTableStatement.java b/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/ast/CKCreateTableStatement.java
@@ -15,8 +15,6 @@
 public class CKCreateTableStatement extends SQLCreateTableStatement {
     protected final List<SQLAssignItem> settings = new ArrayList<SQLAssignItem>();
     private SQLOrderBy orderBy;
-    private SQLExpr partitionBy;
-
     private SQLPrimaryKey primaryKey;
     private SQLExpr sampleBy;
 
@@ -40,18 +38,6 @@ public void setOrderBy(SQLOrderBy x) {
         this.orderBy = x;
     }
 
-    public SQLExpr getPartitionBy() {
-        return partitionBy;
-    }
-
-    public void setPartitionBy(SQLExpr x) {
-        if (x != null) {
-            x.setParent(this);
-        }
-
-        this.partitionBy = x;
-    }
-
     public SQLExpr getSampleBy() {
         return sampleBy;
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/parser/CKCreateTableParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/parser/CKCreateTableParser.java
@@ -2,6 +2,8 @@
 
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLOrderBy;
+import com.alibaba.druid.sql.ast.SQLPartitionBy;
+import com.alibaba.druid.sql.ast.SQLPartitionByList;
 import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
 import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement;
 import com.alibaba.druid.sql.ast.statement.SQLPrimaryKey;
@@ -21,6 +23,30 @@ protected SQLCreateTableStatement newCreateStatement() {
         return new CKCreateTableStatement();
     }
 
+    @Override
+    public SQLPartitionBy parsePartitionBy() {
+        lexer.nextToken();
+        accept(Token.BY);
+        SQLPartitionBy sqlPartitionBy = new SQLPartitionByList();
+        boolean hasParen = false;
+        if (lexer.nextIf(Token.LPAREN)) {
+            hasParen = true;
+        }
+        for (; ; ) {
+            sqlPartitionBy.addColumn(this.exprParser.expr());
+            if (lexer.token() == Token.COMMA) {
+                lexer.nextToken();
+                continue;
+            }
+            break;
+        }
+        if (hasParen) {
+            accept(Token.RPAREN);
+        }
+
+        return sqlPartitionBy;
+    }
+
     protected void parseCreateTableRest(SQLCreateTableStatement stmt) {
         CKCreateTableStatement ckStmt = (CKCreateTableStatement) stmt;
         if (lexer.identifierEquals(FnvHash.Constants.ENGINE)) {
@@ -34,10 +60,7 @@ protected void parseCreateTableRest(SQLCreateTableStatement stmt) {
         }
 
         if (lexer.token() == Token.PARTITION) {
-            lexer.nextToken();
-            accept(Token.BY);
-            SQLExpr expr = this.exprParser.expr();
-            ckStmt.setPartitionBy(expr);
+            ckStmt.setPartitionBy(parsePartitionBy());
         }
 
         if (lexer.token() == Token.ORDER) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/visitor/CKOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/clickhouse/visitor/CKOutputVisitor.java
@@ -87,16 +87,31 @@ public boolean visit(SQLStructDataType.Field x) {
         return false;
     }
 
+    @Override
+    public boolean visit(SQLPartitionByList x) {
+        if (x.getColumns().size() == 1) {
+            x.getColumns().get(0).accept(this);
+        } else {
+            print('(');
+            printAndAccept(x.getColumns(), ", ");
+            print0(")");
+        }
+
+        printPartitionsCountAndSubPartitions(x);
+
+        printSQLPartitions(x.getPartitions());
+        return false;
+    }
     @Override
     public boolean visit(CKCreateTableStatement x) {
         super.visit((SQLCreateTableStatement) x);
 
-        SQLExpr partitionBy = x.getPartitionBy();
-        if (partitionBy != null) {
-            println();
-            print0(ucase ? "PARTITION BY " : "partition by ");
-            partitionBy.accept(this);
-        }
+//        SQLPartitionBy partitionBy = x.getPartitioning();
+//        if (partitionBy != null) {
+//            println();
+//            print0(ucase ? "PARTITION BY " : "partition by ");
+//            partitionBy.accept(this);
+//        }
 
         SQLOrderBy orderBy = x.getOrderBy();
         if (orderBy != null) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -3810,6 +3810,7 @@ public boolean visit(SQLCreateTableStatement x) {
             print0(ucase ? "PARTITION OF " : "partition of ");
             partitionOf.accept(this);
         }
+        printEngine(x);
         printPartitionBy(x);
         printTableOptions(x);
 
@@ -3819,7 +3820,6 @@ public boolean visit(SQLCreateTableStatement x) {
             print0(ucase ? "TABLESPACE " : "tablespace ");
             tablespace.accept(this);
         }
-        printEngine(x);
         SQLSelect select = x.getSelect();
         if (select != null) {
             println();

diff --git a/core/src/test/resources/bvt/parser/clickhouse/4.txt b/core/src/test/resources/bvt/parser/clickhouse/4.txt
@@ -64,4 +64,37 @@ CREATE TABLE db.replicated_orders ON CLUSTER test (
 	size String
 ) ENGINE = ReplicatedMergeTree('/clickhouse/tables/{shard}/replicated_orders', '{replica}')
 PARTITION BY order_date
-ORDER BY (order_date, customer_id, order_id)
+ORDER BY (order_date, customer_id, order_id)
+------------------------------------------------------------------------------------------------------------------------
+CREATE TABLE replicated_orders (
+                    order_id UInt64 CODEC(ZSTD),
+                    customer_id LowCardinality(UInt32),
+                    order_date Date,
+                    product_id LowCardinality(UInt32),
+                    quantity Int32,
+                    price Decimal(10, 2),
+                    array1 Array(UInt8),
+                    tuple1 Tuple(UInt8, String),
+                    map1 Map(String, UInt64),
+                    variant1 Variant(String, Array(UInt8)),
+                    null1 Nullable(String)
+                )
+                ENGINE = ReplicatedMergeTree('/clickhouse/tables/{shard}/replicated_orders', '{replica}')
+                PARTITION BY order_date, quantity
+                ORDER BY (order_date, customer_id, order_id)
+--------------------
+CREATE TABLE replicated_orders (
+	order_id UInt64 CODEC(ZSTD),
+	customer_id LowCardinality(UInt32),
+	order_date Date,
+	product_id LowCardinality(UInt32),
+	quantity Int32,
+	price Decimal(10, 2),
+	array1 Array(UInt8),
+	tuple1 Tuple(UInt8, String),
+	map1 MAP(String, UInt64),
+	variant1 Variant(String, Array(UInt8)),
+	null1 Nullable(String)
+) ENGINE = ReplicatedMergeTree('/clickhouse/tables/{shard}/replicated_orders', '{replica}')
+PARTITION BY (order_date, quantity)
+ORDER BY (order_date, customer_id, order_id)