data_shard.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
* 
*   http://www.apache.org/licenses/LICENSE-2.0
* 
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*
*************************************************************/

#ifndef SINGA_UTILS_DATA_SHARD_H_
#define SINGA_UTILS_DATA_SHARD_H_

#include <google/protobuf/message.h>
#include <fstream>
#include <string>
#include <unordered_set>

namespace singa {

class DataShard {
 public:
  enum {
    // read only mode used in training
    kRead = 0,
    // write mode used in creating shard (will overwrite previous one)
    kCreate = 1,
    // append mode, e.g. used when previous creating crashes
    kAppend = 2
  };

  DataShard(const std::string& folder, int mode);
  DataShard(const std::string& folder, int mode, int capacity);
  ~DataShard();

  bool Next(std::string* key, google::protobuf::Message* val);
  bool Next(std::string* key, std::string* val);
  bool Insert(const std::string& key, const google::protobuf::Message& tuple);
  bool Insert(const std::string& key, const std::string& tuple);
  void SeekToFirst();
  void Flush();
  int Count();
  inline std::string path() { return path_; }

 protected:
  int Next(std::string* key);
  int PrepareForAppend(const std::string& path);
  bool PrepareNextField(int size);

 private:
  char mode_ = 0;
  std::string path_ = "";
  // either ifstream or ofstream
  std::fstream fdat_;
  // to avoid replicated record
  std::unordered_set<std::string> keys_;
  // internal buffer
  char* buf_ = nullptr;
  // offset inside the buf_
  int offset_ = 0;
  // allocated bytes for the buf_
  int capacity_ = 0;
  // bytes in buf_, used in reading
  int bufsize_ = 0;
};

}  // namespace singa

#endif  // SINGA_UTILS_DATA_SHARD_H_