cached_input_split.h

Go to the documentation of this file.

 
#ifndef DMLC_IO_CACHED_INPUT_SPLIT_H_
#define DMLC_IO_CACHED_INPUT_SPLIT_H_
 
#include <dmlc/base.h>
// this code depends on c++11
 
#if DMLC_ENABLE_STD_THREAD
#include <dmlc/threadediter.h>
#include <string>
#include <algorithm>
#include "./input_split_base.h"
 
namespace dmlc {
namespace io {
class CachedInputSplit : public InputSplit {
 public:
  CachedInputSplit(InputSplitBase *base,
                   const char *cache_file,
                   bool reuse_exist_cache = true)
      : buffer_size_(InputSplitBase::kBufferSize),
        cache_file_(cache_file),
        fo_(NULL), fi_(NULL),
        base_(base), tmp_chunk_(NULL),
        iter_preproc_(NULL) {
    if (reuse_exist_cache) {
      if (!this->InitCachedIter()) {
        this->InitPreprocIter();
      }
    } else {
      this->InitPreprocIter();
    }
  }
  // destructor
  virtual ~CachedInputSplit(void) {
    // NOTE delete can handle NULL ptr
    // deletion order matters
    delete iter_preproc_;
    delete fo_;
    iter_cached_.Destroy();
    delete tmp_chunk_;
    delete base_;
    delete fi_;
  }
  virtual void BeforeFirst(void) {
    // if preprocessing did not end
    // pull data from preprocessing module
    if (iter_preproc_ != NULL) {
      if (tmp_chunk_ != NULL) {
        iter_preproc_->Recycle(&tmp_chunk_);
      }
      while (iter_preproc_->Next(&tmp_chunk_)) {
        iter_preproc_->Recycle(&tmp_chunk_);
      }
      // finalize the push out process
      delete iter_preproc_;
      delete fo_;
      iter_preproc_ = NULL;
      fo_ = NULL;
      CHECK(this->InitCachedIter())
          << "Failed to initialize CachedIter";
    } else {
      iter_cached_.BeforeFirst();
    }
    if (tmp_chunk_ != NULL) {
      iter_cached_.Recycle(&tmp_chunk_);
    }
  }
  virtual void ResetPartition(unsigned /*part_index*/, unsigned /*num_parts*/) {
    LOG(FATAL) << "ResetPartition is not supported in CachedInputSplit";
  }
  virtual void HintChunkSize(size_t chunk_size) {
    buffer_size_ = std::max(chunk_size / sizeof(uint32_t), buffer_size_);
  }
  virtual size_t GetTotalSize(void) {
    return base_->GetTotalSize();
  }
  // implement next record
  virtual bool NextRecord(Blob *out_rec) {
    auto *iter = iter_preproc_ != NULL ? iter_preproc_ : &iter_cached_;
    if (tmp_chunk_ == NULL) {
      if (!iter->Next(&tmp_chunk_)) return false;
    }
    while (!base_->ExtractNextRecord(out_rec, tmp_chunk_)) {
      iter->Recycle(&tmp_chunk_);
      if (!iter->Next(&tmp_chunk_)) return false;
    }
    return true;
  }
  // implement next chunk
  virtual bool NextChunk(Blob *out_chunk) {
    auto *iter = iter_preproc_ != NULL ? iter_preproc_ : &iter_cached_;
    if (tmp_chunk_ == NULL) {
      if (!iter->Next(&tmp_chunk_)) return false;
    }
    while (!base_->ExtractNextChunk(out_chunk, tmp_chunk_)) {
      iter->Recycle(&tmp_chunk_);
      if (!iter->Next(&tmp_chunk_)) return false;
    }
    return true;
  }
 
 private:
  size_t buffer_size_;
  std::string cache_file_;
  dmlc::Stream *fo_;
  dmlc::SeekStream *fi_;
  InputSplitBase *base_;
  InputSplitBase::Chunk *tmp_chunk_;
  ThreadedIter<InputSplitBase::Chunk> *iter_preproc_;
  ThreadedIter<InputSplitBase::Chunk> iter_cached_;
  inline void InitPreprocIter(void);
  inline bool InitCachedIter(void);
};
 
inline void CachedInputSplit:: InitPreprocIter(void) {
  fo_ = dmlc::Stream::Create(cache_file_.c_str(), "w");
  iter_preproc_ = new ThreadedIter<InputSplitBase::Chunk>();
  iter_preproc_->set_max_capacity(16);
  iter_preproc_->Init([this](InputSplitBase::Chunk **dptr) {
      if (*dptr == NULL) {
        *dptr = new InputSplitBase::Chunk(buffer_size_);
      }
      auto *p = *dptr;
      if (!base_->NextChunkEx(p)) return false;
      // after loading, save to disk
      size_t size = p->end - p->begin;
      fo_->Write(&size, sizeof(size));
      fo_->Write(p->begin, size);
      return true;
    });
}
 
inline bool CachedInputSplit::InitCachedIter(void) {
  fi_ = dmlc::SeekStream::CreateForRead(cache_file_.c_str(), true);
  if (fi_ == NULL) return false;
  iter_cached_.Init([this](InputSplitBase::Chunk **dptr) {
      if (*dptr == NULL) {
        *dptr = new InputSplitBase::Chunk(buffer_size_);
      }
      auto *p = *dptr;
      // read data from cache file
      size_t size;
      size_t nread = fi_->Read(&size, sizeof(size));
      if (nread == 0) return false;
      CHECK(nread == sizeof(size))
          << cache_file_ << " has invalid cache file format";
      p->data.resize(size / sizeof(size_t) + 1);
      p->begin = reinterpret_cast<char*>(BeginPtr(p->data));
      p->end = p->begin + size;
      CHECK(fi_->Read(p->begin, size) == size)
          << cache_file_ << " has invalid cache file format";
      return true;
    },
    [this]() { fi_->Seek(0); });
  return true;
}
}  // namespace io
}  // namespace dmlc
#endif  // DMLC_USE_CXX11
#endif  // DMLC_IO_CACHED_INPUT_SPLIT_H_