input_split_shuffle.h

Go to the documentation of this file.

#ifndef DMLC_INPUT_SPLIT_SHUFFLE_H_
#define DMLC_INPUT_SPLIT_SHUFFLE_H_

#include <cstdio>
#include <cstring>
#include <vector>
#include <string>
#include <algorithm>
#include <memory>

namespace dmlc {
class InputSplitShuffle : public InputSplit {
 public:
  // destructor
  virtual ~InputSplitShuffle(void) { source_.reset(); }
  // implement BeforeFirst
  virtual void BeforeFirst(void) {
    if (num_shuffle_parts_ > 1) {
      std::shuffle(shuffle_indexes_.begin(), shuffle_indexes_.end(), trnd_);
      int idx = shuffle_indexes_[0] + part_index_ * num_shuffle_parts_;
      source_->ResetPartition(idx, num_parts_ * num_shuffle_parts_);
      cur_shuffle_idx_ = 0;
    } else {
      source_->BeforeFirst();
    }
  }
  virtual void HintChunkSize(size_t chunk_size) {
    source_->HintChunkSize(chunk_size);
  }
  virtual size_t GetTotalSize(void) {
    return source_->GetTotalSize();
  }
  // implement next record
  virtual bool NextRecord(Blob *out_rec) {
    if (num_shuffle_parts_ > 1) {
      if (!source_->NextRecord(out_rec)) {
        if (cur_shuffle_idx_ == num_shuffle_parts_ - 1) {
          return false;
        }
        ++cur_shuffle_idx_;
        int idx =
            shuffle_indexes_[cur_shuffle_idx_] + part_index_ * num_shuffle_parts_;
        source_->ResetPartition(idx, num_parts_ * num_shuffle_parts_);
        return NextRecord(out_rec);
      } else {
        return true;
      }
    } else {
      return source_->NextRecord(out_rec);
    }
  }
  // implement next chunk
  virtual bool NextChunk(Blob* out_chunk) {
    if (num_shuffle_parts_ > 1) {
      if (!source_->NextChunk(out_chunk)) {
        if (cur_shuffle_idx_ == num_shuffle_parts_ - 1) {
          return false;
        }
        ++cur_shuffle_idx_;
        int idx =
            shuffle_indexes_[cur_shuffle_idx_] + part_index_ * num_shuffle_parts_;
        source_->ResetPartition(idx, num_parts_ * num_shuffle_parts_);
        return NextChunk(out_chunk);
      } else {
        return true;
      }
    } else {
      return source_->NextChunk(out_chunk);
    }
  }
  // implement ResetPartition.
  virtual void ResetPartition(unsigned rank, unsigned nsplit) {
    CHECK(nsplit == num_parts_) << "num_parts is not consistent!";
    int idx = shuffle_indexes_[0] + rank * num_shuffle_parts_;
    source_->ResetPartition(idx, nsplit * num_shuffle_parts_);
    cur_shuffle_idx_ = 0;
  }
  InputSplitShuffle(const char* uri,
                    unsigned part_index,
                    unsigned num_parts,
                    const char* type,
                    unsigned num_shuffle_parts,
                    int shuffle_seed)
      : part_index_(part_index),
        num_parts_(num_parts),
        num_shuffle_parts_(num_shuffle_parts),
        cur_shuffle_idx_(0) {
    for (unsigned i = 0; i < num_shuffle_parts_; i++) {
      shuffle_indexes_.push_back(i);
    }
    trnd_.seed(kRandMagic_ + part_index_ + num_parts_ + num_shuffle_parts_ +
               shuffle_seed);
    std::shuffle(shuffle_indexes_.begin(), shuffle_indexes_.end(), trnd_);
    int idx = shuffle_indexes_[cur_shuffle_idx_] + part_index_ * num_shuffle_parts_;
    source_.reset(
        InputSplit::Create(uri, idx , num_parts_ * num_shuffle_parts_, type));
  }
  static InputSplit* Create(const char* uri,
                            unsigned part_index,
                            unsigned num_parts,
                            const char* type,
                            unsigned num_shuffle_parts,
                            int shuffle_seed) {
    CHECK(num_shuffle_parts > 0) << "number of shuffle parts should be greater than zero!";
    return new InputSplitShuffle(
        uri, part_index, num_parts, type, num_shuffle_parts, shuffle_seed);
  }

 private:
  // magic nyumber for seed
  static const int kRandMagic_ = 666;
  std::mt19937 trnd_;
  std::unique_ptr<InputSplit> source_;
  unsigned part_index_;
  unsigned num_parts_;
  unsigned num_shuffle_parts_;
  unsigned cur_shuffle_idx_;
  std::vector<int> shuffle_indexes_;
};
}  // namespace dmlc
#endif  // DMLC_INPUT_SPLIT_SHUFFLE_H_