From e2802d7dd78019c8578f1fd8ab814f98864a3834 Mon Sep 17 00:00:00 2001
From: Arpit Porwal <arp8@meta.com>
Date: Mon, 12 Dec 2022 18:20:16 -0800
Subject: [PATCH] Add json_format Presto function (#3398)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/3398

Add 'json_format' udf to prestosql.

Doc: https://prestodb.io/docs/current/functions/json.html?highlight=json_format#json_format

Reviewed By: kagamiori

Differential Revision: D41643635

fbshipit-source-id: f2cc011a4b035b1fa20a09b1fa380cc91cc753f1
---
 velox/docs/functions/json.rst                 |  7 +++++
 velox/functions/prestosql/JsonFunctions.h     | 12 +++++++
 .../JsonFunctionsRegistration.cpp             |  1 +
 .../prestosql/tests/JsonFunctionsTest.cpp     | 31 +++++++++++++++++++
 4 files changed, 51 insertions(+)

diff --git a/velox/docs/functions/json.rst b/velox/docs/functions/json.rst
index 8ebd16525699..882927620ef1 100644
--- a/velox/docs/functions/json.rst
+++ b/velox/docs/functions/json.rst
@@ -133,6 +133,13 @@ JSON Functions
         SELECT json_size('{"x": [1, 2, 3]}', '$.x'); -- 3
         SELECT json_size('{"x": {"a": 1, "b": 2}}', '$.x.a'); -- 0
 
+.. 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
index 6645e783226e..1285ff9f3ec6 100644
--- 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 void call(
+      out_type<Varchar>& jsonString,
+      const arg_type<Json>& json) {
+    folly::parseJson(json.getString());
+    jsonString.setNoCopy(json);
+  }
+};
+
 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
index 4f559a79717c..68a15589e69c 100644
--- a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
+++ b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
@@ -35,6 +35,7 @@ void registerJsonFunctions() {
   registerFunction<JsonArrayContainsFunction, bool, Json, Varchar>(
       {"json_array_contains"});
   registerFunction<JsonSizeFunction, int64_t, Json, Varchar>({"json_size"});
+  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
index 313ee1515b16..658cc095287b 100644
--- a/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
+++ b/velox/functions/prestosql/tests/JsonFunctionsTest.cpp
@@ -58,6 +58,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());