aboutsummaryrefslogtreecommitdiffstats
path: root/yql/essentials/udfs/examples/lists/lists_udf.cpp
blob: 35dbb57dfa35407d4cac1f99253c840d296f69bd (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 <yql/essentials/public/udf/udf_value.h>
#include <yql/essentials/public/udf/udf_value_builder.h>
#include <yql/essentials/public/udf/udf_type_builder.h>
#include <yql/essentials/public/udf/udf_registrator.h>

#include <util/generic/yexception.h>
#include <vector>
#include <array>

using namespace NKikimr;
using namespace NUdf;

namespace {

//////////////////////////////////////////////////////////////////////////////
// TNumbersList
//////////////////////////////////////////////////////////////////////////////
class TNumbers: public TBoxedValue
{
public:
    static TStringRef Name() {
        static auto name = TStringRef::Of("Numbers");
        return name;
    }

private:
    TUnboxedValue Run(
            const IValueBuilder* valueBuilder,
            const TUnboxedValuePod* args) const override
    {
        const auto appendPrepend = args[0].AsStringRef();
        const auto count = args[1].Get<ui32>();
        std::vector<TUnboxedValue> list(count);
        ui32 i = 0U;
        if (TStringRef::Of("Append") == appendPrepend) {
            for (auto it = list.begin(); list.end() != it; ++it) {
                *it = TUnboxedValuePod(i++);
            }
        }
        else if (TStringRef::Of("Prepend") == appendPrepend) {
            for (auto it = list.rbegin(); list.rend() != it; ++it) {
                *it = TUnboxedValuePod(i++);
            }
        }

        return valueBuilder->NewList(list.data(), list.size());
    }
};

//////////////////////////////////////////////////////////////////////////////
// TExtend
//////////////////////////////////////////////////////////////////////////////
class TExtend: public TBoxedValue
{
public:
    static TStringRef Name() {
        static auto name = TStringRef::Of("Extend");
        return name;
    }

private:
    TUnboxedValue Run(
            const IValueBuilder* valueBuilder,
            const TUnboxedValuePod* args) const override
    {
        std::array<TUnboxedValue, 2U> list = {{TUnboxedValuePod(args[0]), TUnboxedValuePod(args[1])}};
        return valueBuilder->NewList(list.data(), list.size());
    }
};

//////////////////////////////////////////////////////////////////////////////
// TListsModule
//////////////////////////////////////////////////////////////////////////////
class TListsModule: public IUdfModule
{
public:
    TStringRef Name() const {
        return TStringRef::Of("Lists");
    }

    void CleanupOnTerminate() const final {}

    void GetAllFunctions(IFunctionsSink& sink) const final {
        sink.Add(TNumbers::Name());
        sink.Add(TExtend::Name());
    }

    void BuildFunctionTypeInfo(
            const TStringRef& name,
            TType* userType,
            const TStringRef& typeConfig,
            ui32 flags,
            IFunctionTypeInfoBuilder& builder) const final
    {
        try {
            Y_UNUSED(userType);
            Y_UNUSED(typeConfig);

            bool typesOnly = (flags & TFlags::TypesOnly);

            if (TNumbers::Name() == name) {
                // function signature:
                //    List<ui32> Numbers(String, ui32)
                // runConfig: void
                builder.SimpleSignature<TListType<ui32>(char*, ui32)>();

                if (!typesOnly) {
                    builder.Implementation(new TNumbers);
                }
            }
            else if (TExtend::Name() == name) {
                // function signature:
                //    List<ui32> Numbers(List<ui32>, List<ui32>)
                // runConfig: void
                auto listType = builder.List()->Item<ui32>().Build();
                builder.Returns(listType)
                        .Args()->Add(listType).Add(listType).Done();

                if (!typesOnly) {
                    builder.Implementation(new TExtend);
                }
            }
        } catch (const std::exception& e) {
            builder.SetError(CurrentExceptionMessage());
        }
    }
};

} // namespace

REGISTER_MODULES(TListsModule)