Replace in test_commons

Author: Reranko05

cpp/src/arrow/CMakeLists.txt

@@ -732,8 +732,8 @@ else()
 endif()
 
 set(ARROW_TESTING_SHARED_LINK_LIBS arrow_shared ${ARROW_GTEST_GTEST})
-set(ARROW_TESTING_SHARED_PRIVATE_LINK_LIBS arrow::flatbuffers RapidJSON arrow::simdjson)
-set(ARROW_TESTING_STATIC_LINK_LIBS arrow::flatbuffers RapidJSON arrow::simdjson arrow_static
+set(ARROW_TESTING_SHARED_PRIVATE_LINK_LIBS arrow::flatbuffers arrow::simdjson)
+set(ARROW_TESTING_STATIC_LINK_LIBS arrow::flatbuffers arrow::simdjson arrow_static
                                    ${ARROW_GTEST_GTEST})
 if(ARROW_ENABLE_THREADING)
   list(APPEND ARROW_TESTING_SHARED_PRIVATE_LINK_LIBS arrow::Boost::process)

cpp/src/arrow/json/parser_test.cc

@@ -321,5 +321,17 @@ TEST(BlockParser, AdHoc) {
        R"([{"c":true, "d": "1991-02-03"}, {"c":false, "d":"2019-04-01"}])"});
 }
 
+TEST(JsonTest, PrettyPrintEscapesObjectKeys) {
+  const std::string input = R"({"a\"b":1,"a\\b":2})";
+
+  const std::string expected =
+      "{\n"
+      "    \"a\\\"b\": 1,\n"
+      "    \"a\\\\b\": 2\n"
+      "}";
+
+  EXPECT_EQ(PrettyPrint(input), expected);
+}
+
 }  // namespace json
 }  // namespace arrow

cpp/src/arrow/json/reader_test.cc

@@ -24,6 +24,7 @@
 
 #include "arrow/io/interfaces.h"
 #include "arrow/io/slow.h"
+#include "arrow/json/json_writer_internal.h"
 #include "arrow/json/options.h"
 #include "arrow/json/reader.h"
 #include "arrow/json/test_common.h"
@@ -550,10 +551,14 @@ class StreamingReaderTestBase {
     auto options = GenerateOptions::Defaults();
     options.null_probability = 0;
     for (int i = 0; i < num_rows; ++i) {
-      StringBuffer string_buffer;
-      Writer writer(string_buffer);
+      Writer writer;
       ABORT_NOT_OK(Generate(data_fields, engine, &writer, options));
-      std::string json = string_buffer.GetString();
+
+      auto json_result = writer.GetString();
+      ABORT_NOT_OK(json_result.status());
+      auto json_view = std::move(json_result).ValueOrDie();
+      std::string json(json_view);
+
       rows[i] = Join({"{\"i\":", std::to_string(i), ",\"d\":", json, "}\n"});
       max_row_size = std::max(max_row_size, rows[i].size());
     }

cpp/src/arrow/json/test_common.h

@@ -25,35 +25,31 @@
 #include <utility>
 #include <vector>
 
+#include <simdjson.h>
+
 #include "arrow/array.h"
 #include "arrow/array/builder_binary.h"
 #include "arrow/io/memory.h"
 #include "arrow/json/converter.h"
+#include "arrow/json/json_writer_internal.h"
 #include "arrow/json/options.h"
 #include "arrow/json/parser.h"
-#include "arrow/json/rapidjson_defs.h"
+#include "arrow/result.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/random.h"
 #include "arrow/type.h"
 #include "arrow/util/checked_cast.h"
+#include "arrow/util/simdjson_internal.h"
 #include "arrow/visit_type_inline.h"
 
-#include "rapidjson/document.h"
-#include "rapidjson/prettywriter.h"
-#include "rapidjson/reader.h"
-#include "rapidjson/writer.h"
-
 namespace arrow {
 
 using internal::checked_cast;
 
 namespace json {
 
-namespace rj = arrow::rapidjson;
-
-using rj::StringBuffer;
 using std::string_view;
-using Writer = rj::Writer<StringBuffer>;
+using Writer = JsonWriter;
 
 struct GenerateOptions {
   // Probability of a field being written
@@ -87,35 +83,43 @@ inline static Status Generate(
 
 template <typename Engine>
 struct GenerateImpl {
-  Status Visit(const NullType&) { return OK(writer.Null()); }
+  Status Visit(const NullType&) {
+    writer.Null();
+    return Status::OK();
+  }
 
   Status Visit(const BooleanType&) {
-    return OK(writer.Bool(std::uniform_int_distribution<uint16_t>{}(e) & 1));
+    writer.Bool(std::uniform_int_distribution<uint16_t>{}(e) & 1);
+    return Status::OK();
   }
 
   template <typename T>
   enable_if_physical_unsigned_integer<T, Status> Visit(const T&) {
     auto val = std::uniform_int_distribution<>{}(e);
-    return OK(writer.Uint64(static_cast<typename T::c_type>(val)));
+    writer.Uint64(static_cast<typename T::c_type>(val));
+    return Status::OK();
   }
 
   template <typename T>
   enable_if_physical_signed_integer<T, Status> Visit(const T&) {
     auto val = std::uniform_int_distribution<>{}(e);
-    return OK(writer.Int64(static_cast<typename T::c_type>(val)));
+    writer.Int64(static_cast<typename T::c_type>(val));
+    return Status::OK();
   }
 
   template <typename T>
   enable_if_physical_floating_point<T, Status> Visit(const T&) {
     auto val = std::normal_distribution<typename T::c_type>{0, 1 << 10}(e);
-    return OK(writer.Double(val));
+    writer.Double(val);
+    return Status::OK();
   }
 
   Status GenerateUtf8(const DataType&) {
     auto num_codepoints = std::poisson_distribution<>{4}(e);
     auto seed = std::uniform_int_distribution<uint32_t>{}(e);
     std::string s = RandomUtf8String(seed, num_codepoints);
-    return OK(writer.String(s));
+    writer.String(s);
+    return Status::OK();
   }
 
   template <typename T>
@@ -132,7 +136,8 @@ struct GenerateImpl {
     for (int i = 0; i < size; ++i) {
       RETURN_NOT_OK(Generate(t.value_type(), e, &writer, options));
     }
-    return OK(writer.EndArray(size));
+    writer.EndArray();
+    return Status::OK();
   }
 
   Status Visit(const ListViewType& t) { return NotImplemented(t); }
@@ -162,7 +167,7 @@ struct GenerateImpl {
   }
 
   Engine& e;
-  rj::Writer<rj::StringBuffer>& writer;
+  Writer& writer;
   const GenerateOptions& options;
 };
 
@@ -180,12 +185,9 @@ inline static Status Generate(const std::shared_ptr<DataType>& type, Engine& e,
 template <typename Engine>
 inline static Status Generate(const std::vector<std::shared_ptr<Field>>& fields,
                               Engine& e, Writer* writer, const GenerateOptions& options) {
-  RETURN_NOT_OK(OK(writer->StartObject()));
-
-  int num_fields = 0;
+  writer->StartObject();
   auto write_field = [&](const Field& f) {
-    ++num_fields;
-    writer->Key(f.name().c_str());
+    writer->Key(f.name());
     return Generate(f.type(), e, writer, options);
   };
 
@@ -210,7 +212,8 @@ inline static Status Generate(const std::vector<std::shared_ptr<Field>>& fields,
     }
   }
 
-  return OK(writer->EndObject(num_fields));
+  writer->EndObject();
+  return Status::OK();
 }
 
 inline static Status MakeStream(string_view src_str,
@@ -258,14 +261,26 @@ inline static Status ParseFromString(ParseOptions options, string_view src_str,
 }
 
 static inline std::string PrettyPrint(string_view one_line) {
-  rj::Document document;
+  simdjson::ondemand::parser parser;
 
   // Must pass size to avoid ASAN issues.
-  document.Parse(one_line.data(), one_line.size());
-  rj::StringBuffer sb;
-  rj::PrettyWriter<rj::StringBuffer> writer(sb);
-  document.Accept(writer);
-  return sb.GetString();
+  simdjson::padded_string json(one_line.data(), one_line.size());
+
+  auto document_result =
+      internal::ResolveSimdjsonResult(parser.iterate(json), "Failed to parse JSON");
+  ABORT_NOT_OK(document_result.status());
+  auto document = std::move(document_result).ValueOrDie();
+
+  auto value_result =
+      internal::ResolveSimdjsonResult(document.get_value(), "Failed to get JSON value");
+  ABORT_NOT_OK(value_result.status());
+  auto value = std::move(value_result).ValueOrDie();
+
+  std::string result;
+  result.reserve(one_line.size());
+
+  ABORT_NOT_OK(internal::PrettyPrintJsonValue(value, &result));
+  return result;
 }
 
 template <typename T>

cpp/src/arrow/util/simdjson_internal.h

@@ -27,6 +27,7 @@
 
 #include <simdjson.h>
 
+#include "arrow/json/json_writer_internal.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
 
@@ -294,6 +295,118 @@ Status VisitJsonValue(simdjson::ondemand::value value, ObjectFn&& object_fn,
   return Status::Invalid("Unreachable");
 }
 
+inline Status PrettyPrintJsonValue(simdjson::ondemand::value value, std::string* out,
+                                   int indent = 0) {
+  constexpr int kIndentSize = 4;
+
+  auto append_indent = [&](int level) {
+    out->append(static_cast<size_t>(level * kIndentSize), ' ');
+  };
+
+  ARROW_ASSIGN_OR_RAISE(
+      auto type, ResolveSimdjsonResult(value.type(), "Failed to determine JSON type"));
+
+  switch (type) {
+    case simdjson::ondemand::json_type::object: {
+      ARROW_ASSIGN_OR_RAISE(
+          auto object,
+          ResolveSimdjsonResult(value.get_object(), "Failed to get JSON object"));
+
+      out->append("{");
+
+      bool first = true;
+      for (auto field_result : object) {
+        ARROW_ASSIGN_OR_RAISE(
+            auto field,
+            ResolveSimdjsonResult(field_result, "Failed to iterate JSON object"));
+
+        ARROW_ASSIGN_OR_RAISE(auto key,
+                              ResolveSimdjsonResult(field.unescaped_key(),
+                                                    "Failed to get JSON object key"));
+
+        auto field_value = field.value();
+
+        if (first) {
+          out->append("\n");
+          first = false;
+        } else {
+          out->append(",\n");
+        }
+
+        append_indent(indent + 1);
+
+        json::JsonWriter writer;
+        writer.String(key);
+
+        ARROW_ASSIGN_OR_RAISE(auto escaped_key, writer.GetString());
+        out->append(escaped_key);
+        out->append(": ");
+
+        RETURN_NOT_OK(PrettyPrintJsonValue(field_value, out, indent + 1));
+      }
+
+      if (!first) {
+        out->append("\n");
+        append_indent(indent);
+      }
+
+      out->append("}");
+      return Status::OK();
+    }
+
+    case simdjson::ondemand::json_type::array: {
+      ARROW_ASSIGN_OR_RAISE(
+          auto array,
+          ResolveSimdjsonResult(value.get_array(), "Failed to get JSON array"));
+
+      out->append("[");
+
+      bool first = true;
+      for (auto element_result : array) {
+        ARROW_ASSIGN_OR_RAISE(
+            auto element,
+            ResolveSimdjsonResult(element_result, "Failed to iterate JSON array"));
+
+        if (first) {
+          out->append("\n");
+          first = false;
+        } else {
+          out->append(",\n");
+        }
+
+        append_indent(indent + 1);
+
+        RETURN_NOT_OK(PrettyPrintJsonValue(element, out, indent + 1));
+      }
+
+      if (!first) {
+        out->append("\n");
+        append_indent(indent);
+      }
+
+      out->append("]");
+      return Status::OK();
+    }
+
+    case simdjson::ondemand::json_type::string:
+    case simdjson::ondemand::json_type::boolean:
+    case simdjson::ondemand::json_type::null:
+    case simdjson::ondemand::json_type::number: {
+      ARROW_ASSIGN_OR_RAISE(auto serialized,
+                            ResolveSimdjsonResult(simdjson::to_json_string(value),
+                                                  "Failed to serialize JSON value"));
+
+      out->append(serialized);
+      return Status::OK();
+    }
+
+    case simdjson::ondemand::json_type::unknown:
+      return Status::Invalid("Unknown JSON type");
+  }
+
+  return Status::Invalid("Unreachable");
+}
+
 inline const char* JsonTypeName(simdjson::ondemand::json_type type) {
   switch (type) {
     case simdjson::ondemand::json_type::array: