fix decimal format

liun03 · liun03 · commit df58bc6a4c82 · 2023-06-06T18:08:50.000+08:00
diff --git a/graphmdl-sqlrewrite/src/main/java/io/graphmdl/sqlrewrite/PreAggregationRewrite.java b/graphmdl-sqlrewrite/src/main/java/io/graphmdl/sqlrewrite/PreAggregationRewrite.java
@@ -49,7 +49,8 @@
 import static io.graphmdl.sqlrewrite.Utils.analyzeFrom;
 import static io.graphmdl.sqlrewrite.Utils.toCatalogSchemaTableName;
 import static io.trino.sql.QueryUtil.getQualifiedName;
-import static io.trino.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DOUBLE;
+import static io.trino.sql.SqlFormatter.Dialect.DUCKDB;
+import static io.trino.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DECIMAL;
 import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
@@ -69,12 +70,12 @@ public static Optional<String> rewrite(
             GraphMDL graphMDL)
     {
         try {
-            Statement statement = SQL_PARSER.createStatement(sql, new ParsingOptions(AS_DOUBLE));
+            Statement statement = SQL_PARSER.createStatement(sql, new ParsingOptions(AS_DECIMAL));
             PreAggregationAnalysis aggregationAnalysis = new PreAggregationAnalysis();
             Statement rewritten = (Statement) new Rewriter(sessionContext, converter, graphMDL, aggregationAnalysis).process(statement, Optional.empty());
             if (rewritten instanceof Query
                     && aggregationAnalysis.onlyPreAggregationTables()) {
-                return Optional.of(SqlFormatter.formatSql(rewritten));
+                return Optional.of(SqlFormatter.formatSql(rewritten, DUCKDB));
             }
         }
         catch (Exception e) {
diff --git a/graphmdl-sqlrewrite/src/test/java/io/graphmdl/TestPreAggregationRewrite.java b/graphmdl-sqlrewrite/src/test/java/io/graphmdl/TestPreAggregationRewrite.java
@@ -33,6 +33,7 @@
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static io.graphmdl.base.GraphMDLTypes.DATE;
+import static io.graphmdl.base.GraphMDLTypes.DECIMAL;
 import static io.graphmdl.base.GraphMDLTypes.INTEGER;
 import static io.graphmdl.base.GraphMDLTypes.TIMESTAMP;
 import static io.graphmdl.base.GraphMDLTypes.VARCHAR;
@@ -42,8 +43,9 @@
 import static io.graphmdl.base.dto.TimeGrain.TimeUnit.YEAR;
 import static io.graphmdl.base.dto.TimeGrain.timeGrain;
 import static io.graphmdl.testing.AbstractTestFramework.withDefaultCatalogSchema;
+import static io.trino.sql.SqlFormatter.Dialect.DUCKDB;
 import static io.trino.sql.SqlFormatter.formatSql;
-import static io.trino.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DOUBLE;
+import static io.trino.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DECIMAL;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -95,7 +97,7 @@ public void init()
                                 List.of(
                                         column("author", VARCHAR, null, true),
                                         column("album_name", VARCHAR, null, true, "Album.name")),
-                                List.of(column("price", INTEGER, null, true, "avg(Album.price)")),
+                                List.of(column("price", DECIMAL, null, true, "avg(Album.price)")),
                                 List.of(
                                         timeGrain("p_date", "Album.publish_date", List.of(YEAR)),
                                         timeGrain("r_date", "Album.release_date", List.of(YEAR))),
@@ -300,6 +302,16 @@ public void testTableAliasScope()
                 "with test_a as (with AvgCollection as (select * from table_Collection) select * from AvgCollection) select * from table_AvgCollection");
     }
 
+    @Test
+    public void testDecimalRewrite()
+    {
+        assertRewrite(
+                "SELECT * from AvgCollection where avg = DECIMAL '1.0'",
+                "graphmdl",
+                "test",
+                "SELECT * FROM table_AvgCollection WHERE avg = 1.0");
+    }
+
     @DataProvider(name = "unexpectedStatementProvider")
     public Object[][] unexpectedStatementProvider()
     {
@@ -395,9 +407,9 @@ private void assertRewrite(
                 defaultSchema,
                 tableConverter).orElseThrow(() -> new AssertionError("No rewrite result"));
 
-        Statement expect = sqlParser.createStatement(expectSql, new ParsingOptions(AS_DOUBLE));
-        Statement actualStatement = sqlParser.createStatement(result, new ParsingOptions(AS_DOUBLE));
-        assertThat(result).isEqualTo(formatSql(expect));
+        Statement expect = sqlParser.createStatement(expectSql, new ParsingOptions(AS_DECIMAL));
+        Statement actualStatement = sqlParser.createStatement(result, new ParsingOptions(AS_DECIMAL));
+        assertThat(result).isEqualTo(formatSql(expect, DUCKDB));
         assertThat(actualStatement).isEqualTo(expect);
     }
 
diff --git a/trino-parser/src/main/java/io/trino/sql/ExpressionFormatter.java b/trino-parser/src/main/java/io/trino/sql/ExpressionFormatter.java
@@ -115,6 +115,7 @@
 import static io.trino.sql.RowPatternFormatter.formatPattern;
 import static io.trino.sql.SqlFormatter.Dialect.BIGQUERY;
 import static io.trino.sql.SqlFormatter.Dialect.DEFAULT;
+import static io.trino.sql.SqlFormatter.Dialect.DUCKDB;
 import static io.trino.sql.SqlFormatter.formatName;
 import static io.trino.sql.SqlFormatter.formatSql;
 import static io.trino.sql.tree.ComparisonExpression.Operator.EQUAL;
@@ -138,7 +139,7 @@ public static String formatExpression(Expression expression, Dialect dialect)
 
     private static String formatIdentifier(String s, Dialect dialect)
     {
-        if (dialect == DEFAULT) {
+        if (dialect == DEFAULT || dialect == DUCKDB) {
             return '"' + s.replace("\"", "\"\"") + '"';
         }
         else if (dialect == BIGQUERY) {
@@ -288,7 +289,7 @@ protected String visitArrayConstructor(ArrayConstructor node, Void context)
         protected String visitSubscriptExpression(SubscriptExpression node, Void context)
         {
             String subscript;
-            if (dialect == DEFAULT) {
+            if (dialect == DEFAULT || dialect == DUCKDB) {
                 subscript = formatSql(node.getIndex(), dialect);
             }
             else if (dialect == BIGQUERY) {
@@ -313,12 +314,18 @@ protected String visitLongLiteral(LongLiteral node, Void context)
         @Override
         protected String visitDoubleLiteral(DoubleLiteral node, Void context)
         {
+            if (dialect == DUCKDB) {
+                return String.valueOf(node.getValue());
+            }
             return doubleFormatter.get().format(node.getValue());
         }
 
         @Override
         protected String visitDecimalLiteral(DecimalLiteral node, Void context)
         {
+            if (dialect == DUCKDB) {
+                return node.getValue();
+            }
             // TODO return node value without "DECIMAL '..'" when FeaturesConfig#parseDecimalLiteralsAsDouble switch is removed
             return "DECIMAL '" + node.getValue() + "'";
         }
@@ -338,6 +345,9 @@ protected String visitTimeLiteral(TimeLiteral node, Void context)
         @Override
         protected String visitTimestampLiteral(TimestampLiteral node, Void context)
         {
+            if (dialect == DUCKDB) {
+                return "'" + node.getValue() + "'";
+            }
             return "TIMESTAMP '" + node.getValue() + "'";
         }
 
@@ -612,7 +622,7 @@ protected String visitLikePredicate(LikePredicate node, Void context)
                     .append(" LIKE ")
                     .append(process(node.getPattern(), context));
 
-            if (dialect == DEFAULT) {
+            if (dialect == DEFAULT || dialect == DUCKDB) {
                 node.getEscape().ifPresent(escape -> builder.append(" ESCAPE ")
                         .append(process(escape, context)));
             }
@@ -936,7 +946,7 @@ static String formatStringLiteral(String s, Dialect dialect)
         }
 
         StringBuilder builder = new StringBuilder();
-        if (dialect == DEFAULT) {
+        if (dialect == DEFAULT || dialect == DUCKDB) {
             builder.append("U&");
         }
         builder.append("'");
@@ -952,7 +962,7 @@ static String formatStringLiteral(String s, Dialect dialect)
                 builder.append(ch);
             }
             else {
-                if (dialect == DEFAULT) {
+                if (dialect == DEFAULT || dialect == DUCKDB) {
                     if (codePoint <= 0xFFFF) {
                         builder.append('\\');
                         builder.append(format("%04X", codePoint));
diff --git a/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java b/trino-parser/src/main/java/io/trino/sql/SqlFormatter.java
@@ -165,6 +165,7 @@ public enum Dialect
     {
         DEFAULT,
         BIGQUERY,
+        DUCKDB,
     }
 
     private SqlFormatter() {}

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,7 @@`
`115`	`115`	`import static io.trino.sql.RowPatternFormatter.formatPattern;`
`116`	`116`	`import static io.trino.sql.SqlFormatter.Dialect.BIGQUERY;`
`117`	`117`	`import static io.trino.sql.SqlFormatter.Dialect.DEFAULT;`
	`118`	`+import static io.trino.sql.SqlFormatter.Dialect.DUCKDB;`
`118`	`119`	`import static io.trino.sql.SqlFormatter.formatName;`
`119`	`120`	`import static io.trino.sql.SqlFormatter.formatSql;`
`120`	`121`	`import static io.trino.sql.tree.ComparisonExpression.Operator.EQUAL;`
`@@ -138,7 +139,7 @@ public static String formatExpression(Expression expression, Dialect dialect)`
`138`	`139`
`139`	`140`	`private static String formatIdentifier(String s, Dialect dialect)`
`140`	`141`	`{`
`141`		`- if (dialect == DEFAULT) {`
	`142`	`+ if (dialect == DEFAULT \|\| dialect == DUCKDB) {`
`142`	`143`	`return '"' + s.replace("\"", "\"\"") + '"';`
`143`	`144`	`}`
`144`	`145`	`else if (dialect == BIGQUERY) {`
`@@ -288,7 +289,7 @@ protected String visitArrayConstructor(ArrayConstructor node, Void context)`
`288`	`289`	`protected String visitSubscriptExpression(SubscriptExpression node, Void context)`
`289`	`290`	`{`
`290`	`291`	`String subscript;`
`291`		`- if (dialect == DEFAULT) {`
	`292`	`+ if (dialect == DEFAULT \|\| dialect == DUCKDB) {`
`292`	`293`	`subscript = formatSql(node.getIndex(), dialect);`
`293`	`294`	`}`
`294`	`295`	`else if (dialect == BIGQUERY) {`
`@@ -313,12 +314,18 @@ protected String visitLongLiteral(LongLiteral node, Void context)`
`313`	`314`	`@Override`
`314`	`315`	`protected String visitDoubleLiteral(DoubleLiteral node, Void context)`
`315`	`316`	`{`
	`317`	`+ if (dialect == DUCKDB) {`
	`318`	`+ return String.valueOf(node.getValue());`
	`319`	`+ }`
`316`	`320`	`return doubleFormatter.get().format(node.getValue());`
`317`	`321`	`}`
`318`	`322`
`319`	`323`	`@Override`
`320`	`324`	`protected String visitDecimalLiteral(DecimalLiteral node, Void context)`
`321`	`325`	`{`
	`326`	`+ if (dialect == DUCKDB) {`
	`327`	`+ return node.getValue();`
	`328`	`+ }`
`322`	`329`	`// TODO return node value without "DECIMAL '..'" when FeaturesConfig#parseDecimalLiteralsAsDouble switch is removed`
`323`	`330`	`return "DECIMAL '" + node.getValue() + "'";`
`324`	`331`	`}`
`@@ -338,6 +345,9 @@ protected String visitTimeLiteral(TimeLiteral node, Void context)`
`338`	`345`	`@Override`
`339`	`346`	`protected String visitTimestampLiteral(TimestampLiteral node, Void context)`
`340`	`347`	`{`
	`348`	`+ if (dialect == DUCKDB) {`
	`349`	`+ return "'" + node.getValue() + "'";`
	`350`	`+ }`
`341`	`351`	`return "TIMESTAMP '" + node.getValue() + "'";`
`342`	`352`	`}`
`343`	`353`
`@@ -612,7 +622,7 @@ protected String visitLikePredicate(LikePredicate node, Void context)`
`612`	`622`	`.append(" LIKE ")`
`613`	`623`	`.append(process(node.getPattern(), context));`
`614`	`624`
`615`		`- if (dialect == DEFAULT) {`
	`625`	`+ if (dialect == DEFAULT \|\| dialect == DUCKDB) {`
`616`	`626`	`node.getEscape().ifPresent(escape -> builder.append(" ESCAPE ")`
`617`	`627`	`.append(process(escape, context)));`
`618`	`628`	`}`
`@@ -936,7 +946,7 @@ static String formatStringLiteral(String s, Dialect dialect)`
`936`	`946`	`}`
`937`	`947`
`938`	`948`	`StringBuilder builder = new StringBuilder();`
`939`		`- if (dialect == DEFAULT) {`
	`949`	`+ if (dialect == DEFAULT \|\| dialect == DUCKDB) {`
`940`	`950`	`builder.append("U&");`
`941`	`951`	`}`
`942`	`952`	`builder.append("'");`
`@@ -952,7 +962,7 @@ static String formatStringLiteral(String s, Dialect dialect)`
`952`	`962`	`builder.append(ch);`
`953`	`963`	`}`
`954`	`964`	`else {`
`955`		`- if (dialect == DEFAULT) {`
	`965`	`+ if (dialect == DEFAULT \|\| dialect == DUCKDB) {`
`956`	`966`	`if (codePoint <= 0xFFFF) {`
`957`	`967`	`builder.append('\\');`
`958`	`968`	`builder.append(format("%04X", codePoint));`
Original file line number	Diff line number	Diff line change
`@@ -165,6 +165,7 @@ public enum Dialect`
`165`	`165`	`{`
`166`	`166`	`DEFAULT,`
`167`	`167`	`BIGQUERY,`
	`168`	`+ DUCKDB,`
`168`	`169`	`}`
`169`	`170`
`170`	`171`	`private SqlFormatter() {}`