xgboost
io.h
Go to the documentation of this file.
1 
8 #ifndef XGBOOST_COMMON_IO_H_
9 #define XGBOOST_COMMON_IO_H_
10 
11 #include <dmlc/io.h>
12 #include <rabit/rabit.h>
13 #include <string>
14 #include <cstring>
15 
16 namespace xgboost {
17 namespace common {
20 
25 class PeekableInStream : public dmlc::Stream {
26  public:
27  explicit PeekableInStream(dmlc::Stream* strm)
28  : strm_(strm), buffer_ptr_(0) {}
29 
30  size_t Read(void* dptr, size_t size) override {
31  size_t nbuffer = buffer_.length() - buffer_ptr_;
32  if (nbuffer == 0) return strm_->Read(dptr, size);
33  if (nbuffer < size) {
34  std::memcpy(dptr, dmlc::BeginPtr(buffer_) + buffer_ptr_, nbuffer);
35  buffer_ptr_ += nbuffer;
36  return nbuffer + strm_->Read(reinterpret_cast<char*>(dptr) + nbuffer,
37  size - nbuffer);
38  } else {
39  std::memcpy(dptr, dmlc::BeginPtr(buffer_) + buffer_ptr_, size);
40  buffer_ptr_ += size;
41  return size;
42  }
43  }
44 
45  size_t PeekRead(void* dptr, size_t size) {
46  size_t nbuffer = buffer_.length() - buffer_ptr_;
47  if (nbuffer < size) {
48  buffer_ = buffer_.substr(buffer_ptr_, buffer_.length());
49  buffer_ptr_ = 0;
50  buffer_.resize(size);
51  size_t nadd = strm_->Read(dmlc::BeginPtr(buffer_) + nbuffer, size - nbuffer);
52  buffer_.resize(nbuffer + nadd);
53  std::memcpy(dptr, dmlc::BeginPtr(buffer_), buffer_.length());
54  return buffer_.length();
55  } else {
56  std::memcpy(dptr, dmlc::BeginPtr(buffer_) + buffer_ptr_, size);
57  return size;
58  }
59  }
60 
61  void Write(const void* dptr, size_t size) override {
62  LOG(FATAL) << "Not implemented";
63  }
64 
65  private:
67  dmlc::Stream *strm_;
69  size_t buffer_ptr_;
71  std::string buffer_;
72 };
73 } // namespace common
74 } // namespace xgboost
75 #endif // XGBOOST_COMMON_IO_H_
rabit::utils::MemoryBufferStream MemoryBufferStream
Definition: io.h:19
size_t PeekRead(void *dptr, size_t size)
Definition: io.h:45
size_t Read(void *dptr, size_t size) override
Definition: io.h:30
void Write(const void *dptr, size_t size) override
Definition: io.h:61
Input stream that support additional PeekRead operation, besides read.
Definition: io.h:25
namespace of xgboost
Definition: base.h:79
rabit::utils::MemoryFixSizeBuffer MemoryFixSizeBuffer
Definition: io.h:18
PeekableInStream(dmlc::Stream *strm)
Definition: io.h:27