aboutsummaryrefslogtreecommitdiffstats
path: root/yql/essentials/minikql/comp_nodes/mkql_tobytes.cpp
blob: 825e8386a4d6390d808e3a96461ec8fdb0defc4c (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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
#include "mkql_tobytes.h"
#include <yql/essentials/minikql/computation/mkql_computation_node_codegen.h>  // Y_IGNORE
#include <yql/essentials/minikql/mkql_node_cast.h>
#include <yql/essentials/minikql/mkql_node_builder.h>
#include <yql/essentials/minikql/mkql_string_util.h>
#include <yql/essentials/utils/swap_bytes.h>

namespace NKikimr {
namespace NMiniKQL {

namespace {

template<bool IsOptional, typename Type>
class TToBytesPrimitiveTypeWrapper : public TDecoratorCodegeneratorNode<TToBytesPrimitiveTypeWrapper<IsOptional, Type>> {
using TBaseComputation = TDecoratorCodegeneratorNode<TToBytesPrimitiveTypeWrapper<IsOptional, Type>>;
public:
    TToBytesPrimitiveTypeWrapper(IComputationNode* data)
        : TBaseComputation(data)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const {
        if (IsOptional && !value)
            return NUdf::TUnboxedValuePod();

        const auto& v = value.Get<Type>();
        return NUdf::TUnboxedValuePod::Embedded(NUdf::TStringRef(reinterpret_cast<const char*>(&v), sizeof(v)));
    }
#ifndef MKQL_DISABLE_CODEGEN
    Value* DoGenerateGetValue(const TCodegenContext& ctx, Value* value, BasicBlock*& block) const {
        const uint64_t one[] = {0ULL, sizeof(Type) << 48ULL};
        const auto size = ConstantInt::get(value->getType(), APInt(128, 2, one));
        const uint64_t two[] = {0xFFFFFFFFFFFFFFFFULL, 0xFF00FFFFFFFFFFFFULL};
        const auto mask = ConstantInt::get(value->getType(), APInt(128, 2, two));
        const auto result = BinaryOperator::CreateOr(BinaryOperator::CreateAnd(value, mask, "and", block), size, "or", block);
        if constexpr (IsOptional)
            return SelectInst::Create(IsExists(value, block), result, GetEmpty(ctx.Codegen.GetContext()), "select", block);
        return result;
    }
#endif
};

template<bool IsOptional>
class TToBytesDecimalWrapper : public TMutableComputationNode<TToBytesDecimalWrapper<IsOptional>> {
using TBaseComputation = TMutableComputationNode<TToBytesDecimalWrapper<IsOptional>>;
public:
    TToBytesDecimalWrapper(TComputationMutables& mutables, IComputationNode* data)
        : TBaseComputation(mutables, EValueRepresentation::String)
        , Data(data)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
        const auto& value = this->Data->GetValue(ctx);
        if (IsOptional && !value)
            return NUdf::TUnboxedValuePod();

        return MakeString(NUdf::TStringRef(reinterpret_cast<const char*>(&value), 15));
    }

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

    IComputationNode* const Data;
};

template<bool IsOptional, typename Type>
class TToBytesTzTypeWrapper : public TDecoratorComputationNode<TToBytesTzTypeWrapper<IsOptional, Type>> {
using TBaseComputation = TDecoratorComputationNode<TToBytesTzTypeWrapper<IsOptional, Type>>;
public:
    TToBytesTzTypeWrapper(IComputationNode* data)
        : TBaseComputation(data)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const {
        if (IsOptional && !value)
            return NUdf::TUnboxedValuePod();

        const auto v = NYql::SwapBytes(value.Get<Type>());
        const auto tzId = NYql::SwapBytes(value.GetTimezoneId());
        char buf[sizeof(Type) + sizeof(ui16)];
        std::memcpy(buf, &v, sizeof(v));
        std::memcpy(buf + sizeof(Type), &tzId, sizeof(tzId));
        return NUdf::TUnboxedValuePod::Embedded(NUdf::TStringRef(buf, sizeof(buf)));
    }
};

class TToBytesWrapper : public TDecoratorCodegeneratorNode<TToBytesWrapper> {
using TBaseComputation = TDecoratorCodegeneratorNode<TToBytesWrapper>;
public:
    TToBytesWrapper(IComputationNode* optional)
        : TBaseComputation(optional)
    {}

    NUdf::TUnboxedValuePod DoCalculate(TComputationContext&, const NUdf::TUnboxedValuePod& value) const { return value; }
#ifndef MKQL_DISABLE_CODEGEN
    Value* DoGenerateGetValue(const TCodegenContext&, Value* value, BasicBlock*&) const { return value; }
#endif
};

}

IComputationNode* WrapToBytes(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
    MKQL_ENSURE(callable.GetInputsCount() == 1, "Expected 1 arg");

    bool isOptional;
    const auto dataType = UnpackOptionalData(callable.GetInput(0), isOptional);
    const auto data = LocateNode(ctx.NodeLocator, callable, 0);
    switch(dataType->GetSchemeType()) {
#define MAKE_PRIMITIVE_TYPE_BYTES(type, layout) \
        case NUdf::TDataType<type>::Id: \
            if (isOptional) \
                return new TToBytesPrimitiveTypeWrapper<true, layout>(data); \
            else \
                return new TToBytesPrimitiveTypeWrapper<false, layout>(data);

        KNOWN_FIXED_VALUE_TYPES(MAKE_PRIMITIVE_TYPE_BYTES)
#undef MAKE_PRIMITIVE_TYPE_BYTES
#define MAKE_TZ_TYPE_BYTES(type, layout) \
        case NUdf::TDataType<type>::Id: \
            if (isOptional) \
                return new TToBytesTzTypeWrapper<true, layout>(data); \
            else \
                return new TToBytesTzTypeWrapper<false, layout>(data);

        MAKE_TZ_TYPE_BYTES(NUdf::TTzDate, ui16);
        MAKE_TZ_TYPE_BYTES(NUdf::TTzDatetime, ui32);
        MAKE_TZ_TYPE_BYTES(NUdf::TTzTimestamp, ui64);
        MAKE_TZ_TYPE_BYTES(NUdf::TTzDate32, i32);
        MAKE_TZ_TYPE_BYTES(NUdf::TTzDatetime64, i64);
        MAKE_TZ_TYPE_BYTES(NUdf::TTzTimestamp64, i64);
#undef MAKE_TZ_TYPE_BYTES
        case NUdf::TDataType<NUdf::TDecimal>::Id: {
            if (isOptional) \
                return new TToBytesDecimalWrapper<true>(ctx.Mutables, data);
            else
                return new TToBytesDecimalWrapper<false>(ctx.Mutables, data);
        }
        default:
            break;
    }

    return new TToBytesWrapper(data);
}

}
}