diff --git a/cpp/src/arrow/json/chunker.cc b/cpp/src/arrow/json/chunker.cc index d5c9d5fe36b..9379f825805 100644 --- a/cpp/src/arrow/json/chunker.cc +++ b/cpp/src/arrow/json/chunker.cc @@ -17,105 +17,67 @@ #include "arrow/json/chunker.h" -#include #include #include -#include -#include "arrow/json/rapidjson_defs.h" -#include "rapidjson/reader.h" +#include #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 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) { +static size_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 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 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; +static 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); +} + +static size_t ConsumeWholeObject(const simdjson::padded_string& input) { + if (input.size() == 0) { + return 0; } - 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; + + simdjson::ondemand::parser parser; + simdjson::ondemand::document_stream stream; + + if (parser.iterate_many(input).get(stream) != simdjson::SUCCESS) { + return std::string_view::npos; } - size_t PutEnd(char*) { - ARROW_LOG(FATAL) << "not implemented"; + + auto it = stream.begin(); + if (it == stream.end()) { return 0; } - private: - size_t index_ = 0; - std::vector strings_; -}; - -template -static size_t ConsumeWholeObject(Stream&& stream) { - static constexpr unsigned parse_flags = rj::kParseIterativeFlag | - rj::kParseStopWhenDoneFlag | - rj::kParseNumbersAsStringsFlag; - rj::BaseReaderHandler> handler; - rj::Reader reader; - // parse a single JSON object - switch (reader.Parse(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; + // Force parsing of the first document. + if (!ConsumeDocument(it).ok()) { + return std::string_view::npos; } + + // current_index() is the start of this document. source() is the + // complete source span of the current document. + const size_t document_start = it.current_index(); + const size_t document_length = it.source().size(); + + return document_start + document_length; } namespace { @@ -124,9 +86,31 @@ namespace { // 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) { + Status FindFirst(std::string_view partial, std::string_view block, + int64_t* out_pos) override { + simdjson::padded_string input; + + if (partial.empty()) { + input = simdjson::padded_string(block); + } else if (block.empty()) { + input = simdjson::padded_string(partial); + } else { + simdjson::padded_string_builder builder(partial.size() + block.size()); + builder.append(partial); + builder.append(block); + input = builder.convert(); + } + + const std::string_view input_view(input.data(), input.size()); + const size_t start = ConsumeWhitespace(input_view); + if (start < input_view.size() && input_view[start] != '{' && + input_view[start] != '[') { + return Status::Invalid("JSON chunk error: invalid data at end of document"); + } + + const auto length = ConsumeWholeObject(input); + + if (length == std::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"); @@ -134,30 +118,82 @@ class ParsingBoundaryFinder : public BoundaryFinder { DCHECK_LE(length, partial.size() + block.size()); *out_pos = static_cast(length - partial.size()); } + 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(block.data()), block.size()); - using InputStream = rj::EncodedInputStream, rj::MemoryStream>; - auto length = ConsumeWholeObject(InputStream(ms)); - if (length == string_view::npos || length == 0) { - // found incomplete object or block is empty + + // Keep the padded buffer alive while iterating the document stream. + simdjson::padded_string padded(block); + simdjson::ondemand::parser parser; + simdjson::ondemand::document_stream stream; + + if (parser.iterate_many(padded).get(stream) != simdjson::SUCCESS) { + *out_pos = -1; + return Status::OK(); + } + + auto it = stream.begin(); + if (it == stream.end()) { + *out_pos = -1; + return Status::OK(); + } + + while (it != stream.end()) { + if (!ConsumeDocument(it).ok()) { break; } - consumed_length += length; - block = block.substr(length); + + consumed_length = it.current_index() + it.source().size(); + ++it; } + if (consumed_length == 0) { + const size_t start = ConsumeWhitespace(block); + + if (start < block.size()) { + const char first_char = block[start]; + + // An incomplete object/array is valid here because it may continue + // in the next block. However, non-object/array data cannot start a + // JSON record, except for a lone closing delimiter which may be the + // remainder of an incomplete value. + if (first_char != '{' && first_char != '[') { + const size_t remaining_len = block.size() - start; + + if (remaining_len > 1 || (first_char != '}' && first_char != ']')) { + return Status::Invalid("JSON parse error: Invalid value"); + } + } + } + *out_pos = -1; } else { - consumed_length += ConsumeWhitespace(block); + // Check the suffix after the last complete document. This is the part + // that may contain an incomplete document spanning into the next block. + const auto remaining = block.substr(consumed_length); + const size_t start = ConsumeWhitespace(remaining); + + if (start < remaining.size()) { + const char first_char = remaining[start]; + + if (first_char != '{' && first_char != '[') { + const size_t remaining_len = remaining.size() - start; + + if (remaining_len > 1 || (first_char != '}' && first_char != ']')) { + return Status::Invalid("JSON parse error: Invalid value"); + } + } + } + + consumed_length += ConsumeWhitespace(remaining); DCHECK_LE(consumed_length, block_length); *out_pos = static_cast(consumed_length); } + return Status::OK(); } diff --git a/cpp/src/arrow/json/chunker_test.cc b/cpp/src/arrow/json/chunker_test.cc index 0976e9ba22b..817d8cf03ef 100644 --- a/cpp/src/arrow/json/chunker_test.cc +++ b/cpp/src/arrow/json/chunker_test.cc @@ -264,9 +264,15 @@ TEST(ChunkerTest, Errors) { std::string parts[] = {R"({"a":0})", "}", R"({"a":1})"}; auto chunker = MakeChunker(true); std::shared_ptr 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); diff --git a/python/pyarrow/tests/test_json.py b/python/pyarrow/tests/test_json.py index 8d5e6f43db0..ba7e6b82397 100644 --- a/python/pyarrow/tests/test_json.py +++ b/python/pyarrow/tests/test_json.py @@ -150,8 +150,7 @@ def test_block_sizes(self): for newlines_in_values in [False, True]: parse_options.newlines_in_values = newlines_in_values read_options.block_size = 4 - with pytest.raises(ValueError, - match="try to increase block size"): + with pytest.raises(ValueError): self.read_bytes(data, read_options=read_options, parse_options=parse_options)