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
|
#include <Columns/ColumnNullable.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnVector.h>
#include <DataTypes/DataTypeNullable.h>
#include <DataTypes/DataTypesNumber.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/IFunction.h>
#include <Interpreters/Context.h>
#include "clickhouse_config.h"
#if USE_SIMDJSON
# include <Common/JSONParsers/SimdJSONParser.h>
#elif USE_RAPIDJSON
# include <Common/JSONParsers/RapidJSONParser.h>
#else
# include <Common/JSONParsers/DummyJSONParser.h>
#endif
namespace DB
{
namespace ErrorCodes
{
extern const int ILLEGAL_COLUMN;
}
namespace
{
/// JSONArrayLength(json)
class FunctionJSONArrayLength : public IFunction
{
public:
static constexpr auto name = "JSONArrayLength";
static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionJSONArrayLength>(); }
String getName() const override { return name; }
bool isVariadic() const override { return false; }
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
size_t getNumberOfArguments() const override { return 1; }
bool useDefaultImplementationForConstants() const override { return true; }
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
{
auto args = FunctionArgumentDescriptors{
{"json", &isString<IDataType>, nullptr, "String"},
};
validateFunctionArgumentTypes(*this, arguments, args);
return std::make_shared<DataTypeNullable>(std::make_shared<DataTypeUInt64>());
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
{
const ColumnPtr column = arguments[0].column;
const ColumnString * col = typeid_cast<const ColumnString *>(column.get());
if (!col)
throw Exception(ErrorCodes::ILLEGAL_COLUMN, "First argument of function {} must be string", getName());
auto null_map = ColumnUInt8::create();
auto data = ColumnUInt64::create();
null_map->reserve(input_rows_count);
data->reserve(input_rows_count);
#if USE_SIMDJSON
SimdJSONParser parser;
SimdJSONParser::Element element;
#elif USE_RAPIDJSON
RapidJSONParser parser;
RapidJSONParser::Element element;
#else
DummyJSONParser parser;
DummyJSONParser::Element element;
#endif
for (size_t i = 0; i < input_rows_count; ++i)
{
auto str_ref = col->getDataAt(i);
std::string_view str_view(str_ref.data, str_ref.size);
bool ok = parser.parse(std::move(str_view), element);
if (!ok || !element.isArray())
{
null_map->insertValue(1);
data->insertDefault();
}
else
{
auto array = element.getArray();
null_map->insertValue(0);
data->insertValue(array.size());
}
}
return ColumnNullable::create(std::move(data), std::move(null_map));
}
};
}
REGISTER_FUNCTION(JSONArrayLength)
{
factory.registerFunction<FunctionJSONArrayLength>(FunctionDocumentation{
.description="Returns the number of elements in the outermost JSON array. The function returns NULL if input JSON string is invalid."});
/// For Spark compatibility.
factory.registerAlias("JSON_ARRAY_LENGTH", "JSONArrayLength", FunctionFactory::CaseInsensitive);
}
}
|
