fix TypePrinter for Callable, Set type, add more tests

2 files changed, 90 insertions, 12 deletions

diff --git a/ydb/library/yql/minikql/mkql_type_builder_ut.cpp b/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
index 7b26010789e..ab270216d8e 100644
--- a/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
+++ b/ydb/library/yql/minikql/mkql_type_builder_ut.cpp
@@ -8,13 +8,12 @@ namespace NKikimr {
 
 namespace NMiniKQL {
 
-class TMiniKQLTypeBuilderTest: public TTestBase {
+class TMiniKQLTypeBuilderTest : public TTestBase {
 public:
     TMiniKQLTypeBuilderTest()
         : Env(Alloc)
         , TypeInfoHelper(new TTypeInfoHelper())
-        , FunctionTypeInfoBuilder(Env, TypeInfoHelper, "", nullptr, {})
-    {
+        , FunctionTypeInfoBuilder(Env, TypeInfoHelper, "", nullptr, {}) {
     }
 
 private:
@@ -27,12 +26,15 @@ private:
         UNIT_TEST(TestPgTypeFormat);
         UNIT_TEST(TestSingularTypeFormat);
         UNIT_TEST(TestOptionalTypeFormat);
-        UNIT_TEST(TestDataTypeFormat);
+        UNIT_TEST(TestEnumTypeFormat);
         UNIT_TEST(TestResourceTypeFormat);
         UNIT_TEST(TestTaggedTypeFormat);
         UNIT_TEST(TestContainerTypeFormat);
         UNIT_TEST(TestStructTypeFormat);
         UNIT_TEST(TestTupleTypeFormat);
+        UNIT_TEST(TestVariantTypeFormat);
+        UNIT_TEST(TestCallableTypeFormat);
+        UNIT_TEST(TestDataTypeFormat);
     UNIT_TEST_SUITE_END();
 
     TString FormatType(NUdf::TType* t) {
@@ -101,10 +103,27 @@ private:
         }
         {
             auto s = FormatType(FunctionTypeInfoBuilder.Set()->Key<char*>().Build());
-            UNIT_ASSERT_VALUES_EQUAL(s, "Dict<String,Void>");
+            UNIT_ASSERT_VALUES_EQUAL(s, "Set<String>");
+        }
+        {
+            auto s = FormatType(FunctionTypeInfoBuilder.Stream()->Item<ui64>().Build());
+            UNIT_ASSERT_VALUES_EQUAL(s, "Stream<Uint64>");
         }
     }
 
+    void TestEnumTypeFormat() {
+        ui32 index = 0;
+        auto t = FunctionTypeInfoBuilder.Enum(2)
+            ->AddField("RED", &index)
+            .AddField("GREEN", &index)
+            .AddField("BLUE", &index)
+            .Build();
+
+        auto s = FormatType(t);
+        // todo: fix format: Enum<'BLUE','GREEN','RED'>
+        UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Struct<'BLUE':Void,'GREEN':Void,'RED':Void>>");
+    }
+
     void TestResourceTypeFormat() {
         {
             auto s = FormatType(FunctionTypeInfoBuilder.Resource("my_resource"));
@@ -153,6 +172,52 @@ private:
         }
     }
 
+    void TestVariantTypeFormat() {
+        {
+            ui32 index = 0;
+            auto t = FunctionTypeInfoBuilder.Struct(2)
+                ->AddField<bool>("is_ok", &index)
+                .AddField<char*>("name", &index)
+                .Build();
+
+            auto s = FormatType(FunctionTypeInfoBuilder.Variant()->Over(t).Build());
+            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Struct<'is_ok':Bool,'name':String>>");
+        }
+        {
+            auto t = FunctionTypeInfoBuilder.Tuple(2)
+                ->Add<bool>()
+                .Add<char*>()
+                .Build();
+
+            auto s = FormatType(FunctionTypeInfoBuilder.Variant()->Over(t).Build());
+            UNIT_ASSERT_VALUES_EQUAL(s, "Variant<Tuple<Bool,String>>");
+        }
+    }
+
+    void TestCallableTypeFormat() {
+        {
+            auto t = FunctionTypeInfoBuilder.Callable(0)->Returns<char*>().Build();
+            auto s = FormatType(t);
+            UNIT_ASSERT_VALUES_EQUAL(s, "Callable<()->String>");
+        }
+        {
+            auto t = FunctionTypeInfoBuilder.Callable(1)->Arg<i64>().Returns<char*>().Build();
+            auto s = FormatType(t);
+            UNIT_ASSERT_VALUES_EQUAL(s, "Callable<(Int64)->String>");
+        }
+        {
+            auto t = FunctionTypeInfoBuilder.Callable(2)->Arg<i64>().Arg<char*>().Returns<bool>().Build();
+            auto s = FormatType(t);
+            UNIT_ASSERT_VALUES_EQUAL(s, "Callable<(Int64,String)->Bool>");
+        }
+        {
+            auto arg3 = FunctionTypeInfoBuilder.Optional()->Item<char*>().Build();
+            auto t = FunctionTypeInfoBuilder.Callable(3)->Arg<i64>().Arg<NUdf::TUtf8>().Arg(arg3).OptionalArgs(1).Returns<bool>().Build();
+            auto s = FormatType(t);
+            UNIT_ASSERT_VALUES_EQUAL(s, "Callable<(Int64,Utf8,[String?])->Bool>");
+        }
+    }
+
     void TestDataTypeFormat() {
         {
             auto s = FormatType(FunctionTypeInfoBuilder.SimpleType<i8>());
diff --git a/ydb/library/yql/public/udf/udf_type_printer.cpp b/ydb/library/yql/public/udf/udf_type_printer.cpp
index 7e6fb520492..fa984fa538e 100644
--- a/ydb/library/yql/public/udf/udf_type_printer.cpp
+++ b/ydb/library/yql/public/udf/udf_type_printer.cpp
@@ -67,10 +67,17 @@ void TTypePrinter1::OnTuple(ui32 elementsCount, const TType** elementsTypes) {
 }
 
 void TTypePrinter1::OnDict(const TType* keyType, const TType* valueType) {
-    *Output_ << "Dict<";
+    const bool isSet = TypeHelper1_.GetTypeKind(valueType) == ETypeKind::Void;
+    if (isSet) {
+        *Output_ << "Set<";
+    } else {
+        *Output_ << "Dict<";
+    }
     OutImpl(keyType);
-    *Output_ << ',';
-    OutImpl(valueType);
+    if (!isSet) {
+        *Output_ << ',';
+        OutImpl(valueType);
+    }
     *Output_ << '>';
 }
 
@@ -79,16 +86,22 @@ void TTypePrinter1::OnCallable(const TType* returnType, ui32 argsCount, const TT
     for (ui32 i = 0U; i < argsCount; ++i) {
         if (optionalArgsCount && i == argsCount -  optionalArgsCount)
             *Output_ << '[';
-        if (const std::string_view name = payload->GetArgumentName(i); !name.empty())
-            *Output_ << "'" << name << "':";
+        if (payload) {
+            const std::string_view name = payload->GetArgumentName(i);
+            if (!name.empty())
+                *Output_ << "'" << name << "':";
+        }
         OutImpl(argsTypes[i]);
-        if (ICallablePayload::TArgumentFlags::AutoMap == payload->GetArgumentFlags(i))
-            *Output_ << "{Flags:AutoMap}";
+        if (payload) {
+            if (ICallablePayload::TArgumentFlags::AutoMap == payload->GetArgumentFlags(i))
+                *Output_ << "{Flags:AutoMap}";
+        }
         if (i < argsCount - 1U)
             *Output_ << ',';
     }
     *Output_ << (optionalArgsCount ? "])->" : ")->");
     OutImpl(returnType);
+    *Output_ << ">";
 }
 
 void TTypePrinter1::OnVariant(const TType* underlyingType) {