diff --git a/velox/docs/functions/json.rst b/velox/docs/functions/json.rst index 8ebd16525699f..a2da9c59e1556 100644 --- a/velox/docs/functions/json.rst +++ b/velox/docs/functions/json.rst @@ -124,6 +124,13 @@ JSON Functions SELECT json_array_contains('[1, 2, 3]', 2); +.. function:: json_format(json) -> varchar + + Serializes the input JSON value to JSON text conforming to RFC 7159. + The JSON value can be a JSON object, a JSON array, a JSON string, a JSON number, true, false or null. + + SELECT json_format(JSON '{"a": 1, "b": 2}') + .. function:: json_size(json, value) -> bigint Returns the size of the ``value``. For ``objects`` or ``arrays``, the size diff --git a/velox/functions/prestosql/CMakeLists.txt b/velox/functions/prestosql/CMakeLists.txt index f61c7f704909e..3df3a27062975 100644 --- a/velox/functions/prestosql/CMakeLists.txt +++ b/velox/functions/prestosql/CMakeLists.txt @@ -32,6 +32,7 @@ add_library( FromUnixTime.cpp GreatestLeast.cpp InPredicate.cpp + JsonFunctions.cpp Map.cpp MapEntries.cpp MapKeysAndValues.cpp diff --git a/velox/functions/prestosql/JsonFunctions.cpp b/velox/functions/prestosql/JsonFunctions.cpp new file mode 100644 index 0000000000000..72eea5373bb96 --- /dev/null +++ b/velox/functions/prestosql/JsonFunctions.cpp @@ -0,0 +1,72 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/expression/VectorFunction.h" +#include "velox/functions/prestosql/Comparisons.h" + +namespace facebook::velox::functions { + +namespace { +class JsonFormatFunction : public exec::VectorFunction { + public: + void apply( + const SelectivityVector& rows, + std::vector& args, + const TypePtr& /* outputType */, + exec::EvalCtx& context, + VectorPtr& result) const override { + VectorPtr localResult; + + // Input can be constant or flat. + assert(args.size() > 0); + const auto& arg = args[0]; + if (arg->isConstantEncoding()) { + auto value = arg->as>()->valueAt(0); + localResult = std::make_shared>( + context.pool(), rows.end(), false, VARCHAR(), std::move(value)); + } else { + auto flatInput = arg->asFlatVector(); + + auto stringBuffers = flatInput->stringBuffers(); + VELOX_CHECK_LE(rows.end(), flatInput->size()); + localResult = std::make_shared>( + context.pool(), + VARCHAR(), + nullptr, + rows.end(), + flatInput->values(), + std::move(stringBuffers)); + } + + context.moveOrCopyResult(localResult, rows, result); + } + + static std::vector> signatures() { + // json -> varchar + return {exec::FunctionSignatureBuilder() + .returnType("varchar") + .argumentType("json") + .build()}; + } +}; + +} // namespace + +VELOX_DECLARE_VECTOR_FUNCTION( + udf_json_format, + JsonFormatFunction::signatures(), + std::make_unique()); +} // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp index 4f559a79717c5..015d834689ef3 100644 --- a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp @@ -35,6 +35,7 @@ void registerJsonFunctions() { registerFunction( {"json_array_contains"}); registerFunction({"json_size"}); + VELOX_REGISTER_VECTOR_FUNCTION(udf_json_format, "json_format"); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp index 313ee1515b162..9ea4eb0cf5e0e 100644 --- a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp @@ -58,6 +58,48 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest { } }; +TEST_F(JsonFunctionsTest, jsonFormat) { + const auto jsonFormat = [&](std::optional value) { + return evaluateOnce( + "json_format(c0)", {value}, {JSON()}); + }; + + EXPECT_EQ(jsonFormat(std::nullopt), std::nullopt); + EXPECT_EQ(jsonFormat(R"(true)"), "true"); + EXPECT_EQ(jsonFormat(R"(null)"), "null"); + EXPECT_EQ(jsonFormat(R"(42)"), "42"); + EXPECT_EQ(jsonFormat(R"("abc")"), "\"abc\""); + EXPECT_EQ(jsonFormat(R"([1, 2, 3])"), "[1, 2, 3]"); + EXPECT_EQ(jsonFormat(R"({"k1":"v1"})"), "{\"k1\":\"v1\"}"); + + auto data = makeRowVector({makeFlatVector( + {"This is a long sentence", "This is some other sentence"})}); + + auto result = evaluate("json_format(c0)", data); + auto expected = makeFlatVector( + {"This is a long sentence", "This is some other sentence"}); + facebook::velox::test::assertEqualVectors(expected, result); + + data = makeRowVector({makeConstant("apple", 2)}); + result = evaluate("json_format(c0)", data); + expected = makeFlatVector({{"apple", "apple"}}); + + facebook::velox::test::assertEqualVectors(expected, result); + + data = makeRowVector( + {makeFlatVector({true, false}), + makeFlatVector( + {"This is a long sentence", "This is some other sentence"})}); + + result = evaluate("if(c0, 'foo', json_format(c1))", data); + expected = makeFlatVector({"foo", "This is some other sentence"}); + facebook::velox::test::assertEqualVectors(expected, result); + + result = evaluate("if(c0, json_format(c1), 'bar')", data); + expected = makeFlatVector({"This is a long sentence", "bar"}); + facebook::velox::test::assertEqualVectors(expected, result); +} + TEST_F(JsonFunctionsTest, isJsonScalarSignatures) { auto signatures = getSignatureStrings("is_json_scalar"); ASSERT_EQ(1, signatures.size());