diff --git a/cpp/src/arrow/json/chunker.cc b/cpp/src/arrow/json/chunker.cc index d5c9d5fe36b..f69c45b9386 100644 --- a/cpp/src/arrow/json/chunker.cc +++ b/cpp/src/arrow/json/chunker.cc @@ -17,145 +17,72 @@ #include "arrow/json/chunker.h" -#include +#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 { -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(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; - } - 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 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; - } +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(length - partial.size()); + consumed_length -= partial.size(); + DCHECK_LE(consumed_length, block.size()); + *out_pos = static_cast(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(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 - 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(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 buffer_; + + Result 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(required_size)) { + RETURN_NOT_OK(buffer_->Reserve(static_cast(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 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 MakeChunker(const ParseOptions& options) { +std::unique_ptr MakeChunker(const ParseOptions& options, MemoryPool* pool) { std::shared_ptr delimiter; if (options.newlines_in_values) { - delimiter = std::make_shared(); + delimiter = std::make_shared(pool); } else { delimiter = MakeNewlineBoundaryFinder(); } diff --git a/cpp/src/arrow/json/chunker.h b/cpp/src/arrow/json/chunker.h index 9ed85126da1..fc6c641b832 100644 --- a/cpp/src/arrow/json/chunker.h +++ b/cpp/src/arrow/json/chunker.h @@ -19,6 +19,7 @@ #include +#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 MakeChunker(const ParseOptions& options); +std::unique_ptr MakeChunker(const ParseOptions& options, + MemoryPool* pool = default_memory_pool()); } // namespace json } // namespace arrow diff --git a/cpp/src/arrow/json/chunker_test.cc b/cpp/src/arrow/json/chunker_test.cc index 0976e9ba22b..88661973a72 100644 --- a/cpp/src/arrow/json/chunker_test.cc +++ b/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 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) { diff --git a/cpp/src/arrow/json/parser_benchmark.cc b/cpp/src/arrow/json/parser_benchmark.cc index f26d6a14431..768beed3788 100644 --- a/cpp/src/arrow/json/parser_benchmark.cc +++ b/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(json->size()); } +static void BenchmarkJSONChunkingMultipleBlocks( + benchmark::State& state, // NOLINT non-const reference + const std::shared_ptr& json, ParseOptions options, int64_t block_size) { + auto chunker = MakeChunker(options); + + for (auto _ : state) { + std::shared_ptr 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(json->size()); + state.counters["block_size"] = static_cast(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(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(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); diff --git a/cpp/src/arrow/json/reader.cc b/cpp/src/arrow/json/reader.cc index 5c06e726b99..55d5b8770e0 100644 --- a/cpp/src/arrow/json/reader.cc +++ b/cpp/src/arrow/json/reader.cc @@ -253,8 +253,9 @@ class TableReaderImpl : public TableReader, } Result> Read() override { - auto block_it = MakeChunkingIterator(std::move(buffer_iterator_), - MakeChunker(decode_context_.parse_options())); + auto block_it = MakeChunkingIterator( + std::move(buffer_iterator_), + MakeChunker(decode_context_.parse_options(), decode_context_.pool())); bool did_read = false; while (true) { @@ -395,8 +396,8 @@ class StreamingReaderImpl : public StreamingReader { buffer_gen = [source = std::move(buffer_gen), cpu_executor] { return cpu_executor->TransferAlways(source()); }; - auto chunking_gen = MakeChunkingGenerator(std::move(buffer_gen), - MakeChunker(context->parse_options())); + auto chunking_gen = MakeChunkingGenerator( + std::move(buffer_gen), MakeChunker(context->parse_options(), context->pool())); // At this stage, we want to allow the decoding tasks for each chunked block to run // in parallel on the CPU executor. However: @@ -432,7 +433,7 @@ class StreamingReaderImpl : public StreamingReader { // the IO pool while we process its buffers on the calling thread auto chunking_it = MakeChunkingIterator(MakeGeneratorIterator(std::move(buffer_gen)), - MakeChunker(context->parse_options())); + MakeChunker(context->parse_options(), context->pool())); decoding_gen = MakeDecodingGenerator(std::move(chunking_it), DecodingOperator(context)); } diff --git a/cpp/src/arrow/json/reader_test.cc b/cpp/src/arrow/json/reader_test.cc index ac5bafb5293..7dcd30a0eb0 100644 --- a/cpp/src/arrow/json/reader_test.cc +++ b/cpp/src/arrow/json/reader_test.cc @@ -740,7 +740,6 @@ TEST_P(StreamingReaderTest, PropagateErrorsNonLinewiseChunker) { std::shared_ptr batch; std::shared_ptr reader; - Status status; read_options_.block_size = 10; parse_options_.newlines_in_values = true; @@ -759,20 +758,20 @@ TEST_P(StreamingReaderTest, PropagateErrorsNonLinewiseChunker) { AssertReadNext(reader, &batch); EXPECT_EQ(reader->bytes_processed(), 9); ASSERT_BATCHES_EQUAL(*RecordBatchFromJSON(test_schema, "[{\"i\":0}]"), *batch); - // Chunker doesn't require newline delimiters, so this should be valid + + // The chunker doesn't require newline delimiters between records. AssertReadNext(reader, &batch); EXPECT_EQ(reader->bytes_processed(), 20); ASSERT_BATCHES_EQUAL(*RecordBatchFromJSON(test_schema, "[{\"i\":1}]"), *batch); - EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, - ::testing::StartsWith("Invalid: JSON parse error"), - reader->ReadNext(&batch)); - EXPECT_EQ(reader->bytes_processed(), 20); - // Incoming chunker error from ":2}" shouldn't leak through after the first failure, - // which is a possibility if async tasks are still outstanding due to readahead. + // Depending on readahead and chunking, the malformed record may be reported + // by either of the two next reads. + auto status = reader->ReadNext(&batch) & reader->ReadNext(&batch); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("JSON parse error"), + status); + AssertReadEnd(reader); AssertReadEnd(reader); - EXPECT_EQ(reader->bytes_processed(), 20); } TEST_P(StreamingReaderTest, IgnoreLeadingEmptyBlocks) { diff --git a/python/pyarrow/tests/test_json.py b/python/pyarrow/tests/test_json.py index 8d5e6f43db0..ac2a027cfa2 100644 --- a/python/pyarrow/tests/test_json.py +++ b/python/pyarrow/tests/test_json.py @@ -43,17 +43,24 @@ def generate_col_names(): yield first + second -def make_random_json(num_cols=2, num_rows=10, linesep='\r\n'): - arr = np.random.RandomState(42).randint(0, 1000, size=(num_cols, num_rows)) +def make_random_json(num_cols=2, num_rows=10, linesep='\r\n', + *, extra_whitespace=False): + rnd = np.random.RandomState(42) + arr = rnd.randint(0, 1000, size=(num_cols, num_rows)) col_names = list(itertools.islice(generate_col_names(), num_cols)) lines = [] + max_line_len = 0 for row in arr.T: json_obj = OrderedDict([(k, int(v)) for (k, v) in zip(col_names, row)]) - lines.append(json.dumps(json_obj)) + line = json.dumps(json_obj) + if extra_whitespace: + line += ' ' * rnd.randint(0, 5) + lines.append(line) + max_line_len = max(max_line_len, len(line)) data = linesep.join(lines).encode() columns = [pa.array(col, type=pa.int64()) for col in arr] expected = pa.Table.from_arrays(columns, col_names) - return data, expected + return data, max_line_len, expected def check_options_class_pickling(cls, pickler, **attr_values): @@ -150,8 +157,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, match="try to increase block size"): self.read_bytes(data, read_options=read_options, parse_options=parse_options) @@ -295,39 +301,48 @@ def test_explicit_schema_with_unexpected_behaviour(self): @pytest.mark.numpy def test_small_random_json(self): - data, expected = make_random_json(num_cols=2, num_rows=10) - table = self.read_bytes(data) - assert table.schema == expected.schema - assert table.equals(expected) - assert table.to_pydict() == expected.to_pydict() + for extra_whitespace in (False, True): + data, _, expected = make_random_json(num_cols=2, num_rows=10, + extra_whitespace=extra_whitespace) + table = self.read_bytes(data) + assert table.schema == expected.schema + assert table.equals(expected) + assert table.to_pydict() == expected.to_pydict() @pytest.mark.numpy def test_load_large_json(self): - data, expected = make_random_json(num_cols=2, num_rows=100100) - # set block size is 10MB - read_options = ReadOptions(block_size=1024*1024*10) - table = self.read_bytes(data, read_options=read_options) - assert table.num_rows == 100100 - assert expected.num_rows == 100100 + for extra_whitespace in (False, True): + data, _, expected = make_random_json(num_cols=2, num_rows=100100, + extra_whitespace=extra_whitespace) + # set block size is 10MB + read_options = ReadOptions(block_size=1024*1024*10) + table = self.read_bytes(data, read_options=read_options) + assert table.num_rows == 100100 + assert expected.num_rows == 100100 @pytest.mark.numpy def test_stress_block_sizes(self): # Test a number of small block sizes to stress block stitching - data_base, expected = make_random_json(num_cols=2, num_rows=100) - read_options = ReadOptions() - parse_options = ParseOptions() - - for data in [data_base, data_base.rstrip(b'\r\n')]: - for newlines_in_values in [False, True]: - parse_options.newlines_in_values = newlines_in_values - for block_size in [22, 23, 37]: - read_options.block_size = block_size - table = self.read_bytes(data, read_options=read_options, - parse_options=parse_options) - assert table.schema == expected.schema - if not table.equals(expected): - # Better error output - assert table.to_pydict() == expected.to_pydict() + for extra_whitespace in (False, True): + data_base, max_line_len, expected = make_random_json( + num_cols=2, num_rows=100, + extra_whitespace=extra_whitespace) + block_sizes = [max_line_len, max_line_len + 4, max_line_len + 15] + + read_options = ReadOptions() + parse_options = ParseOptions() + + for data in [data_base, data_base.rstrip(b'\r\n')]: + for newlines_in_values in [False, True]: + parse_options.newlines_in_values = newlines_in_values + for block_size in block_sizes: + read_options.block_size = block_size + table = self.read_bytes(data, read_options=read_options, + parse_options=parse_options) + assert table.schema == expected.schema + if not table.equals(expected): + # Better error output + assert table.to_pydict() == expected.to_pydict() class BaseTestJSONRead(BaseTestJSON): @@ -529,8 +544,9 @@ def test_non_linewise_chunker_bad_middle_block(self): 'n': [1] } - with pytest.raises(pa.ArrowInvalid, - match="JSON parse error *"): + # The first read may succeed depending on chunking and readahead. + with pytest.raises(pa.ArrowInvalid, match="JSON parse error"): + assert reader.read_next_batch().to_pydict() == {'n': [None]} reader.read_next_batch() with pytest.raises(StopIteration):