recordio.h

Go to the documentation of this file.

#ifndef DMLC_RECORDIO_H_
#define DMLC_RECORDIO_H_
#include <cstring>
#include <string>
#include "./io.h"
#include "./logging.h"

namespace dmlc {
class RecordIOWriter {
 public:
  static const uint32_t kMagic = 0xced7230a;
  inline static uint32_t EncodeLRec(uint32_t cflag, uint32_t length) {
    return (cflag << 29U) | length;
  }
  inline static uint32_t DecodeFlag(uint32_t rec) {
    return (rec >> 29U) & 7U;
  }
  inline static uint32_t DecodeLength(uint32_t rec) {
    return rec & ((1U << 29U) - 1U);
  }
  explicit RecordIOWriter(Stream *stream)
      : stream_(stream), seek_stream_(dynamic_cast<SeekStream*>(stream)),
        except_counter_(0) {
    CHECK(sizeof(uint32_t) == 4) << "uint32_t needs to be 4 bytes";
  }
  void WriteRecord(const void *buf, size_t size);
  inline void WriteRecord(const std::string &data) {
    this->WriteRecord(data.c_str(), data.length());
  }
  inline size_t except_counter(void) const {
    return except_counter_;
  }

  inline size_t Tell(void) {
    CHECK(seek_stream_ != NULL) << "The input stream is not seekable";
    return seek_stream_->Tell();
  }

 private:
  Stream *stream_;
  SeekStream *seek_stream_;
  size_t except_counter_;
};
class RecordIOReader {
 public:
  explicit RecordIOReader(Stream *stream)
      : stream_(stream), seek_stream_(dynamic_cast<SeekStream*>(stream)),
        end_of_stream_(false) {
    CHECK(sizeof(uint32_t) == 4) << "uint32_t needs to be 4 bytes";
  }
  bool NextRecord(std::string *out_rec);

  inline void Seek(size_t pos) {
    CHECK(seek_stream_ != NULL) << "The input stream is not seekable";
    seek_stream_->Seek(pos);
  }

  inline size_t Tell(void) {
    CHECK(seek_stream_ != NULL) << "The input stream is not seekable";
    return seek_stream_->Tell();
  }

 private:
  Stream *stream_;
  SeekStream *seek_stream_;
  bool end_of_stream_;
};

class RecordIOChunkReader {
 public:
  explicit RecordIOChunkReader(InputSplit::Blob chunk,
                               unsigned part_index = 0,
                               unsigned num_parts = 1);
  bool NextRecord(InputSplit::Blob *out_rec);

 private:
  std::string temp_;
  char *pbegin_, *pend_;
};

}  // namespace dmlc
#endif  // DMLC_RECORDIO_H_