Arrow 之 IPC

RecordBatchStreamReader

// ----------------------------------------------------------------------
// 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_;
};