GH-50944: [C++] Replace RapidJSON with simdjson in JSON chunker

cpp/src/arrow/json/chunker.cc

@@ -17,145 +17,72 @@
 
 #include "arrow/json/chunker.h"
 
-#include <algorithm>
+#include <cstring>
 #include <string_view>
 #include <utility>
-#include <vector>
 
-#include "arrow/json/rapidjson_defs.h"
-#include "rapidjson/reader.h"
+#include <simdjson.h>
 
 #include "arrow/buffer.h"
 #include "arrow/json/options.h"
 #include "arrow/util/logging_internal.h"
+#include "arrow/util/simdjson_internal.h"
 
 namespace arrow {
-
-using std::string_view;
-
 namespace json {
+namespace {
 
-namespace rj = arrow::rapidjson;
-
-static size_t ConsumeWhitespace(string_view view) {
-#ifdef RAPIDJSON_SIMD
-  auto data = view.data();
-  auto nonws_begin = rj::SkipWhitespace_SIMD(data, data + view.size());
-  return nonws_begin - data;
-#else
-  auto ws_count = view.find_first_not_of(" \t\r\n");
-  if (ws_count == string_view::npos) {
+// XXX We could try to SIMD-accelerate this routine but it's called only
+// once per chunk and also will presumably examine a minimal amount of bytes.
+int64_t ConsumeWhitespace(std::string_view view) {
+  const auto ws_count = view.find_first_not_of(" \t\r\n");
+  if (ws_count == std::string_view::npos) {
     return view.size();
-  } else {
-    return ws_count;
   }
-#endif
+  return static_cast<int64_t>(ws_count);
 }
 
-/// RapidJson custom stream for reading JSON stored in multiple buffers
-/// http://rapidjson.org/md_doc_stream.html#CustomStream
-class MultiStringStream {
- public:
-  using Ch = char;
-  explicit MultiStringStream(std::vector<string_view> strings)
-      : strings_(std::move(strings)) {
-    std::reverse(strings_.begin(), strings_.end());
-  }
-  explicit MultiStringStream(const BufferVector& buffers) : strings_(buffers.size()) {
-    for (size_t i = 0; i < buffers.size(); ++i) {
-      strings_[i] = string_view(*buffers[i]);
-    }
-    std::reverse(strings_.begin(), strings_.end());
-  }
-  char Peek() const {
-    if (strings_.size() == 0) return '\0';
-    return strings_.back()[0];
-  }
-  char Take() {
-    if (strings_.size() == 0) return '\0';
-    char taken = strings_.back()[0];
-    if (strings_.back().size() == 1) {
-      strings_.pop_back();
-    } else {
-      strings_.back() = strings_.back().substr(1);
-    }
-    ++index_;
-    return taken;
-  }
-  size_t Tell() { return index_; }
-  void Put(char) { ARROW_LOG(FATAL) << "not implemented"; }
-  void Flush() { ARROW_LOG(FATAL) << "not implemented"; }
-  char* PutBegin() {
-    ARROW_LOG(FATAL) << "not implemented";
-    return nullptr;
-  }
-  size_t PutEnd(char*) {
-    ARROW_LOG(FATAL) << "not implemented";
-    return 0;
-  }
-
- private:
-  size_t index_ = 0;
-  std::vector<string_view> strings_;
-};
-
-template <typename Stream>
-static size_t ConsumeWholeObject(Stream&& stream) {
-  static constexpr unsigned parse_flags = rj::kParseIterativeFlag |
-                                          rj::kParseStopWhenDoneFlag |
-                                          rj::kParseNumbersAsStringsFlag;
-  rj::BaseReaderHandler<rj::UTF8<>> handler;
-  rj::Reader reader;
-  // parse a single JSON object
-  switch (reader.Parse<parse_flags>(stream, handler).Code()) {
-    case rj::kParseErrorNone:
-      return stream.Tell();
-    case rj::kParseErrorDocumentEmpty:
-      return 0;
-    default:
-      // rapidjson emitted an error, the most recent object was partial
-      return string_view::npos;
-  }
+Status ConsumeDocument(simdjson::ondemand::document_stream::iterator& it) {
+  ARROW_ASSIGN_OR_RAISE(
+      auto document, internal::ResolveSimdjsonResult(*it, "Failed to get JSON document"));
+  ARROW_ASSIGN_OR_RAISE(
+      auto value,
+      internal::ResolveSimdjsonResult(document.get_value(), "Failed to get JSON value"));
+  return internal::ConsumeJsonValue(value);
 }
 
-namespace {
-
 // A BoundaryFinder implementation that assumes JSON objects can contain raw newlines,
 // and uses actual JSON parsing to delimit them.
 class ParsingBoundaryFinder : public BoundaryFinder {
  public:
-  Status FindFirst(string_view partial, string_view block, int64_t* out_pos) override {
-    auto length = ConsumeWholeObject(MultiStringStream({partial, block}));
-    if (length == string_view::npos) {
-      *out_pos = -1;
-    } else if (ARROW_PREDICT_FALSE(length < partial.size())) {
-      return Status::Invalid("JSON chunk error: invalid data at end of document");
+  explicit ParsingBoundaryFinder(MemoryPool* pool) : pool_(pool) {}
+
+  Status FindFirst(std::string_view partial, std::string_view block,
+                   int64_t* out_pos) override {
+    ARROW_ASSIGN_OR_RAISE(auto input, GetPaddedStringView(partial, block));
+    ARROW_ASSIGN_OR_RAISE(auto consumed_length, FindDocument(input, /*find_last=*/false));
+
+    DCHECK_NE(consumed_length, std::string_view::npos);
+    if (consumed_length == 0) {
+      *out_pos = kNoDelimiterFound;
+    } else if (ARROW_PREDICT_FALSE(consumed_length <= partial.size())) {
+      // Something bad happened: partial wasn't supposed to be a full document
+      return Status::Invalid("JSON parse error: invalid data at end of document");
     } else {
-      DCHECK_LE(length, partial.size() + block.size());
-      *out_pos = static_cast<int64_t>(length - partial.size());
+      consumed_length -= partial.size();
+      DCHECK_LE(consumed_length, block.size());
+      *out_pos = static_cast<int64_t>(consumed_length);
     }
     return Status::OK();
   }
 
   Status FindLast(std::string_view block, int64_t* out_pos) override {
-    const size_t block_length = block.size();
-    size_t consumed_length = 0;
-    while (consumed_length < block_length) {
-      rj::MemoryStream ms(reinterpret_cast<const char*>(block.data()), block.size());
-      using InputStream = rj::EncodedInputStream<rj::UTF8<>, rj::MemoryStream>;
-      auto length = ConsumeWholeObject(InputStream(ms));
-      if (length == string_view::npos || length == 0) {
-        // found incomplete object or block is empty
-        break;
-      }
-      consumed_length += length;
-      block = block.substr(length);
-    }
+    ARROW_ASSIGN_OR_RAISE(auto input, GetPaddedStringView(/*partial=*/"", block));
+    ARROW_ASSIGN_OR_RAISE(auto consumed_length, FindDocument(input, /*find_last=*/true));
+
     if (consumed_length == 0) {
-      *out_pos = -1;
+      *out_pos = kNoDelimiterFound;
     } else {
-      consumed_length += ConsumeWhitespace(block);
-      DCHECK_LE(consumed_length, block_length);
       *out_pos = static_cast<int64_t>(consumed_length);
     }
     return Status::OK();
@@ -165,14 +92,94 @@ class ParsingBoundaryFinder : public BoundaryFinder {
                  int64_t* out_pos, int64_t* num_found) override {
     return Status::NotImplemented("ParsingBoundaryFinder::FindNth");
   }
+
+ private:
+  MemoryPool* pool_;
+  simdjson::ondemand::parser parser_;
+  // A persistent buffer to keep padded contents for simdjson.
+  // This should be more efficient than allocating a new padded_string everytime.
+  std::shared_ptr<ResizableBuffer> buffer_;
+
+  Result<simdjson::padded_string_view> GetPaddedStringView(std::string_view partial,
+                                                           std::string_view block) {
+    const auto data_size = partial.size() + block.size();
+    const auto required_size = data_size + simdjson::SIMDJSON_PADDING;
+    if (!buffer_) {
+      ARROW_ASSIGN_OR_RAISE(buffer_, AllocateResizableBuffer(
+                                         /*size=*/0, pool_));
+    }
+    // Ensure the buffer has enough space for this data + padding
+    if (buffer_->capacity() < static_cast<int64_t>(required_size)) {
+      RETURN_NOT_OK(buffer_->Reserve(static_cast<int64_t>(required_size * 3 / 2)));
+    }
+    auto data = buffer_->mutable_data();
+    std::memcpy(data, partial.data(), partial.size());
+    std::memcpy(data + partial.size(), block.data(), block.size());
+    auto view = simdjson::padded_string_view(data, /*len=*/data_size,
+                                             /*capacity=*/buffer_->capacity());
+    DCHECK(view.has_sufficient_padding());
+    return view;
+  }
+
+  // Find the first or last JSON object (depending on `find_last`)
+  // and return the consumed JSON byte length, or 0 if no valid document
+  // can be parsed.
+  Result<size_t> FindDocument(simdjson::padded_string_view input, bool find_last) {
+    simdjson::ondemand::document_stream stream;
+    // XXX Should be pass a specific batch_size?
+    // The default value used by simdjson is 1MB, probably enough for most purposes.
+    RETURN_NOT_OK(ToStatus(parser_.iterate_many(input).get(stream)));
+    auto it = stream.begin();
+    if (it == stream.end()) {
+      // Empty input (only whitespace?)
+      return 0;
+    }
+
+    int64_t consumed_length = 0;
+    if (!find_last) {
+      // Parsing the first document only.
+      if (!ConsumeDocument(it).ok()) {
+        // Could be either a partial document or invalid JSON, we'll let
+        // followup chunker or parser calls decide.
+        return 0;
+      }
+      // current_index() is the start of the current document;
+      // source() is the complete source span of the current document.
+      consumed_length = it.current_index() + it.source().size();
+    } else {
+      while (it != stream.end()) {
+        if (!ConsumeDocument(it).ok()) {
+          // Could be either a partial document or invalid JSON, we'll let
+          // followup chunker or parser calls decide.
+          break;
+        }
+        consumed_length = it.current_index() + it.source().size();
+        ++it;
+      }
+    }
+    if (consumed_length > 0) {
+      // If we found at least one document, also consume its trailing whitespace
+      // to avoid stray bytes at the end of the stream.
+      consumed_length += ConsumeWhitespace(input.substr(consumed_length));
+    }
+    return consumed_length;
+  }
+
+  static Status ToStatus(simdjson::error_code error,
+                         std::string_view error_prefix = "JSON parse error: ") {
+    if (error == simdjson::SUCCESS) {
+      return Status::OK();
+    }
+    return Status::Invalid(std::string(error_prefix) + simdjson::error_message(error));
+  }
 };
 
 }  // namespace
 
-std::unique_ptr<Chunker> MakeChunker(const ParseOptions& options) {
+std::unique_ptr<Chunker> MakeChunker(const ParseOptions& options, MemoryPool* pool) {
   std::shared_ptr<BoundaryFinder> delimiter;
   if (options.newlines_in_values) {
-    delimiter = std::make_shared<ParsingBoundaryFinder>();
+    delimiter = std::make_shared<ParsingBoundaryFinder>(pool);
   } else {
     delimiter = MakeNewlineBoundaryFinder();
   }

cpp/src/arrow/json/chunker.h

@@ -19,6 +19,7 @@
 
 #include <memory>
 
+#include "arrow/type_fwd.h"
 #include "arrow/util/delimiting.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/visibility.h"
@@ -29,7 +30,8 @@ namespace json {
 struct ParseOptions;
 
 ARROW_EXPORT
-std::unique_ptr<Chunker> MakeChunker(const ParseOptions& options);
+std::unique_ptr<Chunker> MakeChunker(const ParseOptions& options,
+                                     MemoryPool* pool = default_memory_pool());
 
 }  // namespace json
 }  // namespace arrow

cpp/src/arrow/json/chunker_test.cc

@@ -264,14 +264,20 @@ TEST(ChunkerTest, Errors) {
   std::string parts[] = {R"({"a":0})", "}", R"({"a":1})"};
   auto chunker = MakeChunker(true);
   std::shared_ptr<Buffer> whole, rest, completion;
+
   ASSERT_OK(chunker->Process(Buffer::FromString(parts[0] + parts[1]), &whole, &rest));
-  ASSERT_EQ(std::string_view(*whole), parts[0]);
-  ASSERT_EQ(std::string_view(*rest), parts[1]);
+
+  // simdjson rejects the malformed stream as a whole, so no complete chunk
+  // is emitted before the trailing invalid data.
+  ASSERT_TRUE(whole);
+  ASSERT_EQ(std::string_view(*whole), "");
+  ASSERT_EQ(std::string_view(*rest), parts[0] + parts[1]);
+
   auto status =
       chunker->ProcessWithPartial(rest, Buffer::FromString(parts[2]), &completion, &rest);
   ASSERT_RAISES(Invalid, status);
   EXPECT_THAT(status.message(),
-              ::testing::StartsWith("JSON chunk error: invalid data at end of document"));
+              ::testing::StartsWith("JSON parse error: invalid data at end of document"));
 }
 
 TEST_P(BaseChunkerTest, StraddlingEmpty) {

cpp/src/arrow/json/parser_benchmark.cc

@@ -101,6 +101,38 @@ static void BenchmarkJSONChunking(benchmark::State& state,  // NOLINT non-const
   state.counters["json_size"] = static_cast<double>(json->size());
 }
 
+static void BenchmarkJSONChunkingMultipleBlocks(
+    benchmark::State& state,  // NOLINT non-const reference
+    const std::shared_ptr<Buffer>& json, ParseOptions options, int64_t block_size) {
+  auto chunker = MakeChunker(options);
+
+  for (auto _ : state) {
+    std::shared_ptr<Buffer> partial;
+    int64_t offset = 0;
+    while (offset < json->size()) {
+      const auto slice_size = std::min(block_size, json->size() - offset);
+      auto block = SliceBuffer(json, offset, slice_size);
+      offset += slice_size;
+      if (offset == json->size()) {
+        // Final block
+        ABORT_NOT_OK(chunker->ProcessFinal(partial, block, &block, &partial));
+        ABORT_NOT_OK(chunker->Process(partial, &block, &partial));
+      } else if (partial) {
+        // Continuation block
+        ABORT_NOT_OK(chunker->ProcessWithPartial(partial, block, &block, &partial));
+        ABORT_NOT_OK(chunker->Process(partial, &block, &partial));
+      } else {
+        // First block
+        ABORT_NOT_OK(chunker->Process(block, &block, &partial));
+      }
+    }
+  }
+
+  state.SetBytesProcessed(state.iterations() * json->size());
+  state.counters["json_size"] = static_cast<double>(json->size());
+  state.counters["block_size"] = static_cast<double>(block_size);
+}
+
 static void ChunkJSONPrettyPrinted(
     benchmark::State& state) {  // NOLINT non-const reference
   const int32_t num_rows = 5000;
@@ -113,6 +145,19 @@ static void ChunkJSONPrettyPrinted(
   BenchmarkJSONChunking(state, std::make_shared<Buffer>(json), options);
 }
 
+static void ChunkJSONPrettyPrintedMultipleBlocks(
+    benchmark::State& state) {  // NOLINT non-const reference
+  const int32_t num_rows = 5000;
+
+  auto options = ParseOptions::Defaults();
+  options.newlines_in_values = true;
+  options.explicit_schema = schema(TestFields());
+
+  auto json = GenerateTestData(options.explicit_schema, num_rows, /*pretty=*/true);
+  BenchmarkJSONChunkingMultipleBlocks(state, std::make_shared<Buffer>(json), options,
+                                      /*block_size=*/json.length() / 8);
+}
+
 static void ChunkJSONLineDelimited(
     benchmark::State& state) {  // NOLINT non-const reference
   const int32_t num_rows = 5000;
@@ -223,6 +268,7 @@ static void ParseJSONFields(benchmark::State& state) {  // NOLINT non-const refe
 }
 
 BENCHMARK(ChunkJSONPrettyPrinted);
+BENCHMARK(ChunkJSONPrettyPrintedMultipleBlocks);
 BENCHMARK(ChunkJSONLineDelimited);
 BENCHMARK(ParseJSONBlockWithSchema);