diff --git a/velox/docs/functions/json.rst b/velox/docs/functions/json.rst index 91f84ffe1f3c..873239dc84b9 100644 --- a/velox/docs/functions/json.rst +++ b/velox/docs/functions/json.rst @@ -116,6 +116,14 @@ JSON Functions SELECT json_array_length('[1, 2, 3]'); +.. function:: json_array_contains(json, value) -> boolean + + Determine if ``value`` exists in ``json`` (a string containing a JSON + array). ``value`` could be a boolean, bigint, double, or varchar. + Returns NULL if ``json`` is not an array:: + + SELECT json_array_contains('[1, 2, 3]', 2); + ============ JSON Vectors ============ diff --git a/velox/functions/prestosql/JsonFunctions.h b/velox/functions/prestosql/JsonFunctions.h index d6e3da1ac29d..f502938a9a51 100644 --- a/velox/functions/prestosql/JsonFunctions.h +++ b/velox/functions/prestosql/JsonFunctions.h @@ -74,4 +74,44 @@ struct JsonArrayLengthFunction { } }; +template +struct JsonArrayContainsFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + template + FOLLY_ALWAYS_INLINE bool + call(bool& result, const arg_type& json, const TInput& value) { + auto parsedJson = folly::parseJson(json); + if (!parsedJson.isArray()) { + return false; + } + + result = false; + for (const auto& v : parsedJson) { + if constexpr (std::is_same_v) { + if (v.isBool() && v == value) { + result = true; + break; + } + } else if constexpr (std::is_same_v) { + if (v.isInt() && v == value) { + result = true; + break; + } + } else if constexpr (std::is_same_v) { + if (v.isDouble() && v == value) { + result = true; + break; + } + } else { + if (v.isString() && v == value) { + result = true; + break; + } + } + } + return true; + } +}; + } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp index fa92c22efb39..d181a1a6aa29 100644 --- a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp @@ -24,6 +24,14 @@ void registerJsonFunctions() { {"json_extract_scalar"}); registerFunction( {"json_array_length"}); + registerFunction( + {"json_array_contains"}); + registerFunction( + {"json_array_contains"}); + registerFunction( + {"json_array_contains"}); + registerFunction( + {"json_array_contains"}); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp index 60890c664a3c..5f994086579c 100644 --- a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp @@ -30,6 +30,13 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest { std::optional json_array_length(std::optional json) { return evaluateOnce("json_array_length(c0)", json); } + + template + std::optional json_array_contains( + std::optional json, + std::optional value) { + return evaluateOnce("json_array_contains(c0, c1)", json, value); + } }; TEST_F(JsonFunctionsTest, isJsonScalar) { @@ -67,6 +74,182 @@ TEST_F(JsonFunctionsTest, jsonArrayLength) { EXPECT_EQ(json_array_length(R"({"k1":[0,1,2], "k2":"v1"})"), std::nullopt); } +TEST_F(JsonFunctionsTest, jsonArrayContainsBool) { + EXPECT_EQ(json_array_contains(R"([])", true), false); + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", false), false); + EXPECT_EQ(json_array_contains(R"([1.2, 2.3, 3.4])", true), false); + EXPECT_EQ( + json_array_contains(R"(["hello", "presto", "world"])", false), + false); + EXPECT_EQ(json_array_contains(R"(1)", true), std::nullopt); + EXPECT_EQ( + json_array_contains(R"("thefoxjumpedoverthefence")", false), + std::nullopt); + EXPECT_EQ(json_array_contains(R"("")", false), std::nullopt); + EXPECT_EQ(json_array_contains(R"(true)", true), std::nullopt); + EXPECT_EQ( + json_array_contains(R"({"k1":[0,1,2], "k2":"v1"})", true), + std::nullopt); + + EXPECT_EQ(json_array_contains(R"([true, false])", true), true); + EXPECT_EQ(json_array_contains(R"([true, true])", false), false); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", true), true); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", false), + false); + EXPECT_EQ( + json_array_contains( + 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( + 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(R"([])", 0), false); + EXPECT_EQ(json_array_contains(R"([1.2, 2.3, 3.4])", 2), false); + EXPECT_EQ(json_array_contains(R"([1.2, 2.0, 3.4])", 2), false); + EXPECT_EQ( + json_array_contains(R"(["hello", "presto", "world"])", 2), + false); + EXPECT_EQ( + json_array_contains(R"([false, false, false])", 17), false); + EXPECT_EQ(json_array_contains(R"(1)", 1), std::nullopt); + EXPECT_EQ( + json_array_contains(R"("thefoxjumpedoverthefence")", 1), + std::nullopt); + EXPECT_EQ(json_array_contains(R"("")", 1), std::nullopt); + EXPECT_EQ(json_array_contains(R"(true)", 1), std::nullopt); + EXPECT_EQ( + json_array_contains(R"({"k1":[0,1,2], "k2":"v1"})", 1), + std::nullopt); + + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", 1), true); + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", 4), false); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", 123), + true); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", 456), + false); + EXPECT_EQ( + json_array_contains( + 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( + 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(R"([])", 2.3), false); + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", 2.3), false); + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", 2.0), false); + EXPECT_EQ( + json_array_contains(R"(["hello", "presto", "world"])", 2.3), + false); + EXPECT_EQ( + json_array_contains(R"([false, false, false])", 2.3), false); + EXPECT_EQ(json_array_contains(R"(1)", 2.3), std::nullopt); + EXPECT_EQ( + json_array_contains(R"("thefoxjumpedoverthefence")", 2.3), + std::nullopt); + EXPECT_EQ(json_array_contains(R"("")", 2.3), std::nullopt); + EXPECT_EQ(json_array_contains(R"(true)", 2.3), std::nullopt); + EXPECT_EQ( + json_array_contains(R"({"k1":[0,1,2], "k2":"v1"})", 2.3), + std::nullopt); + + EXPECT_EQ(json_array_contains(R"([1.2, 2.3, 3.4])", 2.3), true); + EXPECT_EQ(json_array_contains(R"([1.2, 2.3, 3.4])", 2.4), false); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", 123.456), + true); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", 456.789), + false); + EXPECT_EQ( + json_array_contains( + 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( + 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(R"([])", ""), false); + EXPECT_EQ(json_array_contains(R"([1, 2, 3])", "1"), false); + EXPECT_EQ( + json_array_contains(R"([1.2, 2.3, 3.4])", "2.3"), false); + EXPECT_EQ( + json_array_contains(R"([true, false])", R"("true")"), false); + EXPECT_EQ(json_array_contains(R"(1)", "1"), std::nullopt); + EXPECT_EQ( + json_array_contains(R"("thefoxjumpedoverthefence")", "1"), + std::nullopt); + EXPECT_EQ(json_array_contains(R"("")", "1"), std::nullopt); + EXPECT_EQ(json_array_contains(R"(true)", "1"), std::nullopt); + EXPECT_EQ( + json_array_contains(R"({"k1":[0,1,2], "k2":"v1"})", "1"), + std::nullopt); + + EXPECT_EQ( + json_array_contains( + R"(["hello", "presto", "world"])", "presto"), + true); + EXPECT_EQ( + json_array_contains( + R"(["hello", "presto", "world"])", "nation"), + false); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", "abc"), + true); + EXPECT_EQ( + json_array_contains(R"([123, 123.456, true, "abc"])", "def"), + false); + EXPECT_EQ( + json_array_contains( + 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( + R"(["hello", "presto", "world", "hello", "presto", "world", "hello", "presto", "world", "hello", +"presto", "world", "hello", "presto", "world", "hello", "presto", "world"])", + "hola"), + false); + EXPECT_EQ( + json_array_contains( + R"(["hello", "presto", "world", 1, 2, 3, true, false, 1.2, 2.3, {"k1":[0,1,2], "k2":"v1"}])", + "world"), + true); + EXPECT_EQ( + json_array_contains( + R"(["the fox jumped over the fence", "hello presto world"])", + "hello velox world"), + false); + EXPECT_EQ( + json_array_contains( + R"(["the fox jumped over the fence", "hello presto world"])", + "the fox jumped over the fence"), + true); +} + } // namespace } // namespace facebook::velox::functions::prestosql