input_split_base.h

Go to the documentation of this file.

 
#ifndef DMLC_IO_INPUT_SPLIT_BASE_H_
#define DMLC_IO_INPUT_SPLIT_BASE_H_
 
#include <dmlc/io.h>
#include <dmlc/filesystem.h>
#include <cstdio>
#include <cstring>
#include <vector>
#include <string>
#include <algorithm>
 
namespace dmlc {
namespace io {
class InputSplitBase : public InputSplit {
 public:
  struct Chunk {
    char *begin;
    char *end;
    std::vector<uint32_t> data;
    explicit Chunk(size_t buffer_size)
        : begin(NULL), end(NULL),
          data(buffer_size + 1) {}
    // load chunk from split
    bool Load(InputSplitBase *split, size_t buffer_size);
    // append to chunk
    bool Append(InputSplitBase *split, size_t buffer_size);
  };
  // 16 MB
  static const size_t kBufferSize = 2UL << 20UL;
  // destructor
  virtual ~InputSplitBase(void);
  // implement BeforeFirst
  virtual void BeforeFirst(void);
  virtual void HintChunkSize(size_t chunk_size) {
    buffer_size_ = std::max(chunk_size / sizeof(uint32_t), buffer_size_);
  }
  virtual size_t GetTotalSize(void) {
    return file_offset_.back();
  }
  // implement next record
  virtual bool NextRecord(Blob *out_rec) {
    while (!ExtractNextRecord(out_rec, &tmp_chunk_)) {
      if (!NextChunkEx(&tmp_chunk_)) return false;
    }
    return true;
  }
  // implement next chunk
  virtual bool NextChunk(Blob *out_chunk) {
    while (!ExtractNextChunk(out_chunk, &tmp_chunk_)) {
      if (!NextChunkEx(&tmp_chunk_)) return false;
    }
    return true;
  }
  // implement ResetPartition.
  virtual void ResetPartition(unsigned rank, unsigned nsplit);
  virtual bool ReadChunk(void *buf, size_t *size);
  bool ExtractNextChunk(Blob *out_rchunk, Chunk *chunk);
  virtual bool ExtractNextRecord(Blob *out_rec, Chunk *chunk) = 0;
  virtual bool IsTextParser(void) = 0;
  virtual bool NextChunkEx(Chunk *chunk) {
    if (!chunk->Load(this, buffer_size_)) return false;
    return true;
  }
  virtual bool NextBatchEx(Chunk *chunk, size_t /*n_records*/) {
    return NextChunkEx(chunk);
  }
 
 protected:
  FileSystem *filesys_;
  std::vector<size_t> file_offset_;
  size_t offset_curr_;
  size_t offset_begin_;
  size_t offset_end_;
  std::vector<FileInfo> files_;
  SeekStream *fs_;
  size_t file_ptr_;
  size_t file_ptr_end_;
  Chunk tmp_chunk_;
  size_t buffer_size_;
  // constructor
  InputSplitBase()
      : fs_(NULL),
        tmp_chunk_(kBufferSize),
        buffer_size_(kBufferSize),
        align_bytes_(8) {}
  void Init(FileSystem *fs,
            const char *uri,
            size_t align_bytes,
            const bool recurse_directories = false);
  // to be implemented by child class
  virtual size_t SeekRecordBegin(Stream *fi) = 0;
  virtual const char*
  FindLastRecordBegin(const char *begin, const char *end) = 0;
 
  std::vector<URI> ConvertToURIs(const std::string& uri);
  size_t Read(void *ptr, size_t size);
 
 private:
  size_t align_bytes_;
  std::string overflow_;
  void InitInputFileInfo(const std::string& uri,
                         const bool recurse_directories);
  std::string StripEnd(std::string str, char ch);
};
}  // namespace io
}  // namespace dmlc
#endif  // DMLC_IO_INPUT_SPLIT_BASE_H_