Merge commit '57b5d7873f4f0953357e9d98e9c60cff8373d7ec'

subtree/rabit/rabit-learn/io/hdfs-inl.h

@@ -0,0 +1,140 @@
+#ifndef RABIT_LEARN_IO_HDFS_INL_H_
+#define RABIT_LEARN_IO_HDFS_INL_H_
+/*!
+ * \file hdfs-inl.h
+ * \brief HDFS I/O
+ * \author Tianqi Chen
+ */
+#include <string>
+#include <vector>
+#include <hdfs.h>
+#include <errno.h>
+#include "./io.h"
+#include "./line_split-inl.h"
+
+/*! \brief io interface */
+namespace rabit {
+namespace io {
+class HDFSStream : public utils::ISeekStream {
+ public:
+  HDFSStream(hdfsFS fs, const char *fname, const char *mode)
+      : fs_(fs), at_end_(false) {
+    int flag;
+    if (!strcmp(mode, "r")) {
+      flag = O_RDONLY;
+    } else if (!strcmp(mode, "w"))  {
+      flag = O_WRONLY;
+    } else if (!strcmp(mode, "a"))  {
+      flag = O_WRONLY | O_APPEND;
+    } else {
+      utils::Error("HDFSStream: unknown flag %s", mode);
+    }
+    fp_ = hdfsOpenFile(fs_, fname, flag, 0, 0, 0);
+    utils::Check(fp_ != NULL,
+                 "HDFSStream: fail to open %s", fname);
+  }
+  virtual ~HDFSStream(void) {
+    this->Close();
+  }
+  virtual size_t Read(void *ptr, size_t size) {
+    tSize nread = hdfsRead(fs_, fp_, ptr, size);
+    if (nread == -1) {
+      int errsv = errno;
+      utils::Error("HDFSStream.Read Error:%s", strerror(errsv));
+    }
+    if (nread == 0) {
+      at_end_ = true;
+    }
+    return static_cast<size_t>(nread);
+  }
+  virtual void Write(const void *ptr, size_t size) {
+    const char *buf = reinterpret_cast<const char*>(ptr);
+    while (size != 0) {
+      tSize nwrite = hdfsWrite(fs_, fp_, buf, size);
+      if (nwrite == -1) {
+        int errsv = errno;
+        utils::Error("HDFSStream.Write Error:%s", strerror(errsv));
+      }
+      size_t sz = static_cast<size_t>(nwrite);
+      buf += sz; size -= sz;
+    }
+  }
+  virtual void Seek(size_t pos) {
+    if (hdfsSeek(fs_, fp_, pos) != 0) {
+      int errsv = errno;
+      utils::Error("HDFSStream.Seek Error:%s", strerror(errsv));
+    }
+  }
+  virtual size_t Tell(void) {
+    tOffset offset = hdfsTell(fs_, fp_);
+    if (offset == -1) {
+      int errsv = errno;
+      utils::Error("HDFSStream.Tell Error:%s", strerror(errsv));
+    }
+    return static_cast<size_t>(offset);
+  }
+  virtual bool AtEnd(void) const {
+    return at_end_;
+  }
+  inline void Close(void) {
+    if (fp_ != NULL) {
+      if (hdfsCloseFile(fs_, fp_) == -1) {
+        int errsv = errno;
+        utils::Error("HDFSStream.Close Error:%s", strerror(errsv));
+      }
+      fp_ = NULL;
+    }
+  }  
+  
+ private:
+  hdfsFS fs_;
+  hdfsFile fp_;
+  bool at_end_;
+};
+
+/*! \brief line split from normal file system */
+class HDFSSplit : public LineSplitBase {
+ public:
+  explicit HDFSSplit(const char *uri, unsigned rank, unsigned nsplit) {
+    fs_ = hdfsConnect("default", 0);
+    std::vector<std::string> paths;
+    LineSplitBase::SplitNames(&paths, uri, "#");
+    // get the files
+    std::vector<size_t> fsize;
+    for (size_t  i = 0; i < paths.size(); ++i) {
+      hdfsFileInfo *info = hdfsGetPathInfo(fs_, paths[i].c_str());
+      if (info->mKind == 'D') {
+        int nentry;
+        hdfsFileInfo *files = hdfsListDirectory(fs_, info->mName, &nentry);
+        for (int i = 0; i < nentry; ++i) {
+          if (files[i].mKind == 'F') {
+            fsize.push_back(files[i].mSize);
+            fnames_.push_back(std::string(files[i].mName));
+          }
+        }
+        hdfsFreeFileInfo(files, nentry);
+      } else {
+        fsize.push_back(info->mSize);
+        fnames_.push_back(std::string(info->mName));
+      }
+      hdfsFreeFileInfo(info, 1);
+    }
+    LineSplitBase::Init(fsize, rank, nsplit);
+  }
+  virtual ~HDFSSplit(void) {}
+  
+ protected:
+  virtual utils::ISeekStream *GetFile(size_t file_index) {
+    utils::Assert(file_index < fnames_.size(), "file index exceed bound"); 
+    return new HDFSStream(fs_, fnames_[file_index].c_str(), "r");
+  }
+
+ private:
+  // hdfs handle
+  hdfsFS fs_;
+  // file names
+  std::vector<std::string> fnames_;
+};
+}  // namespace io
+}  // namespace rabit
+#endif  // RABIT_LEARN_IO_HDFS_INL_H_