summaryrefslogtreecommitdiffstats
path: root/yql/essentials/udfs/common/json/json_udf.cpp
blob: 92316947df6d03b65580e7768ddef0118d6a7cfa (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
#include <yql/essentials/public/udf/udf_helpers.h>

#include <library/cpp/json/easy_parse/json_easy_parser.h>

using namespace NKikimr;
using namespace NUdf;

namespace {
class TGetField: public TBoxedValue {
public:
    typedef bool TTypeAwareMarker;

public:
    static TStringRef Name() {
        return TStringRef::Of("GetField");
    }

    TUnboxedValue Run(
        const IValueBuilder* valueBuilder,
        const TUnboxedValuePod* args) const override {
        if (!args[0]) {
            return valueBuilder->NewEmptyList();
        }

        const TString json(args[0].AsStringRef());
        const TString field(args[1].AsStringRef());

        if (field.empty()) {
            return valueBuilder->NewEmptyList();
        }

        NJson::TJsonParser parser;
        parser.AddField(field, false);

        TVector<TString> result;
        parser.Parse(json, &result);

        TUnboxedValue* items = nullptr;
        const auto list = valueBuilder->NewArray(result.size(), items);
        for (const TString& item : result) {
            *items++ = valueBuilder->NewString(item);
        }

        return list;
    }

    static bool DeclareSignature(
        const TStringRef& name,
        TType* userType,
        IFunctionTypeInfoBuilder& builder,
        bool typesOnly) {
        if (Name() == name) {
            bool useString = true;
            bool isOptional = true;
            if (userType) {
                // support of an overload with Json/Json? input type
                auto typeHelper = builder.TypeInfoHelper();
                auto userTypeInspector = TTupleTypeInspector(*typeHelper, userType);
                if (!userTypeInspector || userTypeInspector.GetElementsCount() < 1) {
                    builder.SetError("Missing or invalid user type.");
                    return true;
                }

                auto argsTypeTuple = userTypeInspector.GetElementType(0);
                auto argsTypeInspector = TTupleTypeInspector(*typeHelper, argsTypeTuple);
                if (!argsTypeInspector) {
                    builder.SetError("Invalid user type - expected tuple.");
                    return true;
                }

                if (argsTypeInspector.GetElementsCount() != 2) {
                    builder.SetError("Invalid user type - expected two arguments.");
                    return true;
                }

                auto inputType = argsTypeInspector.GetElementType(0);
                auto optInspector = TOptionalTypeInspector(*typeHelper, inputType);
                auto dataType = inputType;
                if (optInspector) {
                    dataType = optInspector.GetItemType();
                } else {
                    isOptional = false;
                }

                auto dataInspector = TDataTypeInspector(*typeHelper, dataType);
                if (dataInspector && dataInspector.GetTypeId() == TDataType<TJson>::Id) {
                    useString = false;
                    builder.UserType(userType);
                }
            }

            auto retType = builder.List()->Item<char*>().Build();
            if (useString) {
                builder.Args()->Add(builder.Optional()->Item<char*>().Build()).Add<char*>().Done().Returns(retType);
            } else {
                auto type = builder.SimpleType<TJson>();
                if (isOptional) {
                    builder.Args()->Add(builder.Optional()->Item(type).Build()).Add<char*>().Done().Returns(retType);
                } else {
                    builder.Args()->Add(type).Add<char*>().Done().Returns(retType);
                }
            }

            if (!typesOnly) {
                builder.Implementation(new TGetField);
            }

            builder.IsStrict();
            return true;
        } else {
            return false;
        }
    }
};
} // namespace

SIMPLE_MODULE(TJsonModule,
              TGetField)

REGISTER_MODULES(TJsonModule)