Add json_format Presto function (#3398)

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

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,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<VectorPtr>& 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<ConstantVector<StringView>>()->valueAt(0);
+      localResult = std::make_shared<ConstantVector<StringView>>(
+          context.pool(), rows.end(), false, VARCHAR(), std::move(value));
+    } else {
+      auto flatInput = arg->asFlatVector<StringView>();
+
+      auto stringBuffers = flatInput->stringBuffers();
+      VELOX_CHECK_LE(rows.end(), flatInput->size());
+      localResult = std::make_shared<FlatVector<StringView>>(
+          context.pool(),
+          VARCHAR(),
+          nullptr,
+          rows.end(),
+          flatInput->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,48 @@ 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(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<StringView>(
+      {"This is a long sentence", "This is some other sentence"})});
+
+  auto result = evaluate("json_format(c0)", data);
+  auto expected = makeFlatVector<StringView>(
+      {"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<StringView>({{"apple", "apple"}});
+
+  facebook::velox::test::assertEqualVectors(expected, result);
+
+  data = makeRowVector(
+      {makeFlatVector<bool>({true, false}),
+       makeFlatVector<StringView>(
+           {"This is a long sentence", "This is some other sentence"})});
+
+  result = evaluate("if(c0, 'foo', json_format(c1))", data);
+  expected = makeFlatVector<StringView>({"foo", "This is some other sentence"});
+  facebook::velox::test::assertEqualVectors(expected, result);
+
+  result = evaluate("if(c0, json_format(c1), 'bar')", data);
+  expected = makeFlatVector<StringView>({"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());