diff --git a/cpp/src/arrow/buffer.cc b/cpp/src/arrow/buffer.cc index ab20ce7fb9f0..f069277e7c37 100644 --- a/cpp/src/arrow/buffer.cc +++ b/cpp/src/arrow/buffer.cc @@ -87,7 +87,7 @@ Result> SliceMutableBufferSafe(std::shared_ptr b return SliceMutableBuffer(std::move(buffer), offset, length); } -std::string Buffer::ToHexString() { +std::string Buffer::ToHexString() const { return HexEncode(data(), static_cast(size())); } diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h index a2c841208273..07f9931eba73 100644 --- a/cpp/src/arrow/buffer.h +++ b/cpp/src/arrow/buffer.h @@ -121,7 +121,7 @@ class ARROW_EXPORT Buffer { /// \brief Construct a new std::string with a hexadecimal representation of the buffer. /// \return std::string - std::string ToHexString(); + std::string ToHexString() const; /// Return true if both buffers are the same size and contain the same bytes /// up to the number of compared bytes diff --git a/cpp/src/arrow/ipc/message.cc b/cpp/src/arrow/ipc/message.cc index 1fef961ff8b3..8522475782e3 100644 --- a/cpp/src/arrow/ipc/message.cc +++ b/cpp/src/arrow/ipc/message.cc @@ -285,17 +285,21 @@ std::string FormatMessageType(MessageType type) { namespace { -Status ReadFieldsSubset(int64_t offset, int32_t metadata_length, - io::RandomAccessFile* file, +Status ReadFieldsSubset(int64_t offset, io::RandomAccessFile* file, const FieldsLoaderFunction& fields_loader, const std::shared_ptr& metadata, int64_t required_size, std::shared_ptr& body) { + DCHECK_GE(static_cast(metadata->size()), sizeof(int32_t)); + const auto continuation = util::SafeLoadAs(metadata->data()); + // Either 8 bytes (32-bit continuation indicator + 32-bit little-endian length prefix) + // or 4 bytes for legacy IPC without continuation indicator + const auto continuation_size = (continuation == internal::kIpcContinuationToken) + ? 2 * sizeof(int32_t) + : sizeof(int32_t); + const flatbuf::Message* message = nullptr; - uint8_t continuation_metadata_size = sizeof(int32_t) + sizeof(int32_t); - // skip 8 bytes (32-bit continuation indicator + 32-bit little-endian length prefix) - RETURN_NOT_OK(internal::VerifyMessage(metadata->data() + continuation_metadata_size, - metadata->size() - continuation_metadata_size, - &message)); + RETURN_NOT_OK(internal::VerifyMessage(metadata->data() + continuation_size, + metadata->size() - continuation_size, &message)); auto batch = message->header_as_RecordBatch(); if (batch == nullptr) { return Status::IOError( @@ -305,8 +309,8 @@ Status ReadFieldsSubset(int64_t offset, int32_t metadata_length, RETURN_NOT_OK(fields_loader(batch, &io_recorded_random_access_file)); const auto& read_ranges = io_recorded_random_access_file.GetReadRanges(); for (const auto& range : read_ranges) { - auto read_result = file->ReadAt(offset + metadata_length + range.offset, range.length, - body->mutable_data() + range.offset); + auto read_result = file->ReadAt(offset + metadata->size() + range.offset, + range.length, body->mutable_data() + range.offset); if (!read_result.ok()) { return Status::IOError("Failed to read message body, error ", read_result.status().ToString()); @@ -315,6 +319,86 @@ Status ReadFieldsSubset(int64_t offset, int32_t metadata_length, return Status::OK(); } +struct ReadMessageState { + std::unique_ptr result; + std::shared_ptr listener; + std::shared_ptr decoder; + + ReadMessageState() + : listener(std::make_shared(&result)), + decoder(std::make_shared(listener)) {} + + // ReadMessageState points into itself, so it shouldn't be moved + ReadMessageState(ReadMessageState&&) = delete; + ReadMessageState& operator=(ReadMessageState&&) = delete; +}; + +// A common continuation callback for ReadMessage and ReadMessageAsync overloads +static Result> ReadMessageContinued( + int64_t offset, int32_t metadata_length, std::optional body_length, + std::shared_ptr metadata, io::RandomAccessFile* file, + const FieldsLoaderFunction& fields_loader, ReadMessageState* state) { + MessageDecoder* decoder = state->decoder.get(); + if (body_length.has_value()) { + // If body length was known, the given buffer should contain exactly the metadata + // followed by the body. + DCHECK_EQ(metadata->size(), metadata_length + *body_length); + } + ARROW_RETURN_NOT_OK(decoder->Consume(SliceBuffer(metadata, 0, metadata_length))); + if (decoder->buffered_size() > 0) { + return Status::Invalid("Message metadata too long by ", decoder->buffered_size(), + " bytes"); + } + switch (decoder->state()) { + case MessageDecoder::State::INITIAL: + return std::move(state->result); + case MessageDecoder::State::METADATA_LENGTH: + return Status::Invalid("metadata length is missing. File offset: ", offset, + ", metadata length: ", metadata_length); + case MessageDecoder::State::METADATA: + return Status::Invalid("flatbuffer size ", decoder->next_required_size(), + " invalid. File offset: ", offset, + ", metadata length: ", metadata_length); + case MessageDecoder::State::BODY: { + std::shared_ptr body; + if (fields_loader) { + // Selective field loading: allocate a body buffer and read only the + // requested field ranges into it. + DCHECK_NE(file, nullptr); + ARROW_ASSIGN_OR_RAISE( + body, AllocateBuffer(decoder->next_required_size(), default_memory_pool())); + RETURN_NOT_OK(ReadFieldsSubset(offset, file, fields_loader, + SliceBuffer(metadata, 0, metadata_length), + decoder->next_required_size(), body)); + } else if (body_length.has_value()) { + // Body was already read as part of the combined IO; just slice it out. + if (*body_length != decoder->next_required_size()) { + // The streaming decoder got out of sync with the actual advertised + // metadata and body size, which signals an invalid IPC file. + return Status::IOError("Invalid IPC file: advertised body size is ", + *body_length, ", but message decoder expects to read ", + decoder->next_required_size(), " bytes instead"); + } + body = SliceBuffer(metadata, metadata_length, + std::min(*body_length, metadata->size() - metadata_length)); + } else { + // Body length was unknown; do a separate IO to read the body. + DCHECK_NE(file, nullptr); + ARROW_ASSIGN_OR_RAISE( + body, file->ReadAt(offset + metadata_length, decoder->next_required_size(), + /*allow_short_read=*/false)); + } + + RETURN_NOT_OK(decoder->Consume(body)); + return std::move(state->result); + } + case MessageDecoder::State::EOS: + return Status::Invalid("Unexpected empty message in IPC file format"); + default: + return Status::Invalid("Unexpected state: ", state->decoder->state()); + } +} + } // namespace Result> ReadMessage(std::shared_ptr metadata, @@ -330,6 +414,10 @@ Result> ReadMessage(std::shared_ptr metadata, } ARROW_RETURN_NOT_OK(decoder.Consume(metadata)); + if (decoder.buffered_size() > 0) { + return Status::Invalid("Message metadata too long by ", decoder.buffered_size(), + " bytes"); + } switch (decoder.state()) { case MessageDecoder::State::INITIAL: @@ -368,69 +456,20 @@ Result> ReadMessage(std::shared_ptr metadata, static Result> ReadMessageInternal( int64_t offset, int32_t metadata_length, std::optional body_length, io::RandomAccessFile* file, const FieldsLoaderFunction& fields_loader) { - std::unique_ptr result; - auto listener = std::make_shared(&result); - MessageDecoder decoder(listener); + ReadMessageState state; - if (metadata_length < decoder.next_required_size()) { + if (metadata_length < state.decoder->next_required_size()) { return Status::Invalid("metadata_length should be at least ", - decoder.next_required_size()); + state.decoder->next_required_size()); } - // When body_length is known, read metadata + body in one IO call. // Otherwise, read only metadata first. ARROW_ASSIGN_OR_RAISE(std::shared_ptr metadata, file->ReadAt(offset, metadata_length + body_length.value_or(0), /*allow_short_read=*/false)); - ARROW_RETURN_NOT_OK(decoder.Consume(SliceBuffer(metadata, 0, metadata_length))); - - switch (decoder.state()) { - case MessageDecoder::State::INITIAL: - return result; - case MessageDecoder::State::METADATA_LENGTH: - return Status::Invalid("metadata length is missing. File offset: ", offset, - ", metadata length: ", metadata_length); - case MessageDecoder::State::METADATA: - return Status::Invalid("flatbuffer size ", decoder.next_required_size(), - " invalid. File offset: ", offset, - ", metadata length: ", metadata_length); - case MessageDecoder::State::BODY: { - std::shared_ptr body; - if (fields_loader) { - // Selective field loading: allocate a body buffer and read only the - // requested field ranges into it. - ARROW_ASSIGN_OR_RAISE( - body, AllocateBuffer(decoder.next_required_size(), default_memory_pool())); - RETURN_NOT_OK(ReadFieldsSubset(offset, metadata_length, file, fields_loader, - SliceBuffer(metadata, 0, metadata_length), - decoder.next_required_size(), body)); - } else if (body_length.has_value()) { - // Body was already read as part of the combined IO; just slice it out. - if (*body_length != decoder.next_required_size()) { - // The streaming decoder got out of sync with the actual advertised - // metadata and body size, which signals an invalid IPC file. - return Status::IOError("Invalid IPC file: advertised body size is ", - *body_length, ", but message decoder expects to read ", - decoder.next_required_size(), " bytes instead"); - } - body = SliceBuffer(metadata, metadata_length, - std::min(*body_length, metadata->size() - metadata_length)); - } else { - // Body length was unknown; do a separate IO to read the body. - ARROW_ASSIGN_OR_RAISE( - body, file->ReadAt(offset + metadata_length, decoder.next_required_size(), - /*allow_short_read=*/false)); - } - - RETURN_NOT_OK(decoder.Consume(body)); - return result; - } - case MessageDecoder::State::EOS: - return Status::Invalid("Unexpected empty message in IPC file format"); - default: - return Status::Invalid("Unexpected state: ", decoder.state()); - } + return ReadMessageContinued(offset, metadata_length, body_length, metadata, file, + fields_loader, &state); } Result> ReadMessage(int64_t offset, int32_t metadata_length, @@ -452,14 +491,9 @@ Future> ReadMessageAsync(int64_t offset, int32_t metada int64_t body_length, io::RandomAccessFile* file, const io::IOContext& context) { - struct State { - std::unique_ptr result; - std::shared_ptr listener; - std::shared_ptr decoder; - }; - auto state = std::make_shared(); - state->listener = std::make_shared(&state->result); - state->decoder = std::make_shared(state->listener); + // Make a std::shared_ptr so as to have a stable ReadMessageState pointer, + // since state->listener will point back to state->result. + auto state = std::make_shared(); if (metadata_length < state->decoder->next_required_size()) { return Status::Invalid("metadata_length should be at least ", @@ -469,35 +503,11 @@ Future> ReadMessageAsync(int64_t offset, int32_t metada ->ReadAsync(context, offset, metadata_length + body_length, /*allow_short_read=*/false) .Then([=](std::shared_ptr metadata) -> Result> { - DCHECK_EQ(metadata->size(), metadata_length + body_length); - ARROW_RETURN_NOT_OK( - state->decoder->Consume(SliceBuffer(metadata, 0, metadata_length))); - switch (state->decoder->state()) { - case MessageDecoder::State::INITIAL: - return std::move(state->result); - case MessageDecoder::State::METADATA_LENGTH: - return Status::Invalid("metadata length is missing. File offset: ", offset, - ", metadata length: ", metadata_length); - case MessageDecoder::State::METADATA: - return Status::Invalid("flatbuffer size ", - state->decoder->next_required_size(), - " invalid. File offset: ", offset, - ", metadata length: ", metadata_length); - case MessageDecoder::State::BODY: { - auto body = SliceBuffer(metadata, metadata_length, body_length); - if (body->size() < state->decoder->next_required_size()) { - return Status::IOError("Expected to be able to read ", - state->decoder->next_required_size(), - " bytes for message body, got ", body->size()); - } - RETURN_NOT_OK(state->decoder->Consume(body)); - return std::move(state->result); - } - case MessageDecoder::State::EOS: - return Status::Invalid("Unexpected empty message in IPC file format"); - default: - return Status::Invalid("Unexpected state: ", state->decoder->state()); - } + // Pass a nullptr file to ensure that no further IO occurs + // (we have fetched all the required bytes). + return ReadMessageContinued(offset, metadata_length, body_length, metadata, + /*file=*/nullptr, + /*fields_loader=*/{}, state.get()); }); } @@ -755,6 +765,8 @@ class MessageDecoder::MessageDecoderImpl { int64_t next_required_size() const { return next_required_size_ - buffered_size_; } + int64_t buffered_size() const { return buffered_size_; } + MessageDecoder::State state() const { return state_; } private: @@ -1020,6 +1032,8 @@ Status MessageDecoder::Consume(std::shared_ptr buffer) { int64_t MessageDecoder::next_required_size() const { return impl_->next_required_size(); } +int64_t MessageDecoder::buffered_size() const { return impl_->buffered_size(); } + MessageDecoder::State MessageDecoder::state() const { return impl_->state(); } // ---------------------------------------------------------------------- diff --git a/cpp/src/arrow/ipc/message.h b/cpp/src/arrow/ipc/message.h index df80b0eba252..c6701f200f45 100644 --- a/cpp/src/arrow/ipc/message.h +++ b/cpp/src/arrow/ipc/message.h @@ -419,6 +419,11 @@ class ARROW_EXPORT MessageDecoder { /// \return the current state State state() const; + /// \brief Return the number of bytes buffered in the decoder. + /// + /// This method is mainly useful for testing and debugging. + int64_t buffered_size() const; + private: class MessageDecoderImpl; std::unique_ptr impl_; @@ -453,9 +458,12 @@ using FieldsLoaderFunction = std::function> ReadMessage( /// Read a length-prefixed message flatbuffer starting at the indicated file /// offset. /// -/// The metadata_length includes at least the length prefix and the flatbuffer +/// The metadata_length includes the IPC encapsulation prefix and the +/// Flatbuffers-serialized message. /// /// \param[in] offset the position in the file where the message starts. The /// first 4 bytes after the offset are the message length @@ -578,6 +587,7 @@ Status DecodeMessage(MessageDecoder* decoder, io::InputStream* stream); /// \param[out] message_length the total size of the payload written including /// padding /// \return Status +ARROW_EXPORT Status WriteMessage(const Buffer& message, const IpcWriteOptions& options, io::OutputStream* file, int32_t* message_length); diff --git a/cpp/src/arrow/ipc/message_internal_test.cc b/cpp/src/arrow/ipc/message_internal_test.cc index 87ce3de39cb7..bc4506eaaed3 100644 --- a/cpp/src/arrow/ipc/message_internal_test.cc +++ b/cpp/src/arrow/ipc/message_internal_test.cc @@ -15,15 +15,27 @@ // specific language governing permissions and limitations // under the License. +#include +#include +#include +#include + #include + #include -#include + +#include "Message_generated.h" #include "arrow/buffer.h" +#include "arrow/io/memory.h" #include "arrow/ipc/dictionary.h" +#include "arrow/ipc/message.h" #include "arrow/ipc/metadata_internal.h" #include "arrow/ipc/options.h" #include "arrow/ipc/reader.h" +#include "arrow/ipc/type_fwd.h" +#include "arrow/result.h" +#include "arrow/status.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/key_value_metadata.h" @@ -103,4 +115,332 @@ TEST(TestMessageInternal, TestEndiannessRoundtrip) { } } +struct SampleMessageParams { + std::shared_ptr custom_metadata = {}; + int64_t body_length = 0; + IpcWriteOptions options = {}; + + std::string ToString() const { + std::stringstream ss; + ss << *this; + return std::move(ss).str(); + } + + friend std::ostream& operator<<(std::ostream& os, const SampleMessageParams& p) { + os << "legacy IPC = " << p.options.write_legacy_ipc_format << ", " + << "body length = " << p.body_length + << ", metadata length = " << (p.custom_metadata ? p.custom_metadata->size() : 0); + return os; + } +}; + +struct SampleMessage { + std::shared_ptr metadata_bytes; // encapsulated with IPC framing + std::shared_ptr body_bytes; + MessageType message_type; + int64_t num_rows = 0; + std::shared_ptr custom_metadata; +}; + +class MessageDecodingTest : public ::testing::Test { + public: + static constexpr int64_t kNumRows = 5; + + std::vector write_options() { + return {IpcWriteOptions{}, IpcWriteOptions{.write_legacy_ipc_format = true}, + IpcWriteOptions{.alignment = 32}}; + } + + std::vector message_params() { + std::vector params; + for (const auto& options : write_options()) { + for (int64_t body_length : {0, 24}) { + params.push_back(SampleMessageParams{ + .custom_metadata = nullptr, .body_length = body_length, .options = options}); + params.push_back(SampleMessageParams{.custom_metadata = GetCustomMetadata(), + .body_length = body_length, + .options = options}); + } + } + return params; + } + + // Return the serialized metadata encapsulated in IPC framing + Result> EncapsulateMetadata( + const std::shared_ptr& metadata_bytes, const IpcWriteOptions& options) { + ARROW_ASSIGN_OR_RAISE(auto out_stream, ::arrow::io::BufferOutputStream::Create()); + int32_t written_bytes = 0; + RETURN_NOT_OK( + WriteMessage(*metadata_bytes, options, out_stream.get(), &written_bytes)); + ARROW_ASSIGN_OR_RAISE(auto out, out_stream->Finish()); + return out; + } + + Result GetSampleMessage(const SampleMessageParams& params) { + // Create a dummy RecordBatch message + auto field_md = + std::vector{FieldMetadata{.length = kNumRows, .null_count = 1, .offset = 0}}; + auto buffer_md = std::vector{BufferMetadata{.offset = 64, .length = 1}, + BufferMetadata{.offset = 72, .length = 10}}; + std::shared_ptr out; + RETURN_NOT_OK(WriteRecordBatchMessage(/*length=*/kNumRows, params.body_length, + params.custom_metadata, field_md, buffer_md, + /*variadic_counts=*/{}, params.options, &out)); + ARROW_ASSIGN_OR_RAISE(out, EncapsulateMetadata(out, params.options)); + // Generate a dummy body of the advertised length + ARROW_ASSIGN_OR_RAISE(auto body_bytes, AllocateBuffer(params.body_length)); + memset(body_bytes->mutable_data(), '!', body_bytes->size()); + return SampleMessage{.metadata_bytes = out, + .body_bytes = std::move(body_bytes), + .message_type = MessageType::RECORD_BATCH, + .num_rows = kNumRows, + .custom_metadata = params.custom_metadata}; + } + + std::shared_ptr GetCustomMetadata() { + return KeyValueMetadata::Make(/*keys=*/{"key1", "key2"}, /*values=*/{"foo", "bar"}); + } + + void CheckSampleMessage(const Message& message, const SampleMessage& sample_message) { + ASSERT_EQ(message.type(), sample_message.message_type); + const auto batch = reinterpret_cast(message.header()); + ASSERT_EQ(batch->length(), sample_message.num_rows); + ASSERT_EQ(message.body_length(), sample_message.body_bytes->size()); + if (message.body_length() > 0) { + AssertBufferEqual(*message.body(), *sample_message.body_bytes, /*verbose=*/true); + } + if (sample_message.custom_metadata && sample_message.custom_metadata->size() > 0) { + ASSERT_NE(message.custom_metadata(), nullptr); + ASSERT_TRUE(message.custom_metadata()->Equals(*sample_message.custom_metadata)); + } else { + ASSERT_EQ(message.custom_metadata(), nullptr); + } + } + + // Return concatenated metadata and body bytes + Result> ConcatenateMessage(const SampleMessage& sample_message, + int64_t padding_size = 0) { + auto padding_buffer = Buffer::FromString(std::string(padding_size, 'x')); + return ConcatenateBuffers({padding_buffer, sample_message.metadata_bytes, + sample_message.body_bytes, padding_buffer}); + } + + void CheckDecoding(const std::shared_ptr& buffer, int64_t chunk_size, + const SampleMessage& sample_message) { + std::unique_ptr message; + auto listener = std::make_shared(&message); + MessageDecoder decoder(listener); + int64_t offset = 0; + ASSERT_EQ(decoder.buffered_size(), 0); + while (offset < buffer->size()) { + // No message was decoded yet + ASSERT_EQ(message, nullptr); + // The decoder is expecting more data, but not more than remaining in our buffer + ASSERT_GT(decoder.next_required_size(), 0); + ASSERT_LE(decoder.next_required_size(), buffer->size() - offset); + const auto to_consume = std::min(chunk_size, buffer->size() - offset); + ASSERT_OK(decoder.Consume(SliceBuffer(buffer, offset, to_consume))); + offset += to_consume; + if (offset >= 4 && offset < buffer->size()) { + // We went past the initial 4-byte continuation + ASSERT_NE(decoder.state(), MessageDecoder::INITIAL); + if (offset >= buffer->size() - sample_message.body_bytes->size()) { + // The offset points in the body + ASSERT_EQ(decoder.state(), MessageDecoder::BODY); + } + } + } + ASSERT_EQ(decoder.buffered_size(), 0); + ASSERT_EQ(decoder.state(), MessageDecoder::INITIAL); + ASSERT_NE(message, nullptr); + CheckSampleMessage(*message, sample_message); + } + + void TestDecoding(const SampleMessage& sample_message) { + ASSERT_OK_AND_ASSIGN(auto buffer, ConcatenateMessage(sample_message)); + for (const auto chunk_size : std::vector{ + 1, 2, 3, buffer->size() / 3, buffer->size() - 1, buffer->size()}) { + ARROW_SCOPED_TRACE("chunk_size = ", chunk_size); + CheckDecoding(buffer, chunk_size, sample_message); + } + } + + void TestDecoding(const SampleMessageParams& params) { + ASSERT_OK_AND_ASSIGN(auto message, GetSampleMessage(params)); + TestDecoding(message); + } + + template + void CheckReadMessageOk(ReadMessageFunc read_message, + const SampleMessageParams& params) { + ASSERT_OK_AND_ASSIGN(auto sample_message, GetSampleMessage(params)); + ASSERT_OK_AND_ASSIGN(auto message, read_message(sample_message)); + CheckSampleMessage(*message, sample_message); + } + + template + void CheckReadMessageTruncated(ReadMessageFunc read_message, + const SampleMessageParams& params, + bool force_truncate_metadata = false) { + ASSERT_OK_AND_ASSIGN(auto sample_message, GetSampleMessage(params)); + if (force_truncate_metadata || sample_message.body_bytes->size() == 0) { + sample_message.metadata_bytes = + SliceBuffer(sample_message.metadata_bytes, /*offset=*/0, + sample_message.metadata_bytes->size() - 1); + sample_message.body_bytes = SliceBuffer(sample_message.body_bytes, /*offset=*/0, + /*length=*/0); + } else { + sample_message.body_bytes = SliceBuffer(sample_message.body_bytes, /*offset=*/0, + sample_message.body_bytes->size() - 1); + } + Status status = read_message(sample_message).status(); + ASSERT_TRUE(status.IsInvalid() || status.IsIOError()) + << "Unexpected status: " << status.ToString(); + } + + template + void CheckReadMessageOversized(ReadMessageFunc read_message, + const SampleMessageParams& params) { + ASSERT_OK_AND_ASSIGN(auto sample_message, GetSampleMessage(params)); + auto trailing_bytes = Buffer::FromString("x"); + if (sample_message.body_bytes->size() > 0) { + ASSERT_OK_AND_ASSIGN( + sample_message.body_bytes, + ConcatenateBuffers({sample_message.body_bytes, trailing_bytes})); + } else { + ASSERT_OK_AND_ASSIGN( + sample_message.metadata_bytes, + ConcatenateBuffers({sample_message.metadata_bytes, trailing_bytes})); + } + Status status = read_message(sample_message).status(); + ASSERT_TRUE(status.IsInvalid() || status.IsIOError()) + << "Unexpected status: " << status.ToString(); + } +}; + +TEST_F(MessageDecodingTest, MessageDecoder) { + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + TestDecoding(params); + } +} + +TEST_F(MessageDecodingTest, ReadMessage1) { + auto read_message = [&](const SampleMessage& sample_message) { + std::shared_ptr body = + sample_message.body_bytes->size() > 0 ? sample_message.body_bytes : nullptr; + return ReadMessage(sample_message.metadata_bytes, body); + }; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + CheckReadMessageTruncated(read_message, params); + CheckReadMessageOversized(read_message, params); + } +} + +TEST_F(MessageDecodingTest, ReadMessage2) { + auto read_message = + [&](const SampleMessage& sample_message) -> Result> { + const int kStreamOffset = 42; + ARROW_ASSIGN_OR_RAISE( + auto stream_buf, + ConcatenateMessage(sample_message, /*padding_size=*/kStreamOffset)); + io::BufferReader reader(stream_buf); + return ReadMessage(kStreamOffset, + static_cast(sample_message.metadata_bytes->size()), + static_cast(sample_message.body_bytes->size()), &reader); + }; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + CheckReadMessageTruncated(read_message, params); + CheckReadMessageOversized(read_message, params); + } +} + +TEST_F(MessageDecodingTest, ReadMessageAsync) { + auto read_message = + [&](const SampleMessage& sample_message) -> Result> { + const int kStreamOffset = 42; + ARROW_ASSIGN_OR_RAISE( + auto stream_buf, + ConcatenateMessage(sample_message, /*padding_size=*/kStreamOffset)); + io::BufferReader reader(stream_buf); + return ReadMessageAsync( + kStreamOffset, static_cast(sample_message.metadata_bytes->size()), + static_cast(sample_message.body_bytes->size()), &reader) + .result(); + }; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + CheckReadMessageTruncated(read_message, params); + CheckReadMessageOversized(read_message, params); + } +} + +TEST_F(MessageDecodingTest, ReadMessage3) { + int padding_size; + + auto read_message = + [&](const SampleMessage& sample_message) -> Result> { + // No padding, so that reading the truncated message actually fails + ARROW_ASSIGN_OR_RAISE(auto stream_buf, + ConcatenateMessage(sample_message, padding_size)); + io::BufferReader reader(stream_buf); + return ReadMessage(/*offset=*/padding_size, + static_cast(sample_message.metadata_bytes->size()), + &reader, /*fields_loader=*/{}); + }; + padding_size = 0; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + CheckReadMessageTruncated(read_message, params); + } + // With a non-zero padding, a truncated message wouldn't fail + padding_size = 42; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + } +} + +TEST_F(MessageDecodingTest, ReadMessage4) { + FieldsLoaderFunction fields_loader = [&](const void* void_batch, + io::RandomAccessFile* file) -> Status { + auto* batch = reinterpret_cast(void_batch); + // Check something about the message header + EXPECT_EQ(batch->length(), kNumRows); + // Read the entire body range from the file + ARROW_ASSIGN_OR_RAISE(auto read_size, file->GetSize()); + return file->ReadAt(/*position=*/0, read_size, /*allow_short_read=*/false).status(); + }; + + int padding_size; + auto read_message = + [&](const SampleMessage& sample_message) -> Result> { + // No padding, so that reading the truncated message actually fails + ARROW_ASSIGN_OR_RAISE(auto stream_buf, + ConcatenateMessage(sample_message, padding_size)); + io::BufferReader reader(stream_buf); + return ReadMessage(/*offset=*/padding_size, + static_cast(sample_message.metadata_bytes->size()), + &reader, fields_loader); + }; + padding_size = 0; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + CheckReadMessageTruncated(read_message, params, /*force_truncate_metadata=*/true); + } + // With a non-zero padding, a truncated message wouldn't fail + padding_size = 42; + for (const auto& params : message_params()) { + ARROW_SCOPED_TRACE("Params: ", params); + CheckReadMessageOk(read_message, params); + } +} + } // namespace arrow::ipc::internal diff --git a/cpp/src/arrow/ipc/options.h b/cpp/src/arrow/ipc/options.h index ec0e2a5b6f90..254a4693866e 100644 --- a/cpp/src/arrow/ipc/options.h +++ b/cpp/src/arrow/ipc/options.h @@ -67,7 +67,7 @@ struct ARROW_EXPORT IpcWriteOptions { /// \brief Compression codec to use for record batch body buffers /// /// May only be UNCOMPRESSED, LZ4_FRAME and ZSTD. - std::shared_ptr codec; + std::shared_ptr codec = {}; /// \brief Minimum space savings percentage required for compression to be applied /// @@ -82,7 +82,7 @@ struct ARROW_EXPORT IpcWriteOptions { /// /// Note that enabling this option may result in unreadable data for Arrow C++ versions /// prior to 12.0.0. - std::optional min_space_savings; + std::optional min_space_savings = {}; /// \brief Use global CPU thread pool to parallelize any computational tasks /// like compression diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc index 22a96111f4fc..76d286dbdcc7 100644 --- a/cpp/src/arrow/ipc/reader.cc +++ b/cpp/src/arrow/ipc/reader.cc @@ -2875,7 +2875,7 @@ Status FuzzIpcFile(const uint8_t* data, int64_t size) { } } - if (maybe_read_result.has_value()) { + if (final_status.ok()) { // IPC file read successful: compare results with IPC stream reader, // if possible. // NOTE: some valid IPC files may not be readable as IPC streams, diff --git a/cpp/src/arrow/testing/gtest_util.cc b/cpp/src/arrow/testing/gtest_util.cc index 0e2cbdb644ac..3ea7d9bb22ec 100644 --- a/cpp/src/arrow/testing/gtest_util.cc +++ b/cpp/src/arrow/testing/gtest_util.cc @@ -237,9 +237,17 @@ void AssertBufferEqual(const Buffer& buffer, std::string_view expected) { } } -void AssertBufferEqual(const Buffer& buffer, const Buffer& expected) { - ASSERT_EQ(buffer.size(), expected.size()) << "Mismatching buffer size"; - ASSERT_TRUE(buffer.Equals(expected)); +void AssertBufferEqual(const Buffer& buffer, const Buffer& expected, bool verbose) { + ASSERT_EQ(buffer.size(), expected.size()) + << "Mismatching buffer size, got " << buffer.size() << ", expected " + << expected.size(); + if (verbose) { + ASSERT_TRUE(buffer.Equals(expected)) + << "Mismatching buffers, got : " << buffer.ToHexString() + << " but expected: " << expected.ToHexString(); + } else { + ASSERT_TRUE(buffer.Equals(expected)); + } } template diff --git a/cpp/src/arrow/testing/gtest_util.h b/cpp/src/arrow/testing/gtest_util.h index b84d253a89e8..a0f433d3a9cd 100644 --- a/cpp/src/arrow/testing/gtest_util.h +++ b/cpp/src/arrow/testing/gtest_util.h @@ -253,7 +253,8 @@ ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, const std::vector& expected); ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, std::string_view expected); -ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, const Buffer& expected); +ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, const Buffer& expected, + bool verbose = false); ARROW_TESTING_EXPORT void AssertTypeEqual(const DataType& lhs, const DataType& rhs, bool check_metadata = false); diff --git a/testing b/testing index 1d74fce2b6bb..9ff285c88565 160000 --- a/testing +++ b/testing @@ -1 +1 @@ -Subproject commit 1d74fce2b6bb30158f254fc292252f4a87fc67a3 +Subproject commit 9ff285c88565f0f6abc855918c6a342e70e4909c