Suport for DuckDb Union data types

TylerBrinks · Aug 23, 2024 · db02202 · db02202
1 parent c95800b
commit db02202
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 4 deletions.
diff --git a/src/SqlParser.Tests/Dialects/DuckDbDialectTests.cs b/src/SqlParser.Tests/Dialects/DuckDbDialectTests.cs
@@ -333,4 +333,28 @@ public void Test_Array_Index()
 
         Assert.Equal(expected, expr);
     }
+
+    [Fact]
+    public void Test_DuckDb_Union_Datatype()
+    {
+        const string sql = "CREATE TABLE tbl1 (one UNION(a INT), two UNION(a INT, b INT), nested UNION(a UNION(b INT)))";
+        var create = VerifiedStatement<Statement.CreateTable>(sql);
+
+        var expected = new Statement.CreateTable(new CreateTable("tbl1", [
+            new ColumnDef("one", new DataType.Union([
+                new UnionField("a", new DataType.Int())
+            ])),
+            new ColumnDef("two", new DataType.Union([
+                new UnionField("a", new DataType.Int()),
+                new UnionField("b", new DataType.Int())
+            ])),
+            new ColumnDef("nested", new DataType.Union([
+                new UnionField("a", new DataType.Union([
+                        new UnionField("b", new DataType.Int())
+                    ]))
+            ])),
+        ]));
+
+        Assert.Equal(expected, create);
+    }
 }
diff --git a/src/SqlParser/Ast/DataType.cs b/src/SqlParser/Ast/DataType.cs
@@ -637,8 +637,8 @@ public override void ToSql(SqlTextWriter writer)
     ///
     /// Map Clickhouse: https://clickhouse.com/docs/en/sql-reference/data-types/map
     /// </summary>
-    /// <param name="dt1"></param>
-    /// <param name="dt2"></param>
+    /// <param name="KeyDataType"></param>
+    /// <param name="ValueDataType"></param>
     public record Map(DataType KeyDataType, DataType ValueDataType) : DataType
     {
         public override void ToSql(SqlTextWriter writer)
@@ -773,7 +773,7 @@ public override void ToSql(SqlTextWriter writer)
     }
     /// Struct
     ///
-    /// Bive: https://docs.cloudera.com/cdw-runtime/cloud/impala-sql-reference/topics/impala-struct.html
+    /// Hive: https://docs.cloudera.com/cdw-runtime/cloud/impala-sql-reference/topics/impala-struct.html
     /// BigQuery: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#struct_type
     public record Struct(Sequence<StructField> Fields) : DataType
     {
@@ -915,6 +915,19 @@ public override void ToSql(SqlTextWriter writer)
         }
     }
     /// <summary>
+    /// Union
+    ///
+    /// DuckDb https://duckdb.org/docs/sql/data_types/union.html
+    /// </summary>
+    /// <param name="Values"></param>
+    public record Union(Sequence<UnionField> Fields) : DataType
+    {
+        public override void ToSql(SqlTextWriter writer)
+        {
+            writer.WriteSql($"UNION({Fields.ToSqlDelimited()})");
+        }
+    }
+    /// <summary>
     /// Unsigned big integer with optional display width e.g. BIGINT UNSIGNED or BIGINT(20) UNSIGNED
     /// </summary>
     public record UnsignedBigInt(ulong? Length = null) : LengthDataType(Length)

diff --git a/src/SqlParser/Ast/UnionField.cs b/src/SqlParser/Ast/UnionField.cs
@@ -0,0 +1,9 @@
+namespace SqlParser.Ast;
+
+public record UnionField(Ident FieldName, DataType FieldType) : IWriteSql, IElement
+{
+    public void ToSql(SqlTextWriter writer)
+    {
+        writer.WriteSql($"{FieldName} {FieldType}");
+    }
+}
diff --git a/src/SqlParser/Parser.cs b/src/SqlParser/Parser.cs
@@ -3874,6 +3874,21 @@ public Statement.Declare ParseMsSqlDeclare()
         return new Statement.Declare(statements);
     }
 
+    public Sequence<UnionField> ParseUnionTypeDef()
+    {
+        ExpectKeyword(Keyword.UNION);
+
+        return ExpectParens(() =>
+        {
+            return ParseCommaSeparated(() =>
+            {
+                var identifier = ParseIdentifier();
+                var fieldType = ParseDataType();
+                return new UnionField(identifier, fieldType);
+            });
+        });
+    }
+
     public DeclareAssignment? ParseMsSqlVariableDeclarationExpression()
     {
         if (PeekToken() is Equal)
@@ -6019,7 +6034,7 @@ public Value ParseIntroducedStringValue()
             Word { Keyword: Keyword.SET } => new DataType.Set(ParseStringValue()),
             Word { Keyword: Keyword.ARRAY } => ParseArray(),
             Word { Keyword: Keyword.STRUCT } when _dialect is BigQueryDialect or GenericDialect => ParseStruct(),
-
+            Word { Keyword: Keyword.UNION } when _dialect is DuckDbDialect or GenericDialect => ParseUnion(), 
             Word { Keyword: Keyword.NULLABLE } when _dialect is ClickHouseDialect or GenericDialect =>
                 ParseSubtype(child => new DataType.Nullable(child)),
             Word { Keyword: Keyword.LOWCARDINALITY } when _dialect is ClickHouseDialect or GenericDialect =>
@@ -6203,6 +6218,12 @@ DataType ParseStruct()
             return new DataType.Struct(fieldDefinitions);
         }
 
+        DataType ParseUnion()
+        {
+            PrevToken();
+            return new DataType.Union(ParseUnionTypeDef());
+        }
+
         DataType ParseUnmatched()
         {
             PrevToken();