summaryrefslogtreecommitdiffstats
path: root/yql/essentials/minikql/comp_nodes/ut/mkql_block_map_join_ut_utils.h
blob: 01cbca7282b9f8e16e4c395bfc117954f11e08f6 (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
#pragma once

#include "mkql_computation_node_ut.h"

#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>

namespace NKikimr {
namespace NMiniKQL {

inline bool IsOptionalOrNull(const TType* type) {
    return type->IsOptional() || type->IsNull() || type->IsPg();
}

TType* MakeBlockTupleType(TProgramBuilder& pgmBuilder, TType* tupleType, bool scalar);
TType* MakeJoinType(TProgramBuilder& pgmBuilder, EJoinKind joinKind,
                    TType* leftStreamType, const TVector<ui32>& leftKeyDrops,
                    TType* rightListType, const TVector<ui32>& rightKeyDrops);

NUdf::TUnboxedValuePod ToBlocks(TComputationContext& ctx, size_t blockSize,
                                const TArrayRef<TType* const> types, const NUdf::TUnboxedValuePod& values);
NUdf::TUnboxedValuePod MakeUint64ScalarBlock(TComputationContext& ctx, size_t blockSize,
                                             const TArrayRef<TType* const> types, const NUdf::TUnboxedValuePod& values);
NUdf::TUnboxedValuePod FromBlocks(TComputationContext& ctx,
                                  const TArrayRef<TType* const> types, const NUdf::TUnboxedValuePod& values);

TComputationNodeFactory GetNodeFactory();
TRuntimeNode ThrottleStream(TProgramBuilder& pgmBuilder, TRuntimeNode stream);
TRuntimeNode DethrottleStream(TProgramBuilder& pgmBuilder, TRuntimeNode stream);

TVector<NUdf::TUnboxedValue> ConvertListToVector(const NUdf::TUnboxedValue& list);
void CompareResults(const TType* type, const NUdf::TUnboxedValue& expected, const NUdf::TUnboxedValue& got);

TVector<TString> GenerateValues(size_t level);
TSet<ui64> GenerateFibonacci(size_t count);

//
// Auxiliary routines to build list nodes from the given vectors.
//

struct TTypeMapperBase {
    TProgramBuilder& Pb;
    TType* ItemType;
    auto GetType() {
        return ItemType;
    }
};

template <typename Type>
struct TTypeMapper: TTypeMapperBase {
    TTypeMapper(TProgramBuilder& pb)
        : TTypeMapperBase{pb, pb.NewDataType(NUdf::TDataType<Type>::Id)}
    {
    }
    auto GetValue(const Type& value) {
        return Pb.NewDataLiteral<Type>(value);
    }
};

template <>
struct TTypeMapper<TString>: TTypeMapperBase {
    TTypeMapper(TProgramBuilder& pb)
        : TTypeMapperBase{pb, pb.NewDataType(NUdf::EDataSlot::String)}
    {
    }
    auto GetValue(const TString& value) {
        return Pb.NewDataLiteral<NUdf::EDataSlot::String>(value);
    }
};

template <typename TNested>
class TTypeMapper<std::optional<TNested>>: TTypeMapper<TNested> {
    using TBase = TTypeMapper<TNested>;

public:
    TTypeMapper(TProgramBuilder& pb)
        : TBase(pb)
    {
    }
    auto GetType() {
        return TBase::Pb.NewOptionalType(TBase::GetType());
    }
    auto GetValue(const std::optional<TNested>& value) {
        if (value == std::nullopt) {
            return TBase::Pb.NewEmptyOptional(GetType());
        } else {
            return TBase::Pb.NewOptional(TBase::GetValue(*value));
        }
    }
};

template <typename Type>
const TVector<const TRuntimeNode> BuildListNodes(TProgramBuilder& pb,
                                                 const TVector<Type>& vector) {
    TTypeMapper<Type> mapper(pb);

    TRuntimeNode::TList listItems;
    std::transform(vector.cbegin(), vector.cend(), std::back_inserter(listItems),
                   [&](const auto value) {
                       return mapper.GetValue(value);
                   });

    return {pb.NewList(mapper.GetType(), listItems)};
}

template <typename Type, typename... Tail>
const TVector<const TRuntimeNode> BuildListNodes(TProgramBuilder& pb,
                                                 const TVector<Type>& vector, Tail... vectors) {
    const auto frontList = BuildListNodes(pb, vector);
    const auto tailLists = BuildListNodes(pb, std::forward<Tail>(vectors)...);
    TVector<const TRuntimeNode> lists;
    lists.reserve(tailLists.size() + 1);
    lists.push_back(frontList.front());
    ;
    for (const auto& list : tailLists) {
        lists.push_back(list);
    }
    return lists;
}

template <typename... TVectors>
const std::pair<TType*, NUdf::TUnboxedValue> ConvertVectorsToTuples(
    TSetup<false>& setup, TVectors... vectors) {
    TProgramBuilder& pb = *setup.PgmBuilder;
    const auto lists = BuildListNodes(pb, std::forward<TVectors>(vectors)...);
    const auto tuplesNode = pb.Zip(lists);
    const auto tuplesNodeType = tuplesNode.GetStaticType();
    const auto tuples = setup.BuildGraph(tuplesNode)->GetValue();
    return std::make_pair(tuplesNodeType, tuples);
}

} // namespace NMiniKQL
} // namespace NKikimr