diff --git a/bolt/functions/prestosql/SIMDJsonFunctions.h b/bolt/functions/prestosql/SIMDJsonFunctions.h index 01996d94e..4015e9090 100644 --- a/bolt/functions/prestosql/SIMDJsonFunctions.h +++ b/bolt/functions/prestosql/SIMDJsonFunctions.h @@ -83,6 +83,23 @@ struct SIMDIsJsonScalarFunction { } }; +template +struct SIMDIsJsonObjectFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE static void call( + bool& result, + const arg_type& json) { + ParserContext ctx(json.data(), json.size()); + try { + ctx.parseDocument(); + result = (ctx.jsonDoc.type() == simdjson::ondemand::json_type::object); + } catch (const simdjson::simdjson_error&) { + result = false; + } + } +}; + template struct SonicIsJsonScalarFunction { BOLT_DEFINE_FUNCTION_TYPES(T); @@ -105,6 +122,23 @@ struct SonicIsJsonScalarFunction { } }; +template +struct SonicIsJsonObjectFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE static void call( + bool& result, + const arg_type& json) { + sonic_json::Document jsonDoc; + jsonDoc.Parse(json.data(), json.size()); + if (jsonDoc.HasParseError()) { + result = false; + return; + } + result = jsonDoc.IsObject(); + } +}; + template class WrapperIsJsonScalarFunction { public: @@ -128,6 +162,29 @@ class WrapperIsJsonScalarFunction { bool useSonic = true; }; +template +class WrapperIsJsonObjectFunction { + public: + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void initialize( + const std::vector& /*inputTypes*/, + const core::QueryConfig& config, + const arg_type* /*json*/) { + useSonic = config.enableSonicIsJsonScalar(); + } + + FOLLY_ALWAYS_INLINE void call(bool& result, const arg_type& json) { + if (useSonic) { + SonicIsJsonObjectFunction::call(result, json); + } else { + SIMDIsJsonObjectFunction::call(result, json); + } + } + + bool useSonic = true; +}; + template struct SIMDJsonArrayContainsFunction { BOLT_DEFINE_FUNCTION_TYPES(T); diff --git a/bolt/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/bolt/functions/prestosql/registration/JsonFunctionsRegistration.cpp index f9127de1c..d4f1c61e3 100644 --- a/bolt/functions/prestosql/registration/JsonFunctionsRegistration.cpp +++ b/bolt/functions/prestosql/registration/JsonFunctionsRegistration.cpp @@ -42,6 +42,11 @@ void registerJsonFunctions(const std::string& prefix) { registerFunction( {prefix + "is_json_scalar"}); + registerFunction( + {prefix + "is_json_object"}); + registerFunction( + {prefix + "is_json_object"}); + registerFunction( {prefix + "json_extract_scalar"}); registerFunction( diff --git a/bolt/functions/prestosql/tests/JsonFunctionsTest.cpp b/bolt/functions/prestosql/tests/JsonFunctionsTest.cpp index 92f70d288..72c639800 100644 --- a/bolt/functions/prestosql/tests/JsonFunctionsTest.cpp +++ b/bolt/functions/prestosql/tests/JsonFunctionsTest.cpp @@ -104,6 +104,16 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest { return jsonResult; } + std::optional isJsonObject(std::optional json) { + auto [jsonVector, varcharVector] = makeVectors(json); + auto jsonResult = + evaluateOnce("is_json_object(c0)", makeRowVector({jsonVector})); + auto varcharResult = evaluateOnce( + "is_json_object(c0)", makeRowVector({varcharVector})); + EXPECT_EQ(jsonResult, varcharResult); + return jsonResult; + } + std::optional jsonArrayLength(std::optional json) { auto [jsonVector, varcharVector] = makeVectors(json); auto jsonResult = evaluateOnce( @@ -314,6 +324,13 @@ TEST_F(JsonFunctionsTest, isJsonScalarSignatures) { ASSERT_EQ(1, signatures.count("(varchar) -> boolean")); } +TEST_F(JsonFunctionsTest, isJsonObjectSignatures) { + auto signatures = getSignatureStrings("is_json_object"); + ASSERT_EQ(2, signatures.size()); + ASSERT_EQ(1, signatures.count("(json) -> boolean")); + ASSERT_EQ(1, signatures.count("(varchar) -> boolean")); +} + TEST_F(JsonFunctionsTest, jsonArrayLengthSignatures) { auto signatures = getSignatureStrings("json_array_length"); ASSERT_EQ(2, signatures.size()); @@ -370,6 +387,21 @@ TEST_F(JsonFunctionsTest, isJsonScalar) { EXPECT_EQ(isJsonScalar(R"({"k1":""})"), false); } +TEST_F(JsonFunctionsTest, isJsonObject) { + EXPECT_EQ(isJsonObject(R"({"k1":"v1"})"), true); + EXPECT_EQ(isJsonObject(R"({"k1":[0,1,2]})"), true); + EXPECT_EQ(isJsonObject(R"({})"), true); + + EXPECT_EQ(isJsonObject(R"(1)"), false); + EXPECT_EQ(isJsonObject(R"("hello")"), false); + EXPECT_EQ(isJsonObject(R"(true)"), false); + EXPECT_EQ(isJsonObject(R"(null)"), false); + EXPECT_EQ(isJsonObject(R"([1,2])"), false); + + EXPECT_EQ(isJsonObject(R"(not_json)"), false); + EXPECT_EQ(isJsonObject(R"()"), false); +} + TEST_F(JsonFunctionsTest, jsonArrayLength) { EXPECT_EQ(jsonArrayLength(R"([])"), 0); EXPECT_EQ(jsonArrayLength(R"([1])"), 1);