file_io.cpp

#include <LightGBM/utils/log.h>
#include <LightGBM/utils/file_io.h>

#include <algorithm>
#include <sstream>
#include <unordered_map>
#ifdef USE_HDFS
#include <hdfs.h>
#endif

namespace LightGBM{

struct LocalFile : VirtualFileReader, VirtualFileWriter {
  LocalFile(const std::string& filename, const std::string& mode) : filename_(filename), mode_(mode) {}
  virtual ~LocalFile() {
    if (file_ != NULL) {
      fclose(file_);
    }
  }

  bool Init() {
    if (file_ == NULL) {
#if _MSC_VER
      fopen_s(&file_, filename_.c_str(), mode_.c_str());
#else
      file_ = fopen(filename_.c_str(), mode_.c_str());
#endif
    }
    return file_ != NULL;
  }

  bool Exists() const {
    LocalFile file(filename_, "rb");
    return file.Init();
  }

  size_t Read(void* buffer, size_t bytes) const {
    return fread(buffer, 1, bytes, file_);
  }

  size_t Write(const void* buffer, size_t bytes) const {
    return fwrite(buffer, bytes, 1, file_) == 1 ? bytes : 0;
  }

private:
  FILE* file_ = NULL;
  const std::string filename_;
  const std::string mode_;
};

const std::string kHdfsProto = "hdfs://";

#ifdef USE_HDFS
struct HdfsFile : VirtualFileReader, VirtualFileWriter {
  HdfsFile(const std::string& filename, int flags) : filename_(filename), flags_(flags) {}
  ~HdfsFile() {
    if (file_ != NULL) {
      hdfsCloseFile(fs_, file_);
    }
  }

  bool Init() {
    if (file_ == NULL) {
      if (fs_ == NULL) {
        fs_ = getHdfsFS(filename_);
      }
      if (fs_ != NULL && (flags_ == O_WRONLY || 0 == hdfsExists(fs_, filename_.c_str()))) {
        file_ = hdfsOpenFile(fs_, filename_.c_str(), flags_, 0, 0, 0);
      }
    }
    return file_ != NULL;
  }

  bool Exists() const {
    if (fs_ == NULL) {
      fs_ = getHdfsFS(filename_);
    }
    return fs_ != NULL && 0 == hdfsExists(fs_, filename_.c_str());
  }

  size_t Read(void* data, size_t bytes) const {
    return FileOperation<void*>(data, bytes, &hdfsRead);
  }

  size_t Write(const void* data, size_t bytes) const {
    return FileOperation<const void*>(data, bytes, &hdfsWrite);
  }

private:
  template <typename BufferType>
  using fileOp = tSize(*)(hdfsFS, hdfsFile, BufferType, tSize);

  template <typename BufferType>
  inline size_t FileOperation(BufferType data, size_t bytes, fileOp<BufferType> op) const {
    char* buffer = (char *)data;
    size_t remain = bytes;
    while (remain != 0) {
      size_t nmax = static_cast<size_t>(std::numeric_limits<tSize>::max());
      tSize ret = op(fs_, file_, buffer, std::min(nmax, remain));
      if (ret > 0) {
        size_t n = static_cast<size_t>(ret);
        remain -= n;
        buffer += n;
      } else if (ret == 0) {
        break;
      } else if (errno != EINTR) {
        Log::Fatal("Failed hdfs file operation [%s]", strerror(errno));
      }
    }
    return bytes - remain;
  }

  static hdfsFS getHdfsFS(const std::string& uri) {
    size_t end = uri.find("/", kHdfsProto.length());
    if (uri.find(kHdfsProto) != 0 || end == std::string::npos) {
      Log::Warning("Bad hdfs uri, no namenode found [%s]", uri.c_str());
      return NULL;
    }
    std::string hostport = uri.substr(kHdfsProto.length(), end - kHdfsProto.length());
    if (fs_cache_.count(hostport) == 0) {
      fs_cache_[hostport] = makeHdfsFs(hostport);
    }
    return fs_cache_[hostport];
  }

  static hdfsFS makeHdfsFs(const std::string& hostport) {
    std::istringstream iss(hostport);
    std::string host;
    tPort port = 0;
    std::getline(iss, host, ':');
    iss >> port;
    hdfsFS fs = iss.eof() ? hdfsConnect(host.c_str(), port) : NULL;
    if (fs == NULL) {
      Log::Warning("Could not connect to hdfs namenode [%s]", hostport.c_str());
    }
    return fs;
  }

  mutable hdfsFS fs_ = NULL;
  hdfsFile file_ = NULL;
  const std::string filename_;
  const int flags_;
  static std::unordered_map<std::string, hdfsFS> fs_cache_;
};

std::unordered_map<std::string, hdfsFS> HdfsFile::fs_cache_ = std::unordered_map<std::string, hdfsFS>();

#define WITH_HDFS(x) x
#else
#define WITH_HDFS(x) Log::Fatal("HDFS Support not enabled.")
#endif // USE_HDFS

std::unique_ptr<VirtualFileReader> VirtualFileReader::Make(const std::string& filename) {
  if (0 == filename.find(kHdfsProto)) {
    WITH_HDFS(return std::unique_ptr<VirtualFileReader>(new HdfsFile(filename, O_RDONLY)));
  } else {
    return std::unique_ptr<VirtualFileReader>(new LocalFile(filename, "rb"));
  }
}

std::unique_ptr<VirtualFileWriter> VirtualFileWriter::Make(const std::string& filename) {
  if (0 == filename.find(kHdfsProto)) {
    WITH_HDFS(return std::unique_ptr<VirtualFileWriter>(new HdfsFile(filename, O_WRONLY)));
  } else {
    return std::unique_ptr<VirtualFileWriter>(new LocalFile(filename, "wb"));
  }
}

bool VirtualFileWriter::Exists(const std::string& filename) {
  if (0 == filename.find(kHdfsProto)) {
    WITH_HDFS(HdfsFile file(filename, O_RDONLY); return file.Exists());
  } else {
      LocalFile file(filename, "rb");
      return file.Exists();
  }
}

}  // namespace LightGBM