Skip to content

Commit

Permalink
Add Json function - json_array_contains
Browse files Browse the repository at this point in the history
  • Loading branch information
pramodsatya committed Aug 16, 2022
1 parent 68f5694 commit e1a8bfd
Show file tree
Hide file tree
Showing 4 changed files with 249 additions and 0 deletions.
8 changes: 8 additions & 0 deletions velox/docs/functions/json.rst
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,14 @@ JSON Functions
SELECT is_json_scalar('1'); *-- true*
SELECT is_json_scalar('[1, 2, 3]'); *-- false*

.. function:: json_array_contains(json, value) -> boolean

Determine if ``value`` exists in ``json`` (a string containing a JSON
array). ``value`` could be a boolean, integer (``int64_t``), double or
a string. Returns NULL if ``json`` is not an array::

SELECT json_array_contains('[1, 2, 3]', 2);

============
JSON Vectors
============
Expand Down
81 changes: 81 additions & 0 deletions velox/functions/prestosql/JsonFunctions.h
Original file line number Diff line number Diff line change
Expand Up @@ -57,4 +57,85 @@ struct JsonExtractScalarFunction {
}
};

template <typename T>
struct JsonArrayContainsFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);

FOLLY_ALWAYS_INLINE bool call(
bool& result,
const arg_type<Varchar>& json,
const arg_type<bool>& value) {
auto parsedJson = folly::parseJson(json);
if (!parsedJson.isArray()) {
return false;
}

result = false;
for (const auto& v : parsedJson) {
if (v.isBool() && v == value) {
result = true;
break;
}
}
return true;
}

FOLLY_ALWAYS_INLINE bool call(
bool& result,
const arg_type<Varchar>& json,
const arg_type<int64_t>& value) {
auto parsedJson = folly::parseJson(json);
if (!parsedJson.isArray()) {
return false;
}

result = false;
for (const auto& v : parsedJson) {
if (v.isInt() && v == value) {
result = true;
break;
}
}
return true;
}

FOLLY_ALWAYS_INLINE bool call(
bool& result,
const arg_type<Varchar>& json,
const arg_type<double>& value) {
auto parsedJson = folly::parseJson(json);
if (!parsedJson.isArray()) {
return false;
}

result = false;
for (const auto& v : parsedJson) {
if (v.isDouble() && v == value) {
result = true;
break;
}
}
return true;
}

FOLLY_ALWAYS_INLINE bool call(
bool& result,
const arg_type<Varchar>& json,
const arg_type<Varchar>& value) {
auto parsedJson = folly::parseJson(json);
if (!parsedJson.isArray()) {
return false;
}

result = false;
for (const auto& v : parsedJson) {
if (v.isString() && v == value) {
result = true;
break;
}
}
return true;
}
};

} // namespace facebook::velox::functions
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,14 @@ void registerJsonFunctions() {
registerFunction<IsJsonScalarFunction, bool, Varchar>({"is_json_scalar"});
registerFunction<JsonExtractScalarFunction, Varchar, Varchar, Varchar>(
{"json_extract_scalar"});
registerFunction<JsonArrayContainsFunction, bool, Varchar, bool>(
{"json_array_contains"});
registerFunction<JsonArrayContainsFunction, bool, Varchar, int64_t>(
{"json_array_contains"});
registerFunction<JsonArrayContainsFunction, bool, Varchar, double>(
{"json_array_contains"});
registerFunction<JsonArrayContainsFunction, bool, Varchar, Varchar>(
{"json_array_contains"});
}

} // namespace facebook::velox::functions
152 changes: 152 additions & 0 deletions velox/functions/prestosql/tests/JsonFunctionsTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,30 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest {
std::optional<bool> is_json_scalar(std::optional<std::string> json) {
return evaluateOnce<bool>("is_json_scalar(c0)", json);
}

std::optional<bool> json_array_contains_bool(
std::optional<std::string> json,
std::optional<bool> value) {
return evaluateOnce<bool>("json_array_contains(c0, c1)", json, value);
}

std::optional<bool> json_array_contains_int(
std::optional<std::string> json,
std::optional<int64_t> value) {
return evaluateOnce<bool>("json_array_contains(c0, c1)", json, value);
}

std::optional<bool> json_array_contains_double(
std::optional<std::string> json,
std::optional<double> value) {
return evaluateOnce<bool>("json_array_contains(c0, c1)", json, value);
}

std::optional<bool> json_array_contains_string(
std::optional<std::string> json,
std::optional<std::string> value) {
return evaluateOnce<bool>("json_array_contains(c0, c1)", json, value);
}
};

TEST_F(JsonFunctionsTest, isJsonScalar) {
Expand All @@ -45,6 +69,134 @@ TEST_F(JsonFunctionsTest, isJsonScalar) {
EXPECT_EQ(is_json_scalar(R"({"k1":""})"), false);
}

TEST_F(JsonFunctionsTest, jsonArrayContainsBool) {
EXPECT_EQ(json_array_contains_bool(R"([])", 0), false);
EXPECT_EQ(json_array_contains_bool(R"([1, 2, 3])", 1), false);
EXPECT_EQ(json_array_contains_bool(R"([1.2, 2.3, 3.4])", 2.3), false);
EXPECT_EQ(
json_array_contains_bool(R"(["hello", "presto", "world"])", R"("hello")"),
false);
EXPECT_EQ(json_array_contains_bool(R"(1)", true), std::nullopt);
EXPECT_EQ(
json_array_contains_bool(R"("thefoxjumpedoverthefence")", false),
std::nullopt);
EXPECT_EQ(json_array_contains_bool(R"("")", false), std::nullopt);
EXPECT_EQ(json_array_contains_bool(R"(true)", true), std::nullopt);
EXPECT_EQ(
json_array_contains_bool(R"({"k1":[0,1,2], "k2":"v1"})", true),
std::nullopt);

EXPECT_EQ(json_array_contains_bool(R"([true, false])", true), true);
EXPECT_EQ(json_array_contains_bool(R"([true, true])", false), false);
EXPECT_EQ(
json_array_contains_bool(
R"([false, false, false, false, false, false, false,
false, false, false, false, false, false, true, false, false, false, false])",
true),
true);
EXPECT_EQ(
json_array_contains_bool(
R"([true, true, true, true, true, true, true,
true, true, true, true, true, true, true, true, true, true, true])",
false),
false);
}

TEST_F(JsonFunctionsTest, jsonArrayContainsInt) {
EXPECT_EQ(json_array_contains_int(R"([])", 0), false);
EXPECT_EQ(json_array_contains_int(R"([1.2, 2.3, 3.4])", 2), false);
EXPECT_EQ(
json_array_contains_int(R"(["hello", "presto", "world"])", 2), false);
EXPECT_EQ(json_array_contains_int(R"([false, false, false])", 17), false);
EXPECT_EQ(json_array_contains_int(R"(1)", 1), std::nullopt);
EXPECT_EQ(
json_array_contains_int(R"("thefoxjumpedoverthefence")", 1),
std::nullopt);
EXPECT_EQ(json_array_contains_int(R"("")", 1), std::nullopt);
EXPECT_EQ(json_array_contains_int(R"(true)", 1), std::nullopt);
EXPECT_EQ(
json_array_contains_int(R"({"k1":[0,1,2], "k2":"v1"})", 1), std::nullopt);

EXPECT_EQ(json_array_contains_int(R"([1, 2, 3])", 1), true);
EXPECT_EQ(json_array_contains_int(R"([1, 2, 3])", 4), false);
EXPECT_EQ(
json_array_contains_int(
R"([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20])",
17),
true);
EXPECT_EQ(
json_array_contains_int(
R"([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20])",
23),
false);
}

TEST_F(JsonFunctionsTest, jsonArrayContainsDouble) {
EXPECT_EQ(json_array_contains_double(R"([])", 2.3), false);
EXPECT_EQ(json_array_contains_double(R"([1, 2, 3])", 2.3), false);
EXPECT_EQ(
json_array_contains_double(R"(["hello", "presto", "world"])", 2.3),
false);
EXPECT_EQ(json_array_contains_double(R"([false, false, false])", 2.3), false);
EXPECT_EQ(json_array_contains_double(R"(1)", 2.3), std::nullopt);
EXPECT_EQ(
json_array_contains_double(R"("thefoxjumpedoverthefence")", 2.3),
std::nullopt);
EXPECT_EQ(json_array_contains_double(R"("")", 2.3), std::nullopt);
EXPECT_EQ(json_array_contains_double(R"(true)", 2.3), std::nullopt);
EXPECT_EQ(
json_array_contains_double(R"({"k1":[0,1,2], "k2":"v1"})", 2.3),
std::nullopt);

EXPECT_EQ(json_array_contains_double(R"([1.2, 2.3, 3.4])", 2.3), true);
EXPECT_EQ(json_array_contains_double(R"([1.2, 2.3, 3.4])", 2.4), false);
EXPECT_EQ(
json_array_contains_double(
R"([1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5])",
4.5),
true);
EXPECT_EQ(
json_array_contains_double(
R"([1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5, 1.2, 2.3, 3.4, 4.5])",
4.3),
false);
}

TEST_F(JsonFunctionsTest, jsonArrayContainsString) {
EXPECT_EQ(json_array_contains_string(R"([])", ""), false);
EXPECT_EQ(json_array_contains_string(R"([1, 2, 3])", "1"), false);
EXPECT_EQ(json_array_contains_string(R"([1.2, 2.3, 3.4])", "2.3"), false);
EXPECT_EQ(json_array_contains_string(R"([true, false])", R"("true")"), false);
EXPECT_EQ(json_array_contains_string(R"(1)", "1"), std::nullopt);
EXPECT_EQ(
json_array_contains_string(R"("thefoxjumpedoverthefence")", "1"),
std::nullopt);
EXPECT_EQ(json_array_contains_string(R"("")", "1"), std::nullopt);
EXPECT_EQ(json_array_contains_string(R"(true)", "1"), std::nullopt);
EXPECT_EQ(
json_array_contains_string(R"({"k1":[0,1,2], "k2":"v1"})", "1"),
std::nullopt);

EXPECT_EQ(
json_array_contains_string(R"(["hello", "presto", "world"])", "presto"),
true);
EXPECT_EQ(
json_array_contains_string(R"(["hello", "presto", "world"])", "nation"),
false);
EXPECT_EQ(
json_array_contains_string(
R"(["hello", "presto", "world", "hello", "presto", "world", "hello", "presto", "world", "hello",
"presto", "world", "hello", "presto", "world", "hello", "presto", "world"])",
"hello"),
true);
EXPECT_EQ(
json_array_contains_string(
R"(["hello", "presto", "world", "hello", "presto", "world", "hello", "presto", "world", "hello",
"presto", "world", "hello", "presto", "world", "hello", "presto", "world"])",
"hola"),
false);
}

} // namespace

} // namespace facebook::velox::functions::prestosql

0 comments on commit e1a8bfd

Please sign in to comment.