threaded_input_split.h

Go to the documentation of this file.

 
#ifndef DMLC_IO_THREADED_INPUT_SPLIT_H_
#define DMLC_IO_THREADED_INPUT_SPLIT_H_
 
#include <dmlc/base.h>
// this code depends on c++11
#if DMLC_ENABLE_STD_THREAD
#include <dmlc/threadediter.h>
#include <algorithm>
#include "./input_split_base.h"
 
namespace dmlc {
namespace io {
class ThreadedInputSplit : public InputSplit {
 public:
  explicit ThreadedInputSplit(InputSplitBase *base, const size_t batch_size)
      : buffer_size_(InputSplitBase::kBufferSize),
        batch_size_(batch_size),
        base_(base), tmp_chunk_(NULL) {
    iter_.set_max_capacity(2);
    // initalize the iterator
    iter_.Init([this](InputSplitBase::Chunk **dptr) {
        if (*dptr == NULL) {
          *dptr = new InputSplitBase::Chunk(buffer_size_);
        }
        return base_->NextBatchEx(*dptr, batch_size_);
      },
      [base]() { base->BeforeFirst(); });
  }
  // destructor
  virtual ~ThreadedInputSplit(void) {
    iter_.Destroy();
    delete tmp_chunk_;
    delete base_;
  }
  virtual void BeforeFirst() {
    iter_.BeforeFirst();
    if (tmp_chunk_ != NULL) {
      iter_.Recycle(&tmp_chunk_);
    }
  }
  virtual void HintChunkSize(size_t chunk_size) {
    buffer_size_ = std::max(chunk_size / sizeof(uint32_t), buffer_size_);
  }
  // implement next record
  virtual bool NextRecord(Blob *out_rec) {
    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) {
    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;
  }
 
  virtual size_t GetTotalSize(void) {
    return base_->GetTotalSize();
  }
 
  virtual void ResetPartition(unsigned part_index, unsigned num_parts) {
    base_->ResetPartition(part_index, num_parts);
    this->BeforeFirst();
  }
 
 private:
  size_t buffer_size_;
  size_t batch_size_;
  InputSplitBase *base_;
  ThreadedIter<InputSplitBase::Chunk> iter_;
  InputSplitBase::Chunk *tmp_chunk_;
};
}  // namespace io
}  // namespace dmlc
#endif  // DMLC_USE_CXX11
#endif  // DMLC_IO_THREADED_INPUT_SPLIT_H_