1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
|
#include "mkql_lookup.h"
#include <yql/essentials/minikql/computation/mkql_computation_node_codegen.h> // Y_IGNORE
#include <yql/essentials/minikql/mkql_node_cast.h>
namespace NKikimr {
namespace NMiniKQL {
namespace {
class TLookupWrapper : public TMutableCodegeneratorPtrNode<TLookupWrapper> {
typedef TMutableCodegeneratorPtrNode<TLookupWrapper> TBaseComputation;
public:
TLookupWrapper(TComputationMutables& mutables, EValueRepresentation kind, IComputationNode* dict, IComputationNode* key)
: TBaseComputation(mutables, kind)
, Dict(dict)
, Key(key)
{
}
NUdf::TUnboxedValue DoCalculate(TComputationContext& ctx) const {
return Dict->GetValue(ctx).Lookup(Key->GetValue(ctx));
}
#ifndef MKQL_DISABLE_CODEGEN
void DoGenerateGetValue(const TCodegenContext& ctx, Value* pointer, BasicBlock*& block) const {
const auto dict = GetNodeValue(Dict, ctx, block);
GetNodeValue(pointer, Key, ctx, block);
const auto keyp = new LoadInst(Type::getInt128Ty(ctx.Codegen.GetContext()), pointer, "key", block);
CallBoxedValueVirtualMethod<NUdf::TBoxedValueAccessor::EMethod::Lookup>(pointer, dict, ctx.Codegen, block, pointer);
ValueUnRef(Key->GetRepresentation(), keyp, ctx, block);
if (Dict->IsTemporaryValue())
CleanupBoxed(dict, ctx, block);
}
#endif
private:
void RegisterDependencies() const final {
DependsOn(Dict);
DependsOn(Key);
}
IComputationNode* const Dict;
IComputationNode* const Key;
};
}
IComputationNode* WrapLookup(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
MKQL_ENSURE(callable.GetInputsCount() == 2, "Expected 2 args");
const auto dict = LocateNode(ctx.NodeLocator, callable, 0);
const auto key = LocateNode(ctx.NodeLocator, callable, 1);
return new TLookupWrapper(ctx.Mutables, GetValueRepresentation(callable.GetType()->GetReturnType()), dict, key);
}
}
}
|
