// ----------------------------------------------------------------------
// RecordBatchStreamReader implementation
class RecordBatchStreamReaderImpl : public RecordBatchStreamReader {
public:
Status Open(std::unique_ptr<MessageReader> message_reader,
const IpcReadOptions& options) {
message_reader_ = std::move(message_reader);
options_ = options;
// Read schema
ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Message> message, ReadNextMessage());
if (!message) {
return Status::Invalid("Tried reading schema message, was null or length 0");
}
RETURN_NOT_OK(UnpackSchemaMessage(*message, options, &dictionary_memo_, &schema_,
&out_schema_, &field_inclusion_mask_,
&swap_endian_));
return Status::OK();
}
Status ReadNext(std::shared_ptr<RecordBatch>* batch) override {
if (!have_read_initial_dictionaries_) {
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();
}
// Continue to read other dictionaries, if any
std::unique_ptr<Message> message;
ARROW_ASSIGN_OR_RAISE(message, ReadNextMessage());
while (message != nullptr && message->type() == MessageType::DICTIONARY_BATCH) {
RETURN_NOT_OK(ReadDictionary(*message));
ARROW_ASSIGN_OR_RAISE(message, ReadNextMessage());
}
if (message == nullptr) {
// End of stream
*batch = nullptr;
return Status::OK();
}
CHECK_HAS_BODY(*message);
ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
return ReadRecordBatchInternal(*message->metadata(), schema_, field_inclusion_mask_,
context, reader.get())
.Value(batch);
}
std::shared_ptr<Schema> schema() const override { return out_schema_; }
ReadStats stats() const override { return stats_; }
private:
Result<std::unique_ptr<Message>> ReadNextMessage() {
ARROW_ASSIGN_OR_RAISE(auto message, message_reader_->ReadNextMessage());
if (message) {
++stats_.num_messages;
switch (message->type()) {
case MessageType::RECORD_BATCH:
++stats_.num_record_batches;
break;
case MessageType::DICTIONARY_BATCH:
++stats_.num_dictionary_batches;
break;
default:
break;
}
}
return std::move(message);
}
// Read dictionary from dictionary batch
Status ReadDictionary(const Message& message) {
DictionaryKind kind;
IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
RETURN_NOT_OK(::arrow::ipc::ReadDictionary(message, context, &kind));
switch (kind) {
case DictionaryKind::New:
break;
case DictionaryKind::Delta:
++stats_.num_dictionary_deltas;
break;
case DictionaryKind::Replacement:
++stats_.num_replaced_dictionaries;
break;
}
return Status::OK();
}
Status ReadInitialDictionaries() {
// We must receive all dictionaries before reconstructing the
// first record batch. Subsequent dictionary deltas modify the memo
std::unique_ptr<Message> message;
// TODO(wesm): In future, we may want to reconcile the ids in the stream with
// those found in the schema
const auto num_dicts = dictionary_memo_.fields().num_dicts();
for (int i = 0; i < num_dicts; ++i) {
ARROW_ASSIGN_OR_RAISE(message, ReadNextMessage());
if (!message) {
if (i == 0) {
/// 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;
} else {
// ARROW-6126, the stream terminated before receiving the expected
// number of dictionaries
return Status::Invalid("IPC stream ended without reading the expected number (",
num_dicts, ") of dictionaries");
}
}
if (message->type() != MessageType::DICTIONARY_BATCH) {
return Status::Invalid("IPC stream did not have the expected number (", num_dicts,
") of dictionaries at the start of the stream");
}
RETURN_NOT_OK(ReadDictionary(*message));
}
have_read_initial_dictionaries_ = true;
return Status::OK();
}
std::unique_ptr<MessageReader> message_reader_;
IpcReadOptions options_;
std::vector<bool> field_inclusion_mask_;
bool have_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;
ReadStats stats_;
DictionaryMemo dictionary_memo_;
std::shared_ptr<Schema> schema_, out_schema_;
bool swap_endian_;
};