Caffe2 - C++ API
A deep learning, cross platform ML framework
leveldb.cc
1 #include "caffe2/core/db.h"
2 #include "caffe2/core/logging.h"
3 #include "caffe2/core/flags.h"
4 #include "leveldb/db.h"
5 #include "leveldb/write_batch.h"
6 
7 CAFFE2_DEFINE_int(caffe2_leveldb_block_size, 65536,
8  "The caffe2 leveldb block size when writing a leveldb.");
9 
10 namespace caffe2 {
11 namespace db {
12 
13 class LevelDBCursor : public Cursor {
14  public:
15  explicit LevelDBCursor(leveldb::DB* db)
16  : iter_(db->NewIterator(leveldb::ReadOptions())) {
17  SeekToFirst();
18  }
19  ~LevelDBCursor() {}
20  void Seek(const string& key) override { iter_->Seek(key); }
21  bool SupportsSeek() override { return true; }
22  void SeekToFirst() override { iter_->SeekToFirst(); }
23  void Next() override { iter_->Next(); }
24  string key() override { return iter_->key().ToString(); }
25  string value() override { return iter_->value().ToString(); }
26  bool Valid() override { return iter_->Valid(); }
27 
28  private:
29  std::unique_ptr<leveldb::Iterator> iter_;
30 };
31 
33  public:
34  explicit LevelDBTransaction(leveldb::DB* db) : db_(db) {
35  CAFFE_ENFORCE(db_);
36  batch_.reset(new leveldb::WriteBatch());
37  }
38  ~LevelDBTransaction() { Commit(); }
39  void Put(const string& key, const string& value) override {
40  batch_->Put(key, value);
41  }
42  void Commit() override {
43  leveldb::Status status = db_->Write(leveldb::WriteOptions(), batch_.get());
44  batch_.reset(new leveldb::WriteBatch());
45  CAFFE_ENFORCE(
46  status.ok(),
47  "Failed to write batch to leveldb. ", status.ToString());
48  }
49 
50  private:
51  leveldb::DB* db_;
52  std::unique_ptr<leveldb::WriteBatch> batch_;
53 
54  DISABLE_COPY_AND_ASSIGN(LevelDBTransaction);
55 };
56 
57 class LevelDB : public DB {
58  public:
59  LevelDB(const string& source, Mode mode) : DB(source, mode) {
60  leveldb::Options options;
61  options.block_size = FLAGS_caffe2_leveldb_block_size;
62  options.write_buffer_size = 268435456;
63  options.max_open_files = 100;
64  options.error_if_exists = mode == NEW;
65  options.create_if_missing = mode != READ;
66  leveldb::DB* db_temp;
67  leveldb::Status status = leveldb::DB::Open(options, source, &db_temp);
68  CAFFE_ENFORCE(
69  status.ok(),
70  "Failed to open leveldb ", source, ". ", status.ToString());
71  db_.reset(db_temp);
72  VLOG(1) << "Opened leveldb " << source;
73  }
74 
75  void Close() override { db_.reset(); }
76  unique_ptr<Cursor> NewCursor() override {
77  return make_unique<LevelDBCursor>(db_.get());
78  }
79  unique_ptr<Transaction> NewTransaction() override {
80  return make_unique<LevelDBTransaction>(db_.get());
81  }
82 
83  private:
84  std::unique_ptr<leveldb::DB> db_;
85 };
86 
87 REGISTER_CAFFE2_DB(LevelDB, LevelDB);
88 // For lazy-minded, one can also call with lower-case name.
89 REGISTER_CAFFE2_DB(leveldb, LevelDB);
90 
91 } // namespace db
92 } // namespace caffe2
An abstract class for accessing a database of key-value pairs.
Definition: db.h:80
string key() override
Returns the current key.
Definition: leveldb.cc:24
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...
Definition: leveldb.cc:26
void SeekToFirst() override
Seek to the first key in the database.
Definition: leveldb.cc:22
An abstract class for the cursor of the database while reading.
Definition: db.h:22
unique_ptr< Cursor > NewCursor() override
Returns a cursor to read the database.
Definition: leveldb.cc:76
unique_ptr< Transaction > NewTransaction() override
Returns a transaction to write data to the database.
Definition: leveldb.cc:79
void Seek(const string &key) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
Definition: leveldb.cc:20
string value() override
Returns the current value.
Definition: leveldb.cc:25
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
void Commit() override
Commits the current writes.
Definition: leveldb.cc:42
void Close() override
Closes the database.
Definition: leveldb.cc:75
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
Definition: leveldb.cc:39
Commandline flags support for Caffe2.
An abstract class for the current database transaction while writing.
Definition: db.h:61
void Next() override
Go to the next location in the database.
Definition: leveldb.cc:23