feat(json): Add support for field_names_in_json_cast_enabled when casting rows to json (facebookincubator#13108)

Summary:

This diff adds support for field_names_in_json_cast_enabled session property which when enabled produces a json object with the field names having same key as the child name in the row. 

For e.g: 

```
presto:gen_ai> set session field_names_in_json_cast_enabled=true;
SET SESSION           
presto:gen_ai> select cast(cast(row(1, array[1,2], 3.0, row('a', 5)) AS row(x int, y array<int>, z double, r row(a varchar, b int))) as json);
                     _col0                     
-----------------------------------------------
 {"r":{"a":"a","b":5},"x":1,"y":[1,2],"z":3.0} 
```

Velox will produce canonicalized json even though currently Presto java doesnt.

Differential Revision: D73468351

Author: Krishna Pai

velox/core/QueryConfig.h

@@ -542,6 +542,11 @@ class QueryConfig {
   static constexpr const char* kStreamingAggregationEagerFlush =
       "streaming_aggregation_eager_flush";
 
+  /// If this is true, then it allows you to get the struct field names
+  /// as json element names when casting a row to json.
+  static constexpr const char* kFieldNamesInJsonCastEnabled =
+      "field_names_in_json_cast_enabled";
+
   bool selectiveNimbleReaderEnabled() const {
     return get<bool>(kSelectiveNimbleReaderEnabled, false);
   }
@@ -991,6 +996,10 @@ class QueryConfig {
     return get<bool>(kStreamingAggregationEagerFlush, false);
   }
 
+  bool isFieldNamesInJsonCastEnabled() const {
+    return get<bool>(kFieldNamesInJsonCastEnabled, false);
+  }
+
   template <typename T>
   T get(const std::string& key, const T& defaultValue) const {
     return config_->get<T>(key, defaultValue);

velox/functions/prestosql/tests/JsonCastTest.cpp

@@ -82,7 +82,9 @@ class JsonCastTest : public functions::test::CastBaseTest {
     auto thirdChild =
         makeNullableFlatVector<TChild3>(child3, fromType->childAt(2));
 
-    auto rowVector = makeRowVector({firstChild, secondChild, thirdChild});
+    auto names = fromType->asRow().names();
+    auto rowVector =
+        makeRowVector(names, {firstChild, secondChild, thirdChild});
     auto expectedVector =
         makeNullableFlatVector<JsonNativeType>(expected, JSON());
 
@@ -215,6 +217,12 @@ class JsonCastTest : public functions::test::CastBaseTest {
 
     return vector;
   }
+
+  void setFieldNamesInJsonCast(bool flag) {
+    queryCtx_->testingOverrideConfigUnsafe({
+        {core::QueryConfig::kFieldNamesInJsonCastEnabled, std::to_string(flag)},
+    });
+  }
 };
 
 TEST_F(JsonCastTest, fromInteger) {
@@ -734,6 +742,84 @@ TEST_F(JsonCastTest, fromRow) {
   testCast(allNullRow, allNullExpected);
 }
 
+TEST_F(JsonCastTest, fieldNamesInJsonCast) {
+  setFieldNamesInJsonCast(true);
+
+  std::vector<std::optional<int64_t>> child1{
+      std::nullopt, 2, 3, std::nullopt, 5};
+  std::vector<std::optional<StringView>> child2{
+      "red"_sv, std::nullopt, "blue"_sv, std::nullopt, "yellow"_sv};
+  std::vector<std::optional<double>> child3{
+      1.1, 2.2, std::nullopt, std::nullopt, 5.5};
+  std::vector<std::optional<JsonNativeType>> expected{
+      R"({"a":null,"b":"red","c":1.1})",
+      R"({"a":2,"b":null,"c":2.2})",
+      R"({"a":3,"b":"blue","c":null})",
+      R"({"a":null,"b":null,"c":null})",
+      R"({"a":5,"b":"yellow","c":5.5})"};
+
+  testCastFromRow<int64_t, StringView, double>(
+      ROW({"a", "b", "c"}, {BIGINT(), VARCHAR(), DOUBLE()}),
+      child1,
+      child2,
+      child3,
+      expected);
+
+  // Tests rows with child rows, and make sure json's are canonicalized.
+  auto child1_1 = makeNullableFlatVector<int64_t>({3, 1, 2});
+  auto child1_2 = makeArrayVectorFromJson<int64_t>({
+      "[1, 2, 3]",
+      "[4, 5]",
+      "[6, 7, 8]",
+  });
+
+  auto child1_3 = makeRowVector(
+      {"b", "a"},
+      {makeNullableFlatVector<int64_t>({5, 4, 3}),
+       makeNullableFlatVector<int64_t>({1, 2, 3})});
+
+  auto rowVector =
+      makeRowVector({"xyz", "abc", "mno"}, {child1_1, child1_2, child1_3});
+
+  // Canonicalized json's.
+  auto expectedVector = makeNullableFlatVector<JsonNativeType>(
+      {
+          R"({"abc":[1,2,3],"mno":{"a":1,"b":5},"xyz":3})",
+          R"({"abc":[4,5],"mno":{"a":2,"b":4},"xyz":1})",
+          R"({"abc":[6,7,8],"mno":{"a":3,"b":3},"xyz":2})",
+      },
+      JSON());
+
+  testCast(rowVector, expectedVector);
+
+  // Ensure Rows containing maps are also canonicalized.
+
+  auto child2_1 = makeNullableFlatVector<int64_t>({3, std::nullopt, 2});
+  auto child2_2 = makeMapVector<std::string, int64_t>(
+      {{{"x", 2}, {"a", 4}}, {{"y", 6}}, {{"z", 8}, {"A", 10}}});
+
+  auto child2_3 = makeRowVector(
+      {"b", "a"},
+      {makeNullableFlatVector<int64_t>({5, 4, 3}),
+       makeNullableFlatVector<int64_t>({1, 2, std::nullopt})});
+
+  auto rowVector2 =
+      makeRowVector({"xyz", "abc", "mno"}, {child2_1, child2_2, child2_3});
+
+  // Canonicalized json's.
+  auto expectedVector2 = makeNullableFlatVector<JsonNativeType>(
+      {
+          R"({"abc":{"a":4,"x":2},"mno":{"a":1,"b":5},"xyz":3})",
+          R"({"abc":{"y":6},"mno":{"a":2,"b":4},"xyz":null})",
+          R"({"abc":{"A":10,"z":8},"mno":{"a":null,"b":3},"xyz":2})",
+      },
+      JSON());
+
+  testCast(rowVector2, expectedVector2);
+
+  setFieldNamesInJsonCast(false);
+}
+
 TEST_F(JsonCastTest, fromNested) {
   // Create map of array vector.
   auto keyVector = makeNullableFlatVector<StringView>(

velox/functions/prestosql/types/JsonCastOperator.cpp

@@ -237,8 +237,9 @@ struct AsJson {
       const SelectivityVector& rows,
       const BufferPtr& elementToTopLevelRows,
       const std::shared_ptr<exec::CastHooks>& hooks,
-      bool isMapKey = false)
-      : decoded_(context) {
+      bool isMapKey = false,
+      std::optional<std::string> fieldName = std::nullopt)
+      : decoded_(context), fieldName_(std::move(fieldName)) {
     VELOX_CHECK(rows.hasSelections());
 
     exec::EvalErrorsPtr oldErrors;
@@ -294,12 +295,20 @@ struct AsJson {
       // Null values are inlined as "null".
       return 4;
     } else {
+      // If we have field names, we need to add the field name to the length.
+      if (fieldName_.has_value()) {
+        return fieldName_->size() + 3 + this->at(i).size();
+      }
       return this->at(i).size();
     }
   }
 
   // Appends the json string of the value at i to a string writer.
   void append(vector_size_t i, exec::StringWriter& proxy) const {
+    if (fieldName_.has_value()) {
+      proxy.append(fmt::format("\"{}\":", fieldName_.value()));
+    }
+
     if (decoded_->isNullAt(i)) {
       proxy.append("null");
     } else {
@@ -353,6 +362,7 @@ struct AsJson {
   exec::LocalDecodedVector decoded_;
   VectorPtr json_;
   const SimpleVector<StringView>* jsonStrings_;
+  std::optional<std::string> fieldName_;
 };
 
 void castToJsonFromArray(
@@ -537,26 +547,44 @@ void castToJsonFromRow(
     const SelectivityVector& rows,
     FlatVector<StringView>& flatResult,
     const std::shared_ptr<exec::CastHooks>& hooks) {
+  using NameJsonPair = std::pair<std::string, AsJson>;
   // input is guaranteed to be in flat encoding when passed in.
   VELOX_CHECK_EQ(input.encoding(), VectorEncoding::Simple::ROW);
   auto inputRow = input.as<RowVector>();
   auto childrenSize = inputRow->childrenSize();
+  auto fieldNamesInJsonCastEnabled = context.execCtx()
+                                         ->queryCtx()
+                                         ->queryConfig()
+                                         .isFieldNamesInJsonCastEnabled();
 
   // Estimates an upperbound of the total length of all Json strings for the
   // input according to the length of all children Json strings and the
   // delimiters to be added.
   size_t childrenStringSize = 0;
-  std::vector<AsJson> childrenAsJson;
+  std::vector<NameJsonPair> jsonChildren;
+
   for (int i = 0; i < childrenSize; ++i) {
-    childrenAsJson.emplace_back(
-        context, inputRow->childAt(i), rows, nullptr, hooks);
+    auto name = inputRow->type()->asRow().nameOf(i);
+    std::optional<std::string> fieldName =
+        fieldNamesInJsonCastEnabled ? std::optional{name} : std::nullopt;
+
+    jsonChildren.emplace_back(
+        name,
+        AsJson{
+            context,
+            inputRow->childAt(i),
+            rows,
+            nullptr,
+            hooks,
+            false,
+            std::move(fieldName)});
 
     context.applyToSelectedNoThrow(rows, [&](auto row) {
       if (inputRow->isNullAt(row)) {
         // "null" will be inlined in the StringView.
         return;
       }
-      childrenStringSize += childrenAsJson[i].lengthAt(row);
+      childrenStringSize += jsonChildren[i].second.lengthAt(row);
     });
   }
 
@@ -565,6 +593,19 @@ void castToJsonFromRow(
       rows.countSelected() * (childrenSize > 0 ? childrenSize + 1 : 2);
   flatResult.getBufferWithSpace(childrenStringSize);
 
+  // Make sure to sort the children based on their field names if
+  // fieldNamesInJsonCastEnabled is true. This is to make sure the output is
+  // canoncialized.
+
+  if (fieldNamesInJsonCastEnabled) {
+    std::sort(
+        jsonChildren.begin(),
+        jsonChildren.end(),
+        [](const NameJsonPair& a, const NameJsonPair& b) {
+          return a.first < b.first;
+        });
+  }
+
   // Constructs Json string of each row from Json strings of its children.
   context.applyToSelectedNoThrow(rows, [&](auto row) {
     if (inputRow->isNullAt(row)) {
@@ -574,14 +615,24 @@ void castToJsonFromRow(
 
     auto proxy = exec::StringWriter(&flatResult, row);
 
-    proxy.append("["_sv);
-    for (int i = 0; i < childrenSize; ++i) {
+    if (fieldNamesInJsonCastEnabled) {
+      proxy.append("{"_sv);
+    } else {
+      proxy.append("["_sv);
+    }
+
+    for (int i = 0; i < jsonChildren.size(); ++i) {
       if (i > 0) {
         proxy.append(","_sv);
       }
-      childrenAsJson[i].append(row, proxy);
+      jsonChildren[i].second.append(row, proxy);
+    }
+
+    if (fieldNamesInJsonCastEnabled) {
+      proxy.append("}"_sv);
+    } else {
+      proxy.append("]"_sv);
     }
-    proxy.append("]"_sv);
 
     proxy.finalize();
   });