ClickHouse

1
#include <Columns/ColumnString.h>
2
#include <DataTypes/DataTypeString.h>
3
#include <Functions/FunctionFactory.h>
4
#include <Functions/FunctionHelpers.h>
5
#include <Functions/IFunction.h>
6
#include <Interpreters/Context.h>
7
#include <IO/ReadBufferFromString.h>
8
#include <Common/FieldVisitorToString.h>
9
#include "config.h"
10

11
#if USE_RAPIDJSON
12

13
#include "rapidjson/document.h"
14
#include "rapidjson/writer.h"
15
#include "rapidjson/stringbuffer.h"
16
#include "rapidjson/filewritestream.h"
17
#include "rapidjson/prettywriter.h"
18
#include "rapidjson/filereadstream.h"
19

20

21
namespace DB
22
{
23

24
namespace ErrorCodes
25
{
26
    extern const int BAD_ARGUMENTS;
27
    extern const int ILLEGAL_COLUMN;
28
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
29
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
30
}
31

32
namespace
33
{
34
    // select jsonMergePatch('{"a":1}','{"name": "joey"}','{"name": "tom"}','{"name": "zoey"}');
35
    //           ||
36
    //           \/
37
    // ┌───────────────────────┐
38
    // │ {"a":1,"name":"zoey"} │
39
    // └───────────────────────┘
40
    class FunctionjsonMergePatch : public IFunction
41
    {
42
    public:
43
        static constexpr auto name = "jsonMergePatch";
44
        static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionjsonMergePatch>(); }
45

46
        String getName() const override { return name; }
47
        bool isVariadic() const override { return true; }
48
        bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
49

50
        size_t getNumberOfArguments() const override { return 0; }
51
        bool useDefaultImplementationForConstants() const override { return true; }
52

53
        DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
54
        {
55
            if (arguments.empty())
56
                throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} requires at least one argument.", getName());
57

58
            for (const auto & arg : arguments)
59
                if (!isString(arg.type))
60
                    throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Function {} requires string arguments", getName());
61

62
            return std::make_shared<DataTypeString>();
63
        }
64

65
        ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
66
        {
67
            chassert(!arguments.empty());
68

69
            rapidjson::Document::AllocatorType allocator;
70
            std::function<void(rapidjson::Value &, const rapidjson::Value &)> merge_objects;
71

72
            merge_objects = [&merge_objects, &allocator](rapidjson::Value & dest, const rapidjson::Value & src) -> void
73
            {
74
                if (!src.IsObject())
75
                    return;
76

77
                for (auto it = src.MemberBegin(); it != src.MemberEnd(); ++it)
78
                {
79
                    rapidjson::Value key(it->name, allocator);
80
                    rapidjson::Value value(it->value, allocator);
81
                    if (dest.HasMember(key))
82
                    {
83
                        if (dest[key].IsObject() && value.IsObject())
84
                            merge_objects(dest[key], value);
85
                        else
86
                            dest[key] = value;
87
                    }
88
                    else
89
                    {
90
                        dest.AddMember(key, value, allocator);
91
                    }
92
                }
93
            };
94

95
            auto parse_json_document = [](const ColumnString & column, rapidjson::Document & document, size_t i)
96
            {
97
                auto str_ref = column.getDataAt(i);
98
                const char * json = str_ref.data;
99

100
                document.Parse(json);
101
                if (document.HasParseError() || !document.IsObject())
102
                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Wrong JSON string to merge. Expected JSON object");
103
            };
104

105
            const bool is_first_const = isColumnConst(*arguments[0].column);
106
            const auto * first_column_arg_string = is_first_const
107
                        ? checkAndGetColumnConstData<ColumnString>(arguments[0].column.get())
108
                        : checkAndGetColumn<ColumnString>(arguments[0].column.get());
109

110
            if (!first_column_arg_string)
111
                throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Arguments of function {} must be strings", getName());
112

113
            std::vector<rapidjson::Document> merged_jsons;
114
            merged_jsons.reserve(input_rows_count);
115

116
            for (size_t i = 0; i < input_rows_count; ++i)
117
            {
118
                auto & merged_json = merged_jsons.emplace_back(rapidjson::Type::kObjectType, &allocator);
119
                if (is_first_const)
120
                    parse_json_document(*first_column_arg_string, merged_json, 0);
121
                else
122
                    parse_json_document(*first_column_arg_string, merged_json, i);
123
            }
124

125
            for (size_t col_idx = 1; col_idx < arguments.size(); ++col_idx)
126
            {
127
                const bool is_const = isColumnConst(*arguments[col_idx].column);
128
                const auto * column_arg_string = is_const
129
                            ? checkAndGetColumnConstData<ColumnString>(arguments[col_idx].column.get())
130
                            : checkAndGetColumn<ColumnString>(arguments[col_idx].column.get());
131

132
                if (!column_arg_string)
133
                    throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Arguments of function {} must be strings", getName());
134

135
                for (size_t i = 0; i < input_rows_count; ++i)
136
                {
137
                    rapidjson::Document document(&allocator);
138
                    if (is_const)
139
                        parse_json_document(*column_arg_string, document, 0);
140
                    else
141
                        parse_json_document(*column_arg_string, document, i);
142
                    merge_objects(merged_jsons[i], document);
143
                }
144
            }
145

146
            auto result = ColumnString::create();
147
            auto & result_string = assert_cast<ColumnString &>(*result);
148
            rapidjson::CrtAllocator buffer_allocator;
149

150
            for (size_t i = 0; i < input_rows_count; ++i)
151
            {
152
                rapidjson::StringBuffer buffer(&buffer_allocator);
153
                rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
154

155
                merged_jsons[i].Accept(writer);
156
                result_string.insertData(buffer.GetString(), buffer.GetSize());
157
            }
158

159
            return result;
160
        }
161
    };
162

163
}
164

165
REGISTER_FUNCTION(jsonMergePatch)
166
{
167
    factory.registerFunction<FunctionjsonMergePatch>(FunctionDocumentation{
168
        .description="Returns the merged JSON object string, which is formed by merging multiple JSON objects."});
169
}
170

171
}
172

173
#endif
174