Moved yql/minikql YQL-19206

init

[nodiff:caesar]
commit_hash:d1182ef7d430ccf7e4d37ed933c7126d7bd5d6e4

1 files changed, 124 insertions, 0 deletions

diff --git a/yql/essentials/minikql/computation/mkql_computation_node_dict_ut.cpp b/yql/essentials/minikql/computation/mkql_computation_node_dict_ut.cpp
new file mode 100644
index 00000000000..ae2e64c77cd
--- /dev/null
+++ b/yql/essentials/minikql/computation/mkql_computation_node_dict_ut.cpp
@@ -0,0 +1,124 @@
+#include <yql/essentials/minikql/mkql_node.h>
+#include <yql/essentials/minikql/mkql_node_cast.h>
+#include <yql/essentials/minikql/mkql_program_builder.h>
+#include <yql/essentials/minikql/mkql_function_registry.h>
+#include <yql/essentials/minikql/computation/mkql_computation_node.h>
+#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
+#include <yql/essentials/minikql/invoke_builtins/mkql_builtins.h>
+#include <yql/essentials/minikql/comp_nodes/mkql_factories.h>
+
+#include <library/cpp/testing/unittest/registar.h>
+
+#include <vector>
+#include <utility>
+#include <algorithm>
+
+namespace NKikimr {
+namespace NMiniKQL {
+
+namespace {
+struct TSetup {
+    TSetup(TScopedAlloc& alloc)
+        : Alloc(alloc)
+    {
+        FunctionRegistry = CreateFunctionRegistry(CreateBuiltinRegistry());
+        RandomProvider = CreateDeterministicRandomProvider(1);
+        TimeProvider = CreateDeterministicTimeProvider(10000000);
+
+        Env.Reset(new TTypeEnvironment(Alloc));
+        PgmBuilder.Reset(new TProgramBuilder(*Env, *FunctionRegistry));
+    }
+
+    THolder<IComputationGraph> BuildGraph(TRuntimeNode pgm, const std::vector<TNode*>& entryPoints = std::vector<TNode*>()) {
+        Explorer.Walk(pgm.GetNode(), *Env);
+        TComputationPatternOpts opts(Alloc.Ref(), *Env, GetBuiltinFactory(),
+            FunctionRegistry.Get(),
+            NUdf::EValidateMode::None, NUdf::EValidatePolicy::Exception, "OFF", EGraphPerProcess::Multi);
+        Pattern = MakeComputationPattern(Explorer, pgm, entryPoints, opts);
+        TComputationOptsFull compOpts = opts.ToComputationOptions(*RandomProvider, *TimeProvider);
+        return Pattern->Clone(compOpts);
+    }
+
+    TIntrusivePtr<IFunctionRegistry> FunctionRegistry;
+    TIntrusivePtr<IRandomProvider> RandomProvider;
+    TIntrusivePtr<ITimeProvider> TimeProvider;
+
+    TScopedAlloc& Alloc;
+    THolder<TTypeEnvironment> Env;
+    THolder<TProgramBuilder> PgmBuilder;
+
+    TExploringNodeVisitor Explorer;
+    IComputationPattern::TPtr Pattern;
+};
+}
+
+Y_UNIT_TEST_SUITE(TestCompactMultiDict) {
+    Y_UNIT_TEST(TestIterate) {
+        TScopedAlloc alloc(__LOCATION__);
+
+        TSetup setup(alloc);
+
+        const std::vector<std::pair<ui32, std::vector<ui32>>> items = {{1, {1, 2}}, {2, {1}}, {3, {0}}, {6, {1, 7}}};
+
+        TProgramBuilder& pgmBuilder = *setup.PgmBuilder;
+        TVector<TRuntimeNode> rItems;
+        for (auto& [k, vv]: items) {
+            for (auto& v: vv) {
+                rItems.push_back(pgmBuilder.NewTuple({pgmBuilder.NewDataLiteral<ui32>(k), pgmBuilder.NewDataLiteral<ui32>(v)}));
+            }
+        }
+        auto ui32Type = pgmBuilder.NewDataType(NUdf::TDataType<ui32>::Id);
+        auto list = pgmBuilder.NewList(pgmBuilder.NewTupleType({ui32Type, ui32Type}), rItems);
+
+        auto dict = pgmBuilder.ToHashedDict(list, /*all*/true,
+            [&pgmBuilder](TRuntimeNode item) { return pgmBuilder.Nth(item, 0); },
+            [&pgmBuilder](TRuntimeNode item) { return pgmBuilder.Nth(item, 1); },
+            /*isCompact*/true,
+            items.size());
+
+        auto graph = setup.BuildGraph(dict, {});
+        NUdf::TUnboxedValue res = graph->GetValue();
+
+        std::vector<ui32> keyVals;
+        for (NUdf::TUnboxedValue keys = res.GetKeysIterator(), v; keys.Next(v);) {
+            keyVals.push_back(v.Get<ui32>());
+        }
+        UNIT_ASSERT_VALUES_EQUAL(keyVals.size(), items.size());
+        std::sort(keyVals.begin(), keyVals.end());
+        UNIT_ASSERT(
+            std::equal(keyVals.begin(), keyVals.end(), items.begin(),
+                [](ui32 l, const std::pair<ui32, std::vector<ui32>>& r) { return l == r.first; }
+            )
+        );
+
+        std::vector<std::vector<ui32>> origPayloads;
+        for (auto& [k, vv]: items) {
+            origPayloads.push_back(vv);
+            std::sort(origPayloads.back().begin(), origPayloads.back().end());
+        }
+        std::sort(origPayloads.begin(), origPayloads.end());
+
+        std::vector<std::vector<ui32>> payloadVals;
+        for (NUdf::TUnboxedValue payloads = res.GetPayloadsIterator(), v; payloads.Next(v);) {
+            payloadVals.emplace_back();
+            for (NUdf::TUnboxedValue i = v.GetListIterator(), p; i.Next(p);) {
+                payloadVals.back().push_back(p.Get<ui32>());
+            }
+            std::sort(payloadVals.back().begin(), payloadVals.back().end());
+        }
+        std::sort(payloadVals.begin(), payloadVals.end());
+        UNIT_ASSERT_VALUES_EQUAL(origPayloads, payloadVals);
+
+        std::vector<std::pair<ui32, std::vector<ui32>>> vals;
+        for (NUdf::TUnboxedValue values = res.GetDictIterator(), k, payloads; values.NextPair(k, payloads);) {
+            vals.emplace_back(k.Get<ui32>(), std::vector<ui32>{});
+            for (NUdf::TUnboxedValue i = payloads.GetListIterator(), p; i.Next(p);) {
+                vals.back().second.push_back(p.Get<ui32>());
+            }
+            std::sort(vals.back().second.begin(), vals.back().second.end());
+        }
+        UNIT_ASSERT_VALUES_EQUAL(items, vals);
+    }
+}
+}
+}