YQL-15767: Uniform representation of not null columns in v1 parser

2 files changed, 128 insertions, 2 deletions

diff --git a/ydb/library/yql/sql/v1/query.cpp b/ydb/library/yql/sql/v1/query.cpp
index 63591772b5..18862728ec 100644
--- a/ydb/library/yql/sql/v1/query.cpp
+++ b/ydb/library/yql/sql/v1/query.cpp
@@ -659,11 +659,16 @@ public:
                 columnsSet.insert(col.Name);
             }
 
+            THashSet<TString> pkColumns;
             for (auto& keyColumn : Params.PkColumns) {
                 if (!columnsSet.contains(keyColumn.Name)) {
                     ctx.Error(keyColumn.Pos) << "Undefined column: " << keyColumn.Name;
                     return false;
                 }
+                if (!pkColumns.insert(keyColumn.Name).second) {
+                    ctx.Error(keyColumn.Pos) << "Duplicated column in PK: " << keyColumn.Name;
+                    return false;
+                }
             }
             for (auto& keyColumn : Params.PartitionByColumns) {
                 if (!columnsSet.contains(keyColumn.Name)) {
@@ -741,13 +746,18 @@ public:
             opts = L(opts, Q(Y(Q("columnFamilies"), Q(columnFamilies))));
         }
 
+        THashSet<TString> notNullColumnsSet;
+        auto notNullColumns = Y();
         auto columns = Y();
         for (auto& col : Params.Columns) {
             auto columnDesc = Y();
             columnDesc = L(columnDesc, BuildQuotedAtom(Pos, col.Name));
             auto type = col.Type;
             if (col.Nullable) {
-                type = Y("OptionalType", type);
+                type = Y("AsOptionalType", type);
+            } else {
+                if (notNullColumnsSet.insert(col.Name).second)
+                    notNullColumns = L(notNullColumns, BuildQuotedAtom(Pos, col.Name));
             }
             columnDesc = L(columnDesc, type);
             if (col.Families) {
@@ -789,6 +799,10 @@ public:
             }
         }
 
+        if (!notNullColumnsSet.empty()) {
+            opts = L(opts, Q(Y(Q("notnull"), Q(notNullColumns))));
+        }
+
         if (!Params.PartitionByColumns.empty()) {
             auto partitionBy = Y();
             for (auto& col : Params.PartitionByColumns) {
@@ -959,7 +973,7 @@ public:
                 columnDesc = L(columnDesc, BuildQuotedAtom(Pos, col.Name));
                 auto type = col.Type;
                 if (col.Nullable) {
-                    type = Y("OptionalType", type);
+                    type = Y("AsOptionalType", type);
                 }
                 columnDesc = L(columnDesc, type);
                 if (col.Families) {
diff --git a/ydb/library/yql/sql/v1/sql_ut.cpp b/ydb/library/yql/sql/v1/sql_ut.cpp
index ff660fbd02..5a93b4a961 100644
--- a/ydb/library/yql/sql/v1/sql_ut.cpp
+++ b/ydb/library/yql/sql/v1/sql_ut.cpp
@@ -14,6 +14,9 @@ enum class EDebugOutput {
     ToCerr,
 };
 
+const ui32 PRETTY_FLAGS = NYql::TAstPrintFlags::PerLine | NYql::TAstPrintFlags::ShortQuote |
+                          NYql::TAstPrintFlags::AdaptArbitraryContent;
+
 TString Err2Str(NYql::TAstParseResult& res, EDebugOutput debug = EDebugOutput::None) {
     TStringStream s;
     res.Issues.PrintTo(s);
@@ -786,6 +789,115 @@ Y_UNIT_TEST_SUITE(SqlParsingOnly) {
         UNIT_ASSERT_VALUES_EQUAL(1, elementStat["primarykey"]);
     }
 
+    Y_UNIT_TEST(CreateTableNonNullableYqlTypeAstCorrect) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a int32 not null);");
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Write!") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                           line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (DataType 'Int32)))) '('notnull '('"a")))))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("Write!"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableNullableYqlTypeAstCorrect) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a int32);");
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Write!") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                           line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (AsOptionalType (DataType 'Int32))))))))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("Write!"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableNonNullablePgTypeAstCorrect) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a pg_int4 not null);");
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Write!") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                           line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (PgType '_int4)))) '('notnull '('"a")))))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("Write!"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableNullablePgTypeAstCorrect) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a pg_int4);");
+        UNIT_ASSERT(res.Root);
+
+            res.Root->PrettyPrintTo(Cout, PRETTY_FLAGS);
+
+            TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+                if (word == "Write!") {
+                    UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                               line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (AsOptionalType (PgType '_int4))))))))__"));
+                }
+            };
+
+            TWordCountHive elementStat = {{TString("Write!"), 0}};
+            VerifyProgram(res, elementStat, verifyLine);
+
+            UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableNullPkColumnsAreAllowed) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a int32, primary key(a));");
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Write!") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                           line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (AsOptionalType (DataType 'Int32))))) '('primarykey '('"a")))))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("Write!"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableNotNullPkColumnsAreIdempotentAstCorrect) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a int32 not null, primary key(a));");
+        UNIT_ASSERT(res.Root);
+
+        TVerifyLineFunc verifyLine = [](const TString& word, const TString& line) {
+            if (word == "Write!") {
+                UNIT_ASSERT_VALUES_UNEQUAL(TString::npos,
+                                           line.find(R"__((Write! world sink (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('"a" (DataType 'Int32)))) '('primarykey '('"a")) '('notnull '('"a")))))__"));
+            }
+        };
+
+        TWordCountHive elementStat = {{TString("Write!"), 0}};
+        VerifyProgram(res, elementStat, verifyLine);
+
+        UNIT_ASSERT_VALUES_EQUAL(1, elementStat["Write!"]);
+    }
+
+    Y_UNIT_TEST(CreateTableDuplicatedPkColumnsFail) {
+        NYql::TAstParseResult res = SqlToYql("USE plato; CREATE TABLE t (a int32 not null, primary key(a, a));");
+        UNIT_ASSERT(!res.Root);
+    }
+
     Y_UNIT_TEST(DeleteFromTableByKey) {
         NYql::TAstParseResult res = SqlToYql("delete from plato.Input where key = 200;", 10, "kikimr");
         UNIT_ASSERT(res.Root);