Add json_format Presto function (#3398)

Summary: Pull Request resolved: #3398 Add 'json_format' udf to prestosql. Doc: https://prestodb.io/docs/current/functions/json.html?highlight=json_format#json_format Differential Revision: D41643635 fbshipit-source-id: 5a56a76e39220b9e4de09dd04b7c35cc1fb13ceb
facebookincubator · Dec 9, 2022 · e939517 · e939517
1 parent 3bc00db
commit e939517
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 0 deletions.
diff --git 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}')
+
 ============
 JSON Vectors
 ============

diff --git a/velox/functions/prestosql/JsonFunctions.h b/velox/functions/prestosql/JsonFunctions.h
@@ -20,6 +20,18 @@
 
 namespace facebook::velox::functions {
 
+template <typename T>
+struct JsonFormatFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Varchar>& jsonString,
+      const arg_type<Json>& json) {
+    folly::parseJson(json.getString());
+    jsonString.setNoCopy(json);
+    return true;
+  }
+};
+
 template <typename T>
 struct IsJsonScalarFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);

diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
@@ -34,6 +34,7 @@ void registerJsonFunctions() {
       {"json_array_contains"});
   registerFunction<JsonArrayContainsFunction, bool, Json, Varchar>(
       {"json_array_contains"});
+  registerFunction<JsonFormatFunction, Varchar, Json>({"json_format"});
 }
 
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
@@ -52,6 +52,37 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest {
   }
 };
 
+TEST_F(JsonFunctionsTest, JsonFormat) {
+  const auto json_format = [&](std::optional<std::string> value) {
+    return evaluateOnce<std::string, std::string>(
+        "json_format(c0)", {value}, {JSON()});
+  };
+
+  EXPECT_EQ(json_format(R"(true)"), "true");
+  EXPECT_EQ(json_format(R"(null)"), "null");
+  EXPECT_EQ(json_format(R"(42)"), "42");
+  EXPECT_EQ(json_format(R"("abc")"), "\"abc\"");
+  EXPECT_EQ(json_format(R"([1, 2, 3])"), "[1, 2, 3]");
+  EXPECT_EQ(json_format(R"({"k1":"v1"})"), "{\"k1\":\"v1\"}");
+
+  // check keys and values are there
+  const std::string jsonStr = json_format(R"({"k1":"v1","k2":"v2"})").value();
+  folly::dynamic object = folly::parseJson(jsonStr);
+
+  std::set<std::string> keys{"k1", "k2"};
+
+  for (const auto& key : object.keys()) {
+    EXPECT_TRUE(keys.find(key.getString()) != keys.end());
+  }
+
+  std::unordered_map<std::string, std::string> jsonMap{
+      {"k1", "v1"}, {"k2", "v2"}};
+
+  for (const auto& key : jsonMap) {
+    EXPECT_EQ(object.at(key.first), jsonMap.at(key.first));
+  }
+}
+
 TEST_F(JsonFunctionsTest, isJsonScalarSignatures) {
   auto signatures = getSignatureStrings("is_json_scalar");
   ASSERT_EQ(1, signatures.size());