Add json_format Presto function (facebookincubator#3525)

Summary:
Pull Request resolved: facebookincubator#3525

Pull Request resolved: facebookincubator#3398
Add 'json_format' udf to prestosql.
Doc: https://prestodb.io/docs/current/functions/json.html?highlight=json_format#json_format

Differential Revision: D42060671

fbshipit-source-id: 6b652de2fda6abff7bbde3465307de5fdf311787

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
 ============

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

velox/functions/prestosql/JsonFunctions.cpp

@@ -0,0 +1,60 @@
+/*
+ * 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<VectorPtr>& args,
+      const TypePtr& /* outputType */,
+      exec::EvalCtx& context,
+      VectorPtr& result) const override {
+    assert(args.size() == 1);
+    auto input = args[0]->asFlatVector<StringView>();
+    auto stringBuffers = input->stringBuffers();
+    auto localResult = std::make_shared<FlatVector<StringView>>(
+        context.pool(),
+        VARCHAR(),
+        input->nulls(),
+        rows.end(),
+        input->values(),
+        std::move(stringBuffers));
+
+    context.moveOrCopyResult(localResult, rows, result);
+  }
+
+  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
+    // json -> varchar
+    return {exec::FunctionSignatureBuilder()
+                .returnType("Varchar")
+                .argumentType("Json")
+                .build()};
+  }
+};
+
+} // namespace
+
+VELOX_DECLARE_VECTOR_FUNCTION(
+    udf_json_format,
+    JsonFormatFunction::signatures(),
+    std::make_unique<JsonFormatFunction>());
+} // namespace facebook::velox::functions

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"});
+  VELOX_REGISTER_VECTOR_FUNCTION(udf_json_format, "json_format");
 }
 
 } // namespace facebook::velox::functions

velox/functions/prestosql/tests/JsonFunctionsTest.cpp

@@ -58,6 +58,37 @@ class JsonFunctionsTest : public functions::test::FunctionBaseTest {
   }
 };
 
+TEST_F(JsonFunctionsTest, JsonFormat) {
+  const auto jsonFormat = [&](std::optional<std::string> value) {
+    return evaluateOnce<std::string, std::string>(
+        "json_format(c0)", {value}, {JSON()});
+  };
+
+  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\"}");
+
+  // check keys and values are there
+  const std::string jsonStr = jsonFormat(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());