Move yql/udfs/test/ to /yql/essentials YQL-19206

commit_hash:cdc59617409c3a4fbd6d01670cae7cb39e979340

6 files changed, 321 insertions, 0 deletions

diff --git a/yql/essentials/udfs/test/simple/simple_udf.cpp b/yql/essentials/udfs/test/simple/simple_udf.cpp
new file mode 100644
index 00000000000..067bdcc28dd
--- /dev/null
+++ b/yql/essentials/udfs/test/simple/simple_udf.cpp
@@ -0,0 +1,267 @@
+#include <util/generic/bt_exception.h>
+#include <util/generic/hash.h>
+#include <util/string/cast.h>
+#include <yql/essentials/public/udf/udf_helpers.h>
+#include <yql/essentials/public/udf/udf_value_builder.h>
+
+using namespace NKikimr;
+using namespace NUdf;
+
+namespace {
+
+SIMPLE_UDF(TCrash, ui64(char*)) {
+    Y_UNUSED(valueBuilder);
+    Y_UNUSED(args);
+    int *ptr = nullptr;
+    *ptr = 1;
+    return TUnboxedValuePod(0);
+}
+
+void Throws(const TString& msg) {
+    ythrow TWithBackTrace<yexception>() << msg;
+}
+
+SIMPLE_UDF(TException, ui64(char*)) {
+    Y_UNUSED(valueBuilder);
+    TString msg(args[0].AsStringRef());
+    Throws(msg);
+    return TUnboxedValuePod(0);
+}
+
+SIMPLE_UDF(TReturnNull, char*(char*)) {
+    Y_UNUSED(valueBuilder);
+    Y_UNUSED(args);
+    return TUnboxedValuePod(NULL);
+}
+
+SIMPLE_UDF(TReturnVoid, void(char*)) {
+    Y_UNUSED(valueBuilder);
+    Y_UNUSED(args);
+    return TUnboxedValuePod::Void();
+}
+
+SIMPLE_UDF(TReturnEmpty, TOptional<char*>(char*)) {
+    Y_UNUSED(valueBuilder);
+    Y_UNUSED(args);
+    return TUnboxedValuePod();
+}
+
+SIMPLE_UDF(TReturnBrokenInt, ui32()) {
+    Y_UNUSED(args);
+    return valueBuilder->NewString("Bunny");
+}
+
+SIMPLE_UDF(TEcho, char*(TOptional<char*>)) {
+    if (args[0]) {
+        return valueBuilder->NewString(args[0].AsStringRef());
+    } else {
+        return valueBuilder->NewString("<empty optional>");
+    }
+}
+
+SIMPLE_UDF_WITH_OPTIONAL_ARGS(TEchoWithPrefix, char*(char*,TOptional<char*>), 1) {
+    if (!args[1])
+        return TUnboxedValuePod(args[0]);
+    return valueBuilder->ConcatStrings(args[1], args[0]);
+}
+
+SIMPLE_UDF_RUN(TEchoWithRunPrefix, char*(char*), TOptional<char*>) {
+    if (!RunConfig)
+        return TUnboxedValuePod(args[0]);
+
+    return valueBuilder->PrependString(RunConfig.AsStringRef(), args[0]);
+}
+
+SIMPLE_UDF(TConst, char*()) {
+    Y_UNUSED(args);
+    return valueBuilder->NewString(TStringBuf("Constant response"));
+}
+
+SIMPLE_UDF(TConcat, char*(char*, char*)) {
+    return valueBuilder->ConcatStrings(args[0], args[1]);
+}
+
+SIMPLE_UDF(TRepeat, char*(char*, ui64)) {
+    TString orig(args[0].AsStringRef());
+    ui64 times = args[1].Get<ui64>();
+    TString res = "";
+    for (ui64 i = 0; i < times; i++) {
+        res += orig;
+    }
+    return valueBuilder->NewString(res);
+}
+
+SIMPLE_UDF(TSleep, ui64(ui64)) {
+    Y_UNUSED(valueBuilder);
+    ui64 time = args[0].Get<ui64>();
+    usleep(time);
+    return TUnboxedValuePod(static_cast<ui64>(0));
+}
+
+using TComplexReturnTypeSignature = TDict<char*, ui32>(char*);
+SIMPLE_UDF(TComplexReturnType, TComplexReturnTypeSignature) {
+    const TStringBuf s = args[0].AsStringRef();
+    THashMap<TString, ui32> stat;
+    for(auto c: s) {
+       ++stat[TString{c}];
+    }
+    auto dictBuilder = valueBuilder->NewDict(ReturnType_, 0);
+    for(const auto& [k, v]: stat) {
+        dictBuilder->Add(valueBuilder->NewString(k), TUnboxedValuePod{v});
+    }
+    return dictBuilder->Build();
+}
+
+extern const char c[] = "C";
+extern const char d[] = "D";
+
+using TNamedC = TNamedArg<ui32, c>;
+using TNamedD = TNamedArg<ui32, d>;
+
+SIMPLE_UDF_WITH_OPTIONAL_ARGS(TNamedArgs, char*(ui32, TOptional<ui32>, TNamedC, TNamedD), 3) {
+    TString res;
+    res += "A=" + ToString(args[0].Get<ui32>());
+    res += " B=" + (args[1] ? ToString(args[1].Get<ui32>()) : "none");
+    res += " C=" + (args[2] ? ToString(args[2].Get<ui32>()) : "none");
+    res += " D=" + (args[3] ? ToString(args[3].Get<ui32>()) : "none");
+    return valueBuilder->NewString(res);
+}
+
+UDF(TIncrement, builder.Args(2)->
+    Add<ui32>().Name("Arg1").Flags(ICallablePayload::TArgumentFlags::AutoMap)
+    .Add(builder.SimpleType<TOptional<ui32>>()).Name("Arg2")
+    .Done().Returns<ui32>().OptionalArgs(1);) {
+    Y_UNUSED(valueBuilder);
+    return TUnboxedValuePod(args[0].Get<ui32>() + args[1].GetOrDefault<ui32>(1));
+}
+
+UDF(TIncrementOpt, builder.Args(2)->
+    Add<ui32>().Name("Arg1").Flags(ICallablePayload::TArgumentFlags::AutoMap)
+    .Add(builder.SimpleType<TOptional<ui32>>()).Name("Arg2")
+    .Done().Returns(builder.SimpleType<TOptional<ui32>>()).OptionalArgs(1);) {
+    Y_UNUSED(valueBuilder);
+    if (const ui32 by = args[1].GetOrDefault<ui32>(0)) {
+        return TUnboxedValuePod(args[0].Get<ui32>() + by);
+    }
+    else {
+        return TUnboxedValuePod();
+    }
+}
+
+UDF_IMPL(TIncrementWithCounters,
+    builder.Args(1)->Add<ui32>().Done().Returns<ui32>();
+    ,
+    mutable ::NKikimr::NUdf::TCounter Counter_;
+    mutable ::NKikimr::NUdf::TScopedProbe Scope_;
+    ,
+    Counter_ = builder.GetCounter("IncrementWithCounters_Calls", true);
+    Scope_ = builder.GetScopedProbe("IncrementWithCounters_Time");
+    ,
+    ""
+    ,
+    ""
+    ,
+    void
+) {
+    Y_UNUSED(valueBuilder);
+    Counter_.Inc();
+    with_lock(Scope_) {
+        return TUnboxedValuePod(args[0].Get<ui32>() + 1);
+    }
+}
+
+class TGenericAsStruct : public TBoxedValue {
+public:
+    typedef bool TTypeAwareMarker;
+
+    TUnboxedValue Run(const IValueBuilder* valueBuilder, const TUnboxedValuePod* args) const final {
+        TUnboxedValue* items = nullptr;
+        auto result = valueBuilder->NewArray(Argc, items);
+        for (size_t i = 0; i < Argc; ++i) {
+            items[i] = std::move(args[i]);
+        }
+        return result;
+    }
+
+    static const TStringRef& Name() {
+        static auto name = TStringRef::Of("GenericAsStruct");
+        return name;
+    }
+
+    TGenericAsStruct(size_t argc)
+        : Argc(argc)
+    {}
+
+    static bool DeclareSignature(const TStringRef& name, TType* userType, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
+        if (Name() == name) {
+            if (!userType) {
+                builder.SetError("Missing user type.");
+                return true;
+            }
+
+            builder.UserType(userType);
+            const auto typeHelper = builder.TypeInfoHelper();
+            const auto userTypeInspector = TTupleTypeInspector(*typeHelper, userType);
+            if (!userTypeInspector || userTypeInspector.GetElementsCount() < 1) {
+                builder.SetError("Invalid user type.");
+                return true;
+            }
+
+            const auto argsTypeTuple = userTypeInspector.GetElementType(0);
+            const auto argsTypeInspector = TTupleTypeInspector(*typeHelper, argsTypeTuple);
+            if (!argsTypeInspector) {
+                builder.SetError("Invalid user type - expected tuple.");
+                return true;
+            }
+
+            const size_t argsCount = argsTypeInspector.GetElementsCount();
+            auto argBuilder = builder.Args(argsCount);
+            auto structBuilder = builder.Struct(argsCount);
+            for (size_t i = 0; i < argsCount; ++i) {
+                auto argType = argsTypeInspector.GetElementType(i);
+                argBuilder->Add(argType);
+                TString name = TStringBuilder() << "arg_" << i;
+                structBuilder->AddField(name, argType, nullptr);
+            }
+
+            argBuilder->Done().Returns(builder.Optional()->Item(structBuilder->Build()).Build());
+
+            if (!typesOnly) {
+                builder.Implementation(new TGenericAsStruct(argsCount));
+            }
+            return true;
+        }
+        else {
+            return false;
+        }
+    }
+private:
+    const size_t Argc;
+};
+
+
+SIMPLE_MODULE(TSimpleUdfModule,
+                TCrash,
+                TException,
+                TReturnNull,
+                TReturnVoid,
+                TReturnEmpty,
+                TReturnBrokenInt,
+                TEcho,
+                TEchoWithPrefix,
+                TEchoWithRunPrefix,
+                TConst,
+                TConcat,
+                TRepeat,
+                TSleep,
+                TComplexReturnType,
+                TNamedArgs,
+                TIncrement,
+                TIncrementOpt,
+                TIncrementWithCounters,
+                TGenericAsStruct
+              )
+
+} // namespace
+
+REGISTER_MODULES(TSimpleUdfModule)
diff --git a/yql/essentials/udfs/test/simple/ya.make b/yql/essentials/udfs/test/simple/ya.make
new file mode 100644
index 00000000000..4a5f4c5b0fe
--- /dev/null
+++ b/yql/essentials/udfs/test/simple/ya.make
@@ -0,0 +1,8 @@
+YQL_UDF_CONTRIB(simple_udf)
+YQL_ABI_VERSION(2 9 0)
+
+SRCS(
+    simple_udf.cpp
+)
+
+END()
diff --git a/yql/essentials/udfs/test/test_import/test_import_udf.cpp b/yql/essentials/udfs/test/test_import/test_import_udf.cpp
new file mode 100644
index 00000000000..158194b792e
--- /dev/null
+++ b/yql/essentials/udfs/test/test_import/test_import_udf.cpp
@@ -0,0 +1,33 @@
+#include <yql/essentials/public/udf/udf_helpers.h>
+#include <yql/essentials/public/udf/udf_value_builder.h>
+
+#include <util/string/cast.h>
+#include <util/generic/string.h>
+
+using namespace NKikimr;
+using namespace NUdf;
+
+namespace {
+
+SIMPLE_UDF(TConcat, char*(char*, char*)) {
+    return valueBuilder->ConcatStrings(args[0], args[1]);
+}
+
+SIMPLE_UDF(TRepeat, char*(char*, ui64)) {
+    TString orig(args[0].AsStringRef());
+    ui64 times = args[1].Get<ui64>();
+    TString res = "";
+    for (ui64 i = 0; i < times; i++) {
+        res += orig;
+    }
+    return valueBuilder->NewString(res);
+}
+
+SIMPLE_MODULE(TTestImportUdfModule,
+                TConcat,
+                TRepeat
+              )
+
+} // namespace
+
+REGISTER_MODULES(TTestImportUdfModule)
diff --git a/yql/essentials/udfs/test/test_import/ya.make b/yql/essentials/udfs/test/test_import/ya.make
new file mode 100644
index 00000000000..1adb06149f1
--- /dev/null
+++ b/yql/essentials/udfs/test/test_import/ya.make
@@ -0,0 +1,8 @@
+YQL_UDF_CONTRIB(test_import_udf)
+YQL_ABI_VERSION(2 9 0)
+
+SRCS(
+    test_import_udf.cpp
+)
+
+END()
diff --git a/yql/essentials/udfs/test/ya.make b/yql/essentials/udfs/test/ya.make
new file mode 100644
index 00000000000..5a6daf544bc
--- /dev/null
+++ b/yql/essentials/udfs/test/ya.make
@@ -0,0 +1,4 @@
+RECURSE(
+    simple
+    test_import
+)
diff --git a/yql/essentials/udfs/ya.make b/yql/essentials/udfs/ya.make
index 4b8f9607379..b8d9b2d2132 100644
--- a/yql/essentials/udfs/ya.make
+++ b/yql/essentials/udfs/ya.make
@@ -1,4 +1,5 @@
 RECURSE(
     examples
     logs
+    test
 )