aboutsummaryrefslogtreecommitdiffstats
path: root/yql/essentials/minikql/comp_nodes/mkql_pickle.cpp
blob: 6448efe1aeb356434d5bdacaaf319f05573d4a9c (plain) (blame)
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#include "mkql_pickle.h"
#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
#include <yql/essentials/minikql/mkql_node_cast.h>
#include <yql/essentials/minikql/computation/mkql_computation_node_pack.h>
#include <yql/essentials/minikql/computation/presort.h>
#include <yql/essentials/minikql/mkql_string_util.h>

namespace NKikimr {
namespace NMiniKQL {

namespace {

template <bool Stable>
class TPickleWrapper : public TMutableComputationNode<TPickleWrapper<Stable>> {
    typedef TMutableComputationNode<TPickleWrapper<Stable>> TBaseComputation;
public:
    TPickleWrapper(TComputationMutables& mutables, TType* type, IComputationNode* data)
        : TBaseComputation(mutables)
        , Type(type)
        , ValuePacker(mutables)
        , Data(data)
    {
    }

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
        return MakeString(ValuePacker.RefMutableObject(ctx, Stable, Type).Pack(Data->GetValue(ctx)));
    }

private:
    void RegisterDependencies() const final {
        this->DependsOn(Data);
    }

    TType* Type;
    TMutableObjectOverBoxedValue<TValuePackerBoxed> ValuePacker;
    IComputationNode *const Data;
};

class TUnpickleWrapper : public TMutableComputationNode<TUnpickleWrapper> {
    typedef TMutableComputationNode<TUnpickleWrapper> TBaseComputation;
public:
    TUnpickleWrapper(TComputationMutables& mutables, TType* type, IComputationNode* data)
        : TBaseComputation(mutables)
        , Type(type)
        , ValuePacker(mutables)
        , Data(data)
    {
    }

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
        auto data = Data->GetValue(ctx);
        auto buffer = data.AsStringRef();
        return ValuePacker.RefMutableObject(ctx, false, Type).Unpack(buffer, ctx.HolderFactory).Release();
    }

private:
    void RegisterDependencies() const final {
        DependsOn(Data);
    }

    TType* const Type;
    TMutableObjectOverBoxedValue<TValuePackerBoxed> ValuePacker;
    IComputationNode *const Data;
};


class TGenericPresortEncoderBoxed : public TComputationValue<TGenericPresortEncoderBoxed>, public TGenericPresortEncoder {
    typedef TComputationValue<TGenericPresortEncoderBoxed> TBase;
public:
    TGenericPresortEncoderBoxed(TMemoryUsageInfo* memInfo, TType* type)
        : TBase(memInfo)
        , TGenericPresortEncoder(type)
    {}
};

template <bool Desc>
class TPresortEncodeWrapper : public TMutableComputationNode<TPresortEncodeWrapper<Desc>> {
    typedef TMutableComputationNode<TPresortEncodeWrapper<Desc>> TBaseComputation;
public:
    TPresortEncodeWrapper(TComputationMutables& mutables, TType* type, IComputationNode* data)
        : TBaseComputation(mutables)
        , Type(type)
        , Encoder(mutables)
        , Data(data)
    {
    }

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
        return MakeString(Encoder.RefMutableObject(ctx, Type).Encode(Data->GetValue(ctx), Desc));
    }

private:
    void RegisterDependencies() const final {
        this->DependsOn(Data);
    }

    TType* Type;
    TMutableObjectOverBoxedValue<TGenericPresortEncoderBoxed> Encoder;
    IComputationNode *const Data;
};

}

IComputationNode* WrapPickle(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
    return new TPickleWrapper<false>(ctx.Mutables, callable.GetInput(0).GetStaticType(), LocateNode(ctx.NodeLocator, callable, 0));
}

IComputationNode* WrapStablePickle(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
    return new TPickleWrapper<true>(ctx.Mutables, callable.GetInput(0).GetStaticType(), LocateNode(ctx.NodeLocator, callable, 0));
}

IComputationNode* WrapUnpickle(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 2, "Expected 2 args");
    MKQL_ENSURE(callable.GetInput(0).IsImmediate() && callable.GetInput(0).GetNode()->GetType()->IsType(), "Expected type");
    return new TUnpickleWrapper(ctx.Mutables, static_cast<TType*>(callable.GetInput(0).GetNode()), LocateNode(ctx.NodeLocator, callable, 1));
}

IComputationNode* WrapAscending(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
    return new TPresortEncodeWrapper<false>(ctx.Mutables, callable.GetInput(0).GetStaticType(), LocateNode(ctx.NodeLocator, callable, 0));
}

IComputationNode* WrapDescending(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");
    return new TPresortEncodeWrapper<true>(ctx.Mutables, callable.GetInput(0).GetStaticType(), LocateNode(ctx.NodeLocator, callable, 0));
}

}
}