You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by we...@apache.org on 2019/07/29 14:45:02 UTC
[arrow] branch master updated: ARROW-6006: [C++] Do not fail to
read empty IPC stream with schema having dictionary types
This is an automated email from the ASF dual-hosted git repository.
wesm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 7f30a53 ARROW-6006: [C++] Do not fail to read empty IPC stream with schema having dictionary types
7f30a53 is described below
commit 7f30a536b78d4fa313b706e0ec34627d58742712
Author: Wes McKinney <we...@apache.org>
AuthorDate: Mon Jul 29 09:44:53 2019 -0500
ARROW-6006: [C++] Do not fail to read empty IPC stream with schema having dictionary types
Closes #4947 from wesm/ARROW-6006 and squashes the following commits:
e8d54aa58 <Wes McKinney> Fix typo
b442390c3 <Wes McKinney> Do not fail to read empty stream with schema having dictionary types
Authored-by: Wes McKinney <we...@apache.org>
Signed-off-by: Wes McKinney <we...@apache.org>
---
cpp/src/arrow/ipc/read-write-test.cc | 23 ++++++++++++++++
cpp/src/arrow/ipc/reader.cc | 53 +++++++++++++++++++++---------------
2 files changed, 54 insertions(+), 22 deletions(-)
diff --git a/cpp/src/arrow/ipc/read-write-test.cc b/cpp/src/arrow/ipc/read-write-test.cc
index 96ead3d..d5a0034 100644
--- a/cpp/src/arrow/ipc/read-write-test.cc
+++ b/cpp/src/arrow/ipc/read-write-test.cc
@@ -829,6 +829,29 @@ TEST_F(TestStreamFormat, DifferentSchema) { TestWriteDifferentSchema(); }
TEST_F(TestFileFormat, DifferentSchema) { TestWriteDifferentSchema(); }
+TEST(TestRecordBatchStreamReader, EmptyStreamWithDictionaries) {
+ // ARROW-6006
+ auto f0 = arrow::field("f0", arrow::dictionary(arrow::int8(), arrow::utf8()));
+ auto schema = arrow::schema({f0});
+
+ std::shared_ptr<io::BufferOutputStream> out;
+ ASSERT_OK(io::BufferOutputStream::Create(0, default_memory_pool(), &out));
+
+ std::shared_ptr<RecordBatchWriter> writer;
+ ASSERT_OK(RecordBatchStreamWriter::Open(out.get(), schema, &writer));
+ ASSERT_OK(writer->Close());
+
+ std::shared_ptr<arrow::Buffer> buffer;
+ ASSERT_OK(out->Finish(&buffer));
+ io::BufferReader buffer_reader(buffer);
+ std::shared_ptr<RecordBatchReader> reader;
+ ASSERT_OK(RecordBatchStreamReader::Open(&buffer_reader, &reader));
+
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(reader->ReadNext(&batch));
+ ASSERT_EQ(nullptr, batch);
+}
+
class TestTensorRoundTrip : public ::testing::Test, public IpcTestFixture {
public:
void SetUp() { pool_ = default_memory_pool(); }
diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
index c39f2d7..8806c61 100644
--- a/cpp/src/arrow/ipc/reader.cc
+++ b/cpp/src/arrow/ipc/reader.cc
@@ -431,20 +431,6 @@ Status ReadDictionary(const Buffer& metadata, DictionaryMemo* dictionary_memo,
return dictionary_memo->AddDictionary(id, dictionary);
}
-static Status ReadMessageAndValidate(MessageReader* reader, bool allow_null,
- std::unique_ptr<Message>* message) {
- RETURN_NOT_OK(reader->ReadNextMessage(message));
-
- if (*message == nullptr) {
- if (!allow_null) {
- return Status::Invalid("Expected message in stream, was null or length 0");
- }
- // End of stream?
- }
-
- return Status::OK();
-}
-
// ----------------------------------------------------------------------
// RecordBatchStreamReader implementation
@@ -464,9 +450,10 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
Status ReadSchema() {
std::unique_ptr<Message> message;
- RETURN_NOT_OK(
- ReadMessageAndValidate(message_reader_.get(), /*allow_null=*/false, &message));
-
+ RETURN_NOT_OK(message_reader_->ReadNextMessage(&message));
+ if (!message) {
+ return Status::Invalid("Tried reading schema message, was null or length 0");
+ }
CHECK_MESSAGE_TYPE(Message::SCHEMA, message->type());
CHECK_HAS_NO_BODY(*message);
if (message->header() == nullptr) {
@@ -491,8 +478,17 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
// TODO(wesm): In future, we may want to reconcile the ids in the stream with
// those found in the schema
for (int i = 0; i < dictionary_memo_.num_fields(); ++i) {
- RETURN_NOT_OK(
- ReadMessageAndValidate(message_reader_.get(), /*allow_null=*/false, &message));
+ RETURN_NOT_OK(message_reader_->ReadNextMessage(&message));
+ if (!message) {
+ /// ARROW-6006: If we fail to find any dictionaries in the stream, then
+ /// it may be that the stream has a schema but no actual data. In such
+ /// case we communicate that we were unable to find the dictionaries
+ /// (but there was no failure otherwise), so the caller can decide what
+ /// to do
+ empty_stream_ = true;
+ break;
+ }
+
if (message->type() != Message::DICTIONARY_BATCH) {
return Status::Invalid(
"IPC stream did not find the expected number of "
@@ -510,9 +506,15 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
RETURN_NOT_OK(ReadInitialDictionaries());
}
+ if (empty_stream_) {
+ // ARROW-6006: Degenerate case where stream contains no data, we do not
+ // bother trying to read a RecordBatch message from the stream
+ *batch = nullptr;
+ return Status::OK();
+ }
+
std::unique_ptr<Message> message;
- RETURN_NOT_OK(
- ReadMessageAndValidate(message_reader_.get(), /*allow_null=*/true, &message));
+ RETURN_NOT_OK(message_reader_->ReadNextMessage(&message));
if (message == nullptr) {
// End of stream
*batch = nullptr;
@@ -537,6 +539,10 @@ class RecordBatchStreamReader::RecordBatchStreamReaderImpl {
bool read_initial_dictionaries_ = false;
+ // Flag to set in case where we fail to observe all dictionaries in a stream,
+ // and so the reader should not attempt to parse any messages
+ bool empty_stream_ = false;
+
DictionaryMemo dictionary_memo_;
std::shared_ptr<Schema> schema_;
};
@@ -785,7 +791,10 @@ Status ReadSchema(io::InputStream* stream, DictionaryMemo* dictionary_memo,
std::shared_ptr<Schema>* out) {
std::unique_ptr<MessageReader> reader = MessageReader::Open(stream);
std::unique_ptr<Message> message;
- RETURN_NOT_OK(ReadMessageAndValidate(reader.get(), /*allow_null=*/false, &message));
+ RETURN_NOT_OK(reader->ReadNextMessage(&message));
+ if (!message) {
+ return Status::Invalid("Tried reading schema message, was null or length 0");
+ }
CHECK_MESSAGE_TYPE(message->type(), Message::SCHEMA);
return ReadSchema(*message, dictionary_memo, out);
}