fix tests

2 files changed, 29 insertions, 71 deletions

diff --git a/ydb/library/yql/minikql/arrow/mkql_functions_ut.cpp b/ydb/library/yql/minikql/arrow/mkql_functions_ut.cpp
index 300abdb5fd..caab7ca81c 100644
--- a/ydb/library/yql/minikql/arrow/mkql_functions_ut.cpp
+++ b/ydb/library/yql/minikql/arrow/mkql_functions_ut.cpp
@@ -1,6 +1,7 @@
 #include <library/cpp/testing/unittest/registar.h>
 
 #include "mkql_functions.h"
+#include <ydb/library/yql/minikql/invoke_builtins/mkql_builtins.h>
 
 namespace NKikimr::NMiniKQL {
 
@@ -8,6 +9,7 @@ Y_UNIT_TEST_SUITE(TMiniKQLArrowFunctions) {
     Y_UNIT_TEST(Add) {
         TScopedAlloc alloc(__LOCATION__);
         TTypeEnvironment env(alloc);
+        auto builtins = CreateBuiltinRegistry();
 
         auto uint64Type = TDataType::Create(NUdf::GetDataTypeInfo(NUdf::EDataSlot::Uint64).TypeId, env);
         auto uint64TypeOpt = TOptionalType::Create(uint64Type, env);
@@ -18,77 +20,32 @@ Y_UNIT_TEST_SUITE(TMiniKQLArrowFunctions) {
         auto scalarTypeOpt = TBlockType::Create(uint64TypeOpt, TBlockType::EShape::Scalar, env);
         auto arrayTypeOpt = TBlockType::Create(uint64TypeOpt, TBlockType::EShape::Many, env);
 
-        TType* outputType;
-        UNIT_ASSERT(!FindArrowFunction("_add_", {}, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("add", {}, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("add", TVector<TType*>{ scalarType }, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("add", TVector<TType*>{ arrayType }, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("add", TVector<TType*>{ scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("add", TVector<TType*>{ arrayTypeOpt }, outputType, env));
-
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayType, arrayType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayType));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarType, arrayType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayType));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayType, scalarType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayType));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarType, scalarType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarType));
-
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayType, arrayTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarType, arrayTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayType, scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarType, scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarTypeOpt));
-
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayTypeOpt, arrayType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarTypeOpt, arrayType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayTypeOpt, scalarType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarTypeOpt, scalarType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarTypeOpt));
-
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayTypeOpt, arrayTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarTypeOpt, arrayTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ arrayTypeOpt, scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayTypeOpt));
-        UNIT_ASSERT(FindArrowFunction("add", TVector<TType*>{ scalarTypeOpt, scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarTypeOpt));
-    }
-
-    Y_UNIT_TEST(IsNull) {
-        TScopedAlloc alloc(__LOCATION__);
-        TTypeEnvironment env(alloc);
-
-        auto bool64Type = TDataType::Create(NUdf::GetDataTypeInfo(NUdf::EDataSlot::Bool).TypeId, env);
-        auto bool64TypeOpt = TOptionalType::Create(bool64Type, env);
-
-        auto scalarType = TBlockType::Create(bool64Type, TBlockType::EShape::Scalar, env);
-        auto arrayType = TBlockType::Create(bool64Type, TBlockType::EShape::Many, env);
-
-        auto scalarTypeOpt = TBlockType::Create(bool64TypeOpt, TBlockType::EShape::Scalar, env);
-        auto arrayTypeOpt = TBlockType::Create(bool64TypeOpt, TBlockType::EShape::Many, env);
-
-        TType* outputType;
-        UNIT_ASSERT(!FindArrowFunction("is_null", {}, outputType, env));
-        UNIT_ASSERT(!FindArrowFunction("is_null", TVector<TType*>{ scalarType, scalarType }, outputType, env));
-
-        UNIT_ASSERT(FindArrowFunction("is_null", TVector<TType*>{ scalarType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarType));
-        UNIT_ASSERT(FindArrowFunction("is_null", TVector<TType*>{ arrayType }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayType));
-
-        UNIT_ASSERT(FindArrowFunction("is_null", TVector<TType*>{ scalarTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*scalarType));
-        UNIT_ASSERT(FindArrowFunction("is_null", TVector<TType*>{ arrayTypeOpt }, outputType, env));
-        UNIT_ASSERT(outputType->Equals(*arrayType));
+        UNIT_ASSERT(!FindArrowFunction("_Add_", {}, scalarType, *builtins));
+        UNIT_ASSERT(!FindArrowFunction("Add", {}, scalarType, *builtins));
+        UNIT_ASSERT(!FindArrowFunction("Add", TVector<TType*>{ scalarType }, scalarType, *builtins));
+        UNIT_ASSERT(!FindArrowFunction("Add", TVector<TType*>{ arrayType }, arrayType, *builtins));
+        UNIT_ASSERT(!FindArrowFunction("Add", TVector<TType*>{ scalarTypeOpt }, scalarType, *builtins));
+        UNIT_ASSERT(!FindArrowFunction("Add", TVector<TType*>{ arrayTypeOpt }, arrayType, *builtins));
+
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayType, arrayType }, arrayType, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarType, arrayType }, arrayType, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayType, scalarType }, arrayType, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarType, scalarType }, scalarType, *builtins));
+
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayType, arrayTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarType, arrayTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayType, scalarTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarType, scalarTypeOpt }, scalarTypeOpt, *builtins));
+
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayTypeOpt, arrayType }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarTypeOpt, arrayType }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayTypeOpt, scalarType }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarTypeOpt, scalarType }, scalarTypeOpt, *builtins));
+
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayTypeOpt, arrayTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarTypeOpt, arrayTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ arrayTypeOpt, scalarTypeOpt }, arrayTypeOpt, *builtins));
+        UNIT_ASSERT(FindArrowFunction("Add", TVector<TType*>{ scalarTypeOpt, scalarTypeOpt }, scalarTypeOpt, *builtins));
     }
 }
 
diff --git a/ydb/library/yql/minikql/arrow/ut/ya.make b/ydb/library/yql/minikql/arrow/ut/ya.make
index 8135fe3815..c79440694b 100644
--- a/ydb/library/yql/minikql/arrow/ut/ya.make
+++ b/ydb/library/yql/minikql/arrow/ut/ya.make
@@ -11,6 +11,7 @@ PEERDIR(
     ydb/library/yql/public/udf
     ydb/library/yql/public/udf/service/exception_policy
     ydb/library/yql/sql/pg_dummy
+    ydb/library/yql/minikql/invoke_builtins/llvm
 )
 
 YQL_LAST_ABI_VERSION()