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: e84f71dfc6aae4e3ec7f1f4faa254bd897571df2
facebookincubator · Dec 6, 2022 · eb5033b · eb5033b
1 parent 4f2b996
commit eb5033b
Show file tree

Hide file tree

Showing 4 changed files with 59 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('{"k1":1,"k2":23,"k3":456}');
+
 ============
 JSON Vectors
 ============

diff --git a/velox/functions/prestosql/JsonFunctions.h b/velox/functions/prestosql/JsonFunctions.h
@@ -20,6 +20,17 @@
 
 namespace facebook::velox::functions {
 
+template <typename T>
+struct JsonFormatFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Varchar>& result,
+      const arg_type<Json>& json) {
+    UDFOutputString::assign(result, folly::toJson(folly::parseJson(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
@@ -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
@@ -24,6 +24,10 @@ namespace {
 
 class JsonFunctionsTest : public functions::test::FunctionBaseTest {
  public:
+  std::optional<std::string> json_format(std::optional<std::string> json) {
+    return evaluateOnce<std::string>("json_format(c0)", json);
+  }
+
   std::optional<bool> is_json_scalar(std::optional<std::string> json) {
     return evaluateOnce<bool>("is_json_scalar(c0)", json);
   }
@@ -52,6 +56,42 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest {
   }
 };
 
+TEST_F(JsonFunctionsTest, JsonFormatSignatures) {
+  auto signatures = getSignatureStrings("json_format");
+  ASSERT_EQ(1, signatures.size());
+  ASSERT_EQ(1, signatures.count("(json) -> varchar"));
+}
+
+TEST_F(JsonFunctionsTest, JsonFormat) {
+  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
+  std::optional<std::string> jsonStr = json_format(R"({"k1":"v1","k2":"v1"})");
+  folly::dynamic object = folly::parseJson(jsonStr.value());
+
+  std::set<std::string> keys;
+  keys.insert("k1");
+  keys.insert("k2");
+
+  for (auto begin = object.keys().begin(); begin != object.keys().end();
+       ++begin) {
+    EXPECT_EQ(keys.find(begin->getString()) != keys.end(), true);
+  }
+
+  std::set<std::string> values;
+  values.insert("v1");
+
+  for (auto begin = object.values().begin(); begin != object.values().end();
+       ++begin) {
+    EXPECT_EQ(values.find(begin->getString()) != values.end(), true);
+  }
+}
+
 TEST_F(JsonFunctionsTest, isJsonScalarSignatures) {
   auto signatures = getSignatureStrings("is_json_scalar");
   ASSERT_EQ(1, signatures.size());