Moved yql/minikql YQL-19206

init

[nodiff:caesar]
commit_hash:d1182ef7d430ccf7e4d37ed933c7126d7bd5d6e4

1 files changed, 399 insertions, 0 deletions

diff --git a/yql/essentials/minikql/mkql_node_builder.cpp b/yql/essentials/minikql/mkql_node_builder.cpp
new file mode 100644
index 0000000000..de027eb30c
--- /dev/null
+++ b/yql/essentials/minikql/mkql_node_builder.cpp
@@ -0,0 +1,399 @@
+#include "mkql_node_builder.h"
+#include "mkql_node_printer.h"
+
+#include <util/generic/algorithm.h>
+
+namespace NKikimr {
+namespace NMiniKQL {
+
+using namespace NDetail;
+
+namespace {
+    TNode* BuildCallableTypePayload(const TVector<TStringBuf>& argNames,
+        const TVector<ui64>& argFlags, TStringBuf payload, const TTypeEnvironment& env) {
+        TStructTypeBuilder itemTypeBuilder(env);
+        itemTypeBuilder.Add("Name", TDataType::Create(NUdf::TDataType<char*>::Id, env));
+        itemTypeBuilder.Add("Flags", TDataType::Create(NUdf::TDataType<ui64>::Id, env));
+        TType* itemType = itemTypeBuilder.Build();
+
+        bool startedNames = false;
+        THashSet<TStringBuf> usedNames;
+        TListLiteralBuilder argsListBuilder(env, itemType);
+        for (ui32 i = 0; i < argNames.size(); ++i) {
+            TStructLiteralBuilder itemBuilder(env);
+            auto name = argNames[i];
+            auto flags = argFlags[i];
+            bool hasName = name.size() > 0;
+            if (startedNames) {
+                MKQL_ENSURE(hasName, "Named arguments already started");
+            } else {
+                startedNames = hasName;
+            }
+
+            if (hasName) {
+                MKQL_ENSURE(usedNames.insert(name).second, "Duplication of argument name: " << name);
+            }
+
+            itemBuilder.Add("Name", TRuntimeNode(BuildDataLiteral(name, NUdf::EDataSlot::String, env), true));
+            itemBuilder.Add("Flags", TRuntimeNode(BuildDataLiteral(NUdf::TUnboxedValuePod(flags), NUdf::EDataSlot::Uint64, env), true));
+            argsListBuilder.Add(TRuntimeNode(itemBuilder.Build(), true));
+        }
+
+        TRuntimeNode argsList = TRuntimeNode(argsListBuilder.Build(), true);
+        TStructLiteralBuilder payloadBuilder(env);
+        payloadBuilder.Add("Args", argsList);
+        payloadBuilder.Add("Payload", TRuntimeNode(BuildDataLiteral(payload, NUdf::EDataSlot::String, env), true));
+        return payloadBuilder.Build();
+    }
+}
+
+TDataLiteral* BuildDataLiteral(const NUdf::TUnboxedValuePod& value, NUdf::TDataTypeId type, const TTypeEnvironment& env) {
+    return TDataLiteral::Create(value, TDataType::Create(type, env), env);
+}
+
+TDataLiteral* BuildDataLiteral(const NUdf::TStringRef& data, NUdf::TDataTypeId type, const TTypeEnvironment& env) {
+    return BuildDataLiteral(env.NewStringValue(data), type, env);
+}
+
+TOptionalLiteral* BuildOptionalLiteral(TRuntimeNode value, const TTypeEnvironment& env) {
+    auto type = TOptionalType::Create(value.GetStaticType(), env);
+    return TOptionalLiteral::Create(value, type, env);
+}
+
+TOptionalLiteral* BuildEmptyOptionalLiteral(TType *itemType, const TTypeEnvironment& env) {
+    auto type = TOptionalType::Create(itemType, env);
+    return TOptionalLiteral::Create(type, env);
+}
+
+TOptionalLiteral* BuildEmptyOptionalDataLiteral(NUdf::TDataTypeId schemeType, const TTypeEnvironment& env) {
+    return BuildEmptyOptionalLiteral(TDataType::Create(schemeType, env), env);
+}
+
+TType* UnpackOptional(TRuntimeNode data, bool& isOptional) {
+    return UnpackOptional(data.GetStaticType(), isOptional);
+}
+
+TType* UnpackOptional(TType* type, bool& isOptional) {
+    isOptional = type->IsOptional();
+    if (!isOptional)
+        return type;
+
+    auto optType = static_cast<TOptionalType*>(type);
+    return optType->GetItemType();
+}
+
+TDataType* UnpackOptionalData(TRuntimeNode data, bool& isOptional) {
+    return UnpackOptionalData(data.GetStaticType(), isOptional);
+}
+
+TDataType* UnpackOptionalData(TType* type, bool& isOptional) {
+    auto unpackedType = UnpackOptional(type, isOptional);
+    MKQL_ENSURE(unpackedType->IsData(),
+        "Expected data or optional of data, actual: " << PrintNode(type, true));
+
+    return static_cast<TDataType*>(unpackedType);
+}
+
+TBlockType::EShape GetResultShape(const TVector<TType*>& types) {
+    TBlockType::EShape result = TBlockType::EShape::Scalar;
+    for (auto& type : types) {
+        MKQL_ENSURE(type->IsBlock(),
+                    "Expected block, actual: " << PrintNode(type, true));
+        if (static_cast<TBlockType*>(type)->GetShape() == TBlockType::EShape::Many) {
+            result = TBlockType::EShape::Many;
+        }
+    }
+    return result;
+}
+
+TTupleLiteralBuilder::TTupleLiteralBuilder(const TTypeEnvironment& env) : Env(env)
+{}
+
+void TTupleLiteralBuilder::Reserve(ui32 size) {
+    Types.reserve(size);
+    Values.reserve(size);
+}
+
+TTupleLiteralBuilder& TTupleLiteralBuilder::Add(TRuntimeNode value) {
+    Types.push_back(value.GetRuntimeType());
+    Values.push_back(value);
+    return *this;
+}
+
+TTupleLiteral* TTupleLiteralBuilder::Build() {
+    const auto& type = TTupleType::Create(Types.size(), Types.data(), Env);
+    return TTupleLiteral::Create(Values.size(), Values.data(), type, Env);
+}
+
+void TTupleLiteralBuilder::Clear() {
+    Values.clear();
+    Types.clear();
+}
+
+TStructTypeBuilder::TStructTypeBuilder(const TTypeEnvironment& env)
+    : Env(&env)
+{
+}
+
+void TStructTypeBuilder::Reserve(ui32 size) {
+    Members.reserve(size);
+}
+
+TStructTypeBuilder& TStructTypeBuilder::Add(const TStringBuf& name, TType* type, ui32* index) {
+    Members.push_back(TStructMember(Env->InternName(name).Str(), type, index));
+    return *this;
+}
+
+TStructType* TStructTypeBuilder::Build() {
+    if (Members.empty())
+        return Env->GetEmptyStructLazy()->GetType();
+
+    Sort(Members.begin(), Members.end());
+    return TStructType::Create(Members.size(), Members.data(), *Env);
+}
+
+void TStructTypeBuilder::FillIndexes() {
+    ui32 index = 0;
+    for (const TStructMember& member: Members) {
+        if (member.Index) {
+            *(member.Index) = index++;
+        }
+    }
+}
+
+void TStructTypeBuilder::Clear() {
+    Members.clear();
+}
+
+TStructLiteralBuilder::TStructLiteralBuilder(const TTypeEnvironment& env)
+    : Env(&env)
+{
+}
+
+void TStructLiteralBuilder::Reserve(ui32 size) {
+    Members.reserve(size);
+    Values.reserve(size);
+}
+
+TStructLiteralBuilder& TStructLiteralBuilder::Add(const TStringBuf& name, TRuntimeNode value) {
+    TType* valueType = value.GetStaticType();
+    Members.push_back(TStructMember(Env->InternName(name).Str(), valueType));
+    Values.push_back(value);
+    return *this;
+}
+
+TStructLiteral* TStructLiteralBuilder::Build() {
+    Y_DEBUG_ABORT_UNLESS(Members.size() == Values.size());
+    if (Members.empty())
+        return Env->GetEmptyStructLazy();
+
+    TVector<std::pair<TStringBuf, ui32>> sortedIndicies(Members.size());
+    for (ui32 i = 0, e = Members.size(); i < e; ++i) {
+        sortedIndicies[i] = std::make_pair(Members[i].Name, i);
+    }
+
+    Sort(sortedIndicies.begin(), sortedIndicies.end(),
+        [](const std::pair<TStringBuf, ui32>& x, const std::pair<TStringBuf, ui32>& y) {
+            return x.first < y.first;
+    });
+
+    TVector<TStructMember> sortedMembers(Members.size());
+    TVector<TRuntimeNode> sortedValues(Members.size());
+
+    for (ui32 i = 0, e = Members.size(); i < e; ++i) {
+        sortedMembers[i] = Members[sortedIndicies[i].second];
+        sortedValues[i] = Values[sortedIndicies[i].second];
+    }
+
+    auto type = TStructType::Create(sortedMembers.size(), sortedMembers.data(), *Env);
+    return TStructLiteral::Create(sortedValues.size(), sortedValues.data(), type, *Env);
+}
+
+void TStructLiteralBuilder::Clear() {
+    Members.clear();
+}
+
+TListLiteralBuilder::TListLiteralBuilder(const TTypeEnvironment& env, TType* type)
+    : Env(&env)
+    , Type(type)
+{}
+
+TListLiteralBuilder& TListLiteralBuilder::Add(TRuntimeNode item) {
+    Items.push_back(item);
+    return *this;
+}
+
+TListLiteral* TListLiteralBuilder::Build() {
+    auto type = TListType::Create(Type, *Env);
+    return TListLiteral::Create(Items.data(), Items.size(), type, *Env);
+}
+
+void TListLiteralBuilder::Clear() {
+    Items.clear();
+}
+
+TDictLiteralBuilder::TDictLiteralBuilder(const TTypeEnvironment& env, TType* keyType, TType* payloadType)
+    : Env(&env)
+    , KeyType(keyType)
+    , PayloadType(payloadType)
+{
+}
+
+void TDictLiteralBuilder::Reserve(ui32 size) {
+    Items.reserve(size);
+}
+
+TDictLiteralBuilder& TDictLiteralBuilder::Add(TRuntimeNode key, TRuntimeNode payload) {
+    Items.push_back(std::make_pair(key, payload));
+    return *this;
+}
+
+TDictLiteral* TDictLiteralBuilder::Build() {
+    auto type = TDictType::Create(KeyType, PayloadType, *Env);
+    return TDictLiteral::Create(Items.size(), Items.data(), type, *Env);
+}
+
+void TDictLiteralBuilder::Clear() {
+    Items.clear();
+}
+
+TCallableTypeBuilder::TCallableTypeBuilder(const TTypeEnvironment& env, const TStringBuf& name, TType* returnType)
+    : Env(&env)
+    , Name(Env->InternName(name))
+    , ReturnType(returnType)
+    , OptionalArgsCount(0)
+    , HasPayload(false)
+{}
+
+void TCallableTypeBuilder::Reserve(ui32 size) {
+    Arguments.reserve(size);
+    ArgNames.reserve(size);
+    ArgFlags.reserve(size);
+}
+
+TCallableTypeBuilder& TCallableTypeBuilder::Add(TType *type) {
+    Arguments.push_back(type);
+    ArgNames.emplace_back();
+    ArgFlags.emplace_back();
+    return *this;
+}
+
+TCallableTypeBuilder& TCallableTypeBuilder::SetArgumentName(const TStringBuf& name) {
+    MKQL_ENSURE(!ArgNames.empty(), "No arguments");
+    HasPayload = true;
+    ArgNames.back() = name;
+    return *this;
+}
+
+TCallableTypeBuilder& TCallableTypeBuilder::SetArgumentFlags(ui64 flags) {
+    MKQL_ENSURE(!ArgFlags.empty(), "No arguments");
+    HasPayload = true;
+    ArgFlags.back() = flags;
+    return *this;
+}
+
+TCallableTypeBuilder& TCallableTypeBuilder::SetOptionalArgs(ui32 count) {
+    OptionalArgsCount = count;
+    return *this;
+}
+
+TCallableTypeBuilder& TCallableTypeBuilder::SetPayload(const TStringBuf& data) {
+    HasPayload = true;
+    FuncPayload = data;
+    return *this;
+}
+
+TCallableType* TCallableTypeBuilder::Build() {
+    TNode* payload = nullptr;
+    if (HasPayload) {
+        payload = BuildCallableTypePayload(ArgNames, ArgFlags, FuncPayload, *Env);
+    }
+
+    auto ret = TCallableType::Create(ReturnType, Name.Str(), Arguments.size(), Arguments.data(), payload, *Env);
+    ret->SetOptionalArgumentsCount(OptionalArgsCount);
+    return ret;
+}
+
+void TCallableTypeBuilder::Clear() {
+    Arguments.clear();
+    ArgNames.clear();
+    ArgFlags.clear();
+    OptionalArgsCount = 0;
+    HasPayload = false;
+    FuncPayload = TStringBuf();
+}
+
+TCallableBuilder::TCallableBuilder(const TTypeEnvironment& env, const TStringBuf& name, TType* returnType, bool disableMerge)
+    : Env(&env)
+    , Name(Env->InternName(name))
+    , ReturnType(returnType)
+    , DisableMerge(disableMerge)
+    , OptionalArgsCount(0)
+    , HasPayload(false)
+{}
+
+void TCallableBuilder::Reserve(ui32 size) {
+    Arguments.reserve(size);
+    Inputs.reserve(size);
+    ArgNames.reserve(size);
+    ArgFlags.reserve(size);
+}
+
+TCallableBuilder& TCallableBuilder::Add(TRuntimeNode input) {
+    TType* inputType = input.GetStaticType();
+    Arguments.push_back(inputType);
+    Inputs.push_back(input);
+    ArgNames.emplace_back();
+    ArgFlags.emplace_back();
+    return *this;
+}
+
+TCallableBuilder& TCallableBuilder::SetOptionalArgs(ui32 count) {
+    OptionalArgsCount = count;
+    return *this;
+}
+
+TCallableBuilder& TCallableBuilder::SetTypePayload(const TStringBuf& data) {
+    HasPayload = true;
+    FuncPayload = data;
+    return *this;
+}
+
+TCallableBuilder& TCallableBuilder::SetArgumentName(const TStringBuf& name) {
+    MKQL_ENSURE(!ArgNames.empty(), "No arguments");
+    HasPayload = true;
+    ArgNames.back() = name;
+    return *this;
+}
+
+TCallableBuilder& TCallableBuilder::SetArgumentFlags(ui64 flags) {
+    MKQL_ENSURE(!ArgFlags.empty(), "No arguments");
+    HasPayload = true;
+    ArgFlags.back() = flags;
+    return *this;
+}
+
+TCallable* TCallableBuilder::Build() {
+    TNode* payload = nullptr;
+    if (HasPayload) {
+        payload = BuildCallableTypePayload(ArgNames, ArgFlags, FuncPayload, *Env);
+    }
+
+    auto type = TCallableType::Create(ReturnType, Name.Str(), Arguments.size(), Arguments.data(), payload, *Env);
+    type->SetOptionalArgumentsCount(OptionalArgsCount);
+    if (DisableMerge)
+        type->DisableMerge();
+    return TCallable::Create(Inputs.size(), Inputs.data(), type, *Env);
+}
+
+void TCallableBuilder::Clear() {
+    Arguments.clear();
+    Inputs.clear();
+    ArgNames.clear();
+    ArgFlags.clear();
+    OptionalArgsCount = 0;
+    FuncPayload = TStringBuf();
+}
+
+}
+}