Add support for named args in ICallableTypeBuilder (#5192)

Author: fedor-miron

ydb/library/yql/minikql/mkql_type_builder.cpp

@@ -476,7 +476,7 @@ class TCallableTypeBuilder: public NUdf::ICallableTypeBuilder
         , ReturnType_(nullptr)
         , OptionalArgs_(0)
     {
-        ArgsTypes_.reserve(argsCount);
+        Args_.reserve(argsCount);
     }
 
     NUdf::ICallableTypeBuilder& Returns(
@@ -503,21 +503,31 @@ class TCallableTypeBuilder: public NUdf::ICallableTypeBuilder
 
     NUdf::ICallableTypeBuilder& Arg(NUdf::TDataTypeId typeId) override {
         auto type = NMiniKQL::TDataType::Create(typeId, Env_);
-        ArgsTypes_.push_back(type);
+        Args_.emplace_back().Type_ = type;
         return *this;
     }
 
     NUdf::ICallableTypeBuilder& Arg(const NUdf::TType* type) override {
         auto mkqlType = const_cast<NMiniKQL::TType*>(static_cast<const NMiniKQL::TType*>(type));
-        ArgsTypes_.push_back(mkqlType);
+        Args_.emplace_back().Type_ = mkqlType;
         return *this;
     }
 
     NUdf::ICallableTypeBuilder& Arg(
             const NUdf::ITypeBuilder& typeBuilder) override
     {
         auto type = static_cast<NMiniKQL::TType*>(typeBuilder.Build());
-        ArgsTypes_.push_back(type);
+        Args_.emplace_back().Type_ = type;
+        return *this;
+    }
+
+    NUdf::ICallableTypeBuilder& Name(const NUdf::TStringRef& name) override {
+        Args_.back().Name_ = Env_.InternName(name);
+        return *this;
+    }
+
+    NUdf::ICallableTypeBuilder& Flags(ui64 flags) override {
+        Args_.back().Flags_ = flags;
         return *this;
     }
 
@@ -529,20 +539,26 @@ class TCallableTypeBuilder: public NUdf::ICallableTypeBuilder
     NUdf::TType* Build() const override {
         Y_ABORT_UNLESS(ReturnType_, "callable returns type is not configured");
 
-        NMiniKQL::TNode* payload = nullptr;
+        NMiniKQL::TCallableTypeBuilder builder(Env_, UdfName, ReturnType_);
+        for (const auto& arg : Args_) {
+            builder.Add(arg.Type_);
+            if (!arg.Name_.Str().empty()) {
+                builder.SetArgumentName(arg.Name_.Str());
+            }
+
+            if (arg.Flags_ != 0) {
+                builder.SetArgumentFlags(arg.Flags_);
+            }
+        }
+        builder.SetOptionalArgs(OptionalArgs_);
 
-        auto callableType = NMiniKQL::TCallableType::Create(
-                    UdfName, ReturnType_,
-                    ArgsTypes_.size(), const_cast<NMiniKQL::TType**>(ArgsTypes_.data()),
-                    payload, Env_);
-        callableType->SetOptionalArgumentsCount(OptionalArgs_);
-        return callableType;
+        return builder.Build();
     }
 
 private:
     const NMiniKQL::TTypeEnvironment& Env_;
     NMiniKQL::TType* ReturnType_;
-    TVector<NMiniKQL::TType*> ArgsTypes_;
+    TVector<NMiniKQL::TArgInfo> Args_;
     ui32 OptionalArgs_;
 };
 

ydb/library/yql/public/udf/udf_type_builder.h

@@ -311,6 +311,14 @@ class ICallableTypeBuilder: public ITypeBuilder
     virtual ICallableTypeBuilder& Arg(const ITypeBuilder& typeBuilder) = 0;
 
     virtual ICallableTypeBuilder& OptionalArgs(ui32 optionalArgs) = 0;
+
+#if UDF_ABI_COMPATIBILITY_VERSION_CURRENT >= UDF_ABI_COMPATIBILITY_VERSION(2, 38)
+    // sets name for the last added argument
+    virtual ICallableTypeBuilder& Name(const TStringRef& name) = 0;
+
+    // sets flags for the last added argument, see ICallablePayload::TArgumentFlags
+    virtual ICallableTypeBuilder& Flags(ui64 flags) = 0;
+#endif
 };
 
 UDF_ASSERT_TYPE_SIZE(ICallableTypeBuilder, 8);

ydb/library/yql/public/udf/udf_version.h

@@ -7,7 +7,7 @@ namespace NYql {
 namespace NUdf {
 
 #define CURRENT_UDF_ABI_VERSION_MAJOR 2
-#define CURRENT_UDF_ABI_VERSION_MINOR 37
+#define CURRENT_UDF_ABI_VERSION_MINOR 38
 #define CURRENT_UDF_ABI_VERSION_PATCH 0
 
 #ifdef USE_CURRENT_UDF_ABI_VERSION

ydb/library/yql/udfs/examples/callables/callables_udf.cpp

@@ -90,6 +90,52 @@ class TMul: public TBoxedValue
     }
 };
 
+extern const char A[] = "a";
+using TNamedA = TNamedArg<i32, A>;
+
+//////////////////////////////////////////////////////////////////////////////
+// TNamedArgUdf
+//////////////////////////////////////////////////////////////////////////////
+
+class TNamedArgUdf: public TBoxedValue {
+public:
+    static TStringRef Name() {
+        static auto name = TStringRef::Of("NamedArgUdf");
+        return name;
+    }
+
+private:
+    TUnboxedValue Run(
+            const IValueBuilder* valueBuilder,
+            const TUnboxedValuePod* args) const override
+    {
+        Y_UNUSED(valueBuilder);
+        auto res = args[0] ? args[0].Get<i32>() : 123;
+        return TUnboxedValuePod(res + 1);
+    }
+};
+
+//////////////////////////////////////////////////////////////////////////////
+// TReturnNamedArgCallable
+//////////////////////////////////////////////////////////////////////////////
+
+class TReturnNamedArgCallable: public TBoxedValue {
+public:
+    static TStringRef Name() {
+        static auto name = TStringRef::Of("ReturnNamedArgCallable");
+        return name;
+    }
+
+    TUnboxedValue Run(
+            const IValueBuilder* valueBuilder,
+            const TUnboxedValuePod* args) const override
+    {
+        Y_UNUSED(valueBuilder);
+        Y_UNUSED(args);
+        return TUnboxedValuePod(new TNamedArgUdf());
+    }
+};
+
 //////////////////////////////////////////////////////////////////////////////
 // TCallablesModule
 //////////////////////////////////////////////////////////////////////////////
@@ -154,6 +200,18 @@ class TCallablesModule: public IUdfModule
                 if (!typesOnly) {
                     builder.Implementation(new TMul);
                 }
+            } else if (TNamedArgUdf::Name() == name) {
+                builder.SimpleSignature<int(TNamedA)>();
+
+                if (!typesOnly) {
+                    builder.Implementation(new TNamedArgUdf());
+                }
+            } else if (TReturnNamedArgCallable::Name() == name) {
+                builder.Returns(builder.Callable()->Returns<int>().Arg<int>().Name(A));
+
+                if (!typesOnly) {
+                    builder.Implementation(new TReturnNamedArgCallable());
+                }
             }
         } catch (const std::exception& e) {
             builder.SetError(CurrentExceptionMessage());

ydb/library/yql/udfs/examples/callables/test/canondata/result.json

@@ -0,0 +1,7 @@
+{
+    "test.test[TestReturnNamedArgCallable]": [
+        {
+            "uri": "file://test.test_TestReturnNamedArgCallable_/results.txt"
+        }
+    ]
+}

ydb/library/yql/udfs/examples/callables/test/canondata/test.test_TestReturnNamedArgCallable_/results.txt

@@ -0,0 +1,28 @@
+[
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "DataType";
+                                    "Int32"
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        "2"
+                    ]
+                ]
+            }
+        ]
+    }
+]

ydb/library/yql/udfs/examples/callables/test/cases/TestReturnNamedArgCallable.sql

@@ -0,0 +1,3 @@
+$returnNamedArg = Callables::ReturnNamedArgCallable();
+
+select $returnNamedArg(1 AS a);

ydb/library/yql/udfs/examples/callables/test/ya.make

@@ -0,0 +1,7 @@
+YQL_UDF_YDB_TEST()
+
+DEPENDS(ydb/library/yql/udfs/examples/callables)
+
+SIZE(MEDIUM)
+
+END()

ydb/library/yql/udfs/examples/callables/ya.make

@@ -1,8 +1,12 @@
 YQL_UDF_YDB(callables_udf)
-YQL_ABI_VERSION(2 9 0)
+YQL_ABI_VERSION(2 38 0)
 
 SRCS(
     callables_udf.cpp
 )
 
 END()
+
+RECURSE_FOR_TESTS(
+    test
+)