[] Support NOT NULL in positional SCHEMA

5 files changed, 44 insertions, 7 deletions

diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 7b51ab5fac6..37163ef819d 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -552,11 +552,22 @@ namespace NTypeAnnImpl {
         }
 
         auto type = input->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
-        if (!EnsureDataType(input->Head().Pos(), *type, ctx.Expr)) {
+        if (input->IsCallable("DataOrOptionalData")) {
+            bool isOptional;
+            const TDataExprType* dataType;
+            if (!EnsureDataOrOptionalOfData(input->Pos(), type, isOptional, dataType, ctx.Expr)) {
+                return IGraphTransformer::TStatus::Error;
+            }
+            output = ctx.Expr.NewCallable(input->Pos(), dataType->GetName(), { input->ChildPtr(1) });
+            if (isOptional) {
+                output = ctx.Expr.NewCallable(input->Pos(), "Just", { output });
+            }
+        } else if (EnsureDataType(input->Head().Pos(), *type, ctx.Expr)) {
+            output = ctx.Expr.NewCallable(input->Pos(), type->Cast<TDataExprType>()->GetName(), { input->ChildPtr(1) });
+        } else {
             return IGraphTransformer::TStatus::Error;
         }
 
-        output = ctx.Expr.NewCallable(input->Pos(), type->Cast<TDataExprType>()->GetName(), { input->ChildPtr(1) });
         return IGraphTransformer::TStatus::Repeat;
     }
 
@@ -4835,6 +4846,23 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         return IGraphTransformer::TStatus::Ok;
     }
 
+    IGraphTransformer::TStatus AsOptionalTypeWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (!EnsureArgsCount(*input, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (auto status = EnsureTypeRewrite(input->HeadRef(), ctx.Expr); status != IGraphTransformer::TStatus::Ok) {
+            return status;
+        }
+
+        if (input->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->IsOptionalOrNull()) {
+            output = input->HeadPtr();
+        } else {
+            output = ctx.Expr.NewCallable(input->Pos(), "OptionalType", { input->HeadPtr() });
+        }
+        return IGraphTransformer::TStatus::Repeat;
+    }
+
     IGraphTransformer::TStatus OptionalWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         Y_UNUSED(output);
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -10932,6 +10960,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
 
     TSyncFunctionsMap::TSyncFunctionsMap() {
         Functions["Data"] = &DataWrapper;
+        Functions["DataOrOptionalData"] = &DataWrapper;
         Functions["DataSource"] = &DataSourceWrapper;
         Functions["Key"] = &KeyWrapper;
         Functions[LeftName] = &LeftWrapper;
@@ -11152,6 +11181,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["StreamType"] = &TypeWrapper<ETypeAnnotationKind::Stream>;
         Functions["FlowType"] = &TypeWrapper<ETypeAnnotationKind::Flow>;
         Functions["Nothing"] = &NothingWrapper;
+        Functions["AsOptionalType"] = &AsOptionalTypeWrapper;
         Functions["List"] = &ListWrapper;
         Functions["DictType"] = &TypeWrapper<ETypeAnnotationKind::Dict>;
         Functions["Dict"] = &DictWrapper;
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_io_discovery.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_io_discovery.cpp
index dc0a7fcc594..f0bb1f488f2 100644
--- a/ydb/library/yql/providers/s3/provider/yql_s3_io_discovery.cpp
+++ b/ydb/library/yql/providers/s3/provider/yql_s3_io_discovery.cpp
@@ -190,7 +190,7 @@ public:
                             ctx.Builder(object.Pos())
                                 .List()
                                     .Atom(0, col)
-                                    .Callable(1, "Data")
+                                    .Callable(1, "DataOrOptionalData")
                                         .Callable(0, "StructMemberType")
                                             .Add(0, generatedColumnsConfig->SchemaTypeNode)
                                             .Atom(1, col)
diff --git a/ydb/library/yql/sql/v1/SQLv1.g.in b/ydb/library/yql/sql/v1/SQLv1.g.in
index 36f9b7a3f15..2e8f1f6b20c 100644
--- a/ydb/library/yql/sql/v1/SQLv1.g.in
+++ b/ydb/library/yql/sql/v1/SQLv1.g.in
@@ -254,7 +254,7 @@ type_name_tag: id | STRING_VALUE | bind_parameter;
 
 struct_arg: type_name_tag COLON type_name_or_bind;
 struct_arg_positional:
-    type_name_tag type_name_or_bind
+    type_name_tag type_name_or_bind (NOT? NULL)?
   | type_name_or_bind AS type_name_tag; //deprecated
 variant_arg: (type_name_tag COLON)? type_name_or_bind;
 callable_arg: variant_arg (LBRACE_CURLY AUTOMAP RBRACE_CURLY)?;
diff --git a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
index 7f8604dfe6e..86d526cff81 100644
--- a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
+++ b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
@@ -697,8 +697,8 @@ Y_UNIT_TEST_SUITE(CheckSqlFormatter) {
 
     Y_UNIT_TEST(WithSchemaEquals) {
         TCases cases = {
-            {"select * from plato.T with (format= csv_with_names, schema=(year int32, month String, day String, a Utf8, b Uint16));",
-             "SELECT\n\t*\nFROM plato.T\n\tWITH (format = csv_with_names, SCHEMA = (year int32, month String, day String, a Utf8, b Uint16));\n"},
+            {"select * from plato.T with (format= csv_with_names, schema=(year int32 Null, month String, day String not   null, a Utf8, b Uint16));",
+             "SELECT\n\t*\nFROM plato.T\n\tWITH (format = csv_with_names, SCHEMA = (year int32 NULL, month String, day String NOT NULL, a Utf8, b Uint16));\n"},
              };
 
         TSetup setup;
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 134b153359b..ac3c3bfac33 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -3056,7 +3056,7 @@ bool TSqlTranslation::TableHintImpl(const TRule_table_hint& rule, TTableHints& h
             bool warn = false;
             auto processItem = [&](const TRule_struct_arg_positional& arg) {
                 // struct_arg_positional:
-                //     type_name_tag type_name_or_bind
+                //     type_name_tag type_name_or_bind (NOT? NULL)?
                 //   | type_name_or_bind AS type_name_tag; //deprecated
                 const bool altCurrent = arg.Alt_case() == TRule_struct_arg_positional::kAltStructArgPositional1;
                 auto& typeNameOrBind = altCurrent ?
@@ -3074,6 +3074,13 @@ bool TSqlTranslation::TableHintImpl(const TRule_table_hint& rule, TTableHints& h
                     warn = true;
                 }
 
+                if (altCurrent) {
+                    bool notNull = arg.GetAlt_struct_arg_positional1().HasBlock3() && arg.GetAlt_struct_arg_positional1().GetBlock3().HasBlock1();
+                    if (!notNull) {
+                        typeNode = new TCallNodeImpl(pos, "AsOptionalType", { typeNode });
+                    }
+                }
+
                 auto& typeNameTag = altCurrent ?
                     arg.GetAlt_struct_arg_positional1().GetRule_type_name_tag1() :
                     arg.GetAlt_struct_arg_positional2().GetRule_type_name_tag3();