Added first_key and last_key to dataset

diff --git a/include/caffe/dataset.hpp b/include/caffe/dataset.hpp
index efe3ffd..90deb31 100644 (file)
--- a/include/caffe/dataset.hpp
+++ b/include/caffe/dataset.hpp
@@ -135,6 +135,8 @@ class Dataset {
   virtual bool open(const string& filename, Mode mode) = 0;
   virtual bool put(const K& key, const V& value) = 0;
   virtual bool get(const K& key, V* value) = 0;
+  virtual bool first_key(K* key) = 0;
+  virtual bool last_key(K* key) = 0;
   virtual bool commit() = 0;
   virtual void close() = 0;
 

diff --git a/include/caffe/leveldb_dataset.hpp b/include/caffe/leveldb_dataset.hpp
index 90f92d9..f6c57fe 100644 (file)
--- a/include/caffe/leveldb_dataset.hpp
+++ b/include/caffe/leveldb_dataset.hpp
@@ -29,6 +29,8 @@ class LeveldbDataset : public Dataset<K, V, KCoder, VCoder> {
   bool open(const string& filename, Mode mode);
   bool put(const K& key, const V& value);
   bool get(const K& key, V* value);
+  bool first_key(K* key);
+  bool last_key(K* key);
   bool commit();
   void close();
 

diff --git a/include/caffe/lmdb_dataset.hpp b/include/caffe/lmdb_dataset.hpp
index 8817df7..71b3224 100644 (file)
--- a/include/caffe/lmdb_dataset.hpp
+++ b/include/caffe/lmdb_dataset.hpp
@@ -33,6 +33,8 @@ class LmdbDataset : public Dataset<K, V, KCoder, VCoder> {
   bool open(const string& filename, Mode mode);
   bool put(const K& key, const V& value);
   bool get(const K& key, V* value);
+  bool first_key(K* key);
+  bool last_key(K* key);
   bool commit();
   void close();
 

diff --git a/src/caffe/leveldb_dataset.cpp b/src/caffe/leveldb_dataset.cpp
index af95604..53df985 100644 (file)
--- a/src/caffe/leveldb_dataset.cpp
+++ b/src/caffe/leveldb_dataset.cpp
@@ -107,6 +107,30 @@ bool LeveldbDataset<K, V, KCoder, VCoder>::get(const K& key, V* value) {
 }
 
 template <typename K, typename V, typename KCoder, typename VCoder>
+bool LeveldbDataset<K, V, KCoder, VCoder>::first_key(K* key) {
+  DLOG(INFO) << "LevelDB: First key";
+
+  CHECK_NOTNULL(db_.get());
+  shared_ptr<leveldb::Iterator> iter(db_->NewIterator(leveldb::ReadOptions()));
+  iter->SeekToFirst();
+  CHECK(iter->Valid());
+  const leveldb::Slice& key_slice = iter->key();
+  return KCoder::deserialize(key_slice.data(), key_slice.size(), key);
+}
+
+template <typename K, typename V, typename KCoder, typename VCoder>
+bool LeveldbDataset<K, V, KCoder, VCoder>::last_key(K* key) {
+  DLOG(INFO) << "LevelDB: Last key";
+
+  CHECK_NOTNULL(db_.get());
+  shared_ptr<leveldb::Iterator> iter(db_->NewIterator(leveldb::ReadOptions()));
+  iter->SeekToLast();
+  CHECK(iter->Valid());
+  const leveldb::Slice& key_slice = iter->key();
+  return KCoder::deserialize(key_slice.data(), key_slice.size(), key);
+}
+
+template <typename K, typename V, typename KCoder, typename VCoder>
 bool LeveldbDataset<K, V, KCoder, VCoder>::commit() {
   DLOG(INFO) << "LevelDB: Commit";
 

diff --git a/src/caffe/lmdb_dataset.cpp b/src/caffe/lmdb_dataset.cpp
index ca96843..23852f0 100644 (file)
--- a/src/caffe/lmdb_dataset.cpp
+++ b/src/caffe/lmdb_dataset.cpp
@@ -155,6 +155,70 @@ bool LmdbDataset<K, V, KCoder, VCoder>::get(const K& key, V* value) {
 }
 
 template <typename K, typename V, typename KCoder, typename VCoder>
+bool LmdbDataset<K, V, KCoder, VCoder>::first_key(K* key) {
+  DLOG(INFO) << "LMDB: First key";
+
+  int retval;
+
+  MDB_txn* iter_txn;
+
+  retval = mdb_txn_begin(env_, NULL, MDB_RDONLY, &iter_txn);
+  CHECK_EQ(MDB_SUCCESS, retval) << "mdb_txn_begin failed "
+      << mdb_strerror(retval);
+
+  MDB_cursor* cursor;
+  retval = mdb_cursor_open(iter_txn, dbi_, &cursor);
+  CHECK_EQ(retval, MDB_SUCCESS) << mdb_strerror(retval);
+  MDB_val mdbkey;
+  MDB_val mdbval;
+  retval = mdb_cursor_get(cursor, &mdbkey, &mdbval, MDB_FIRST);
+  CHECK_EQ(retval, MDB_SUCCESS) << mdb_strerror(retval);
+
+  mdb_cursor_close(cursor);
+  mdb_txn_abort(iter_txn);
+
+  if (!KCoder::deserialize(reinterpret_cast<char*>(mdbkey.mv_data),
+      mdbkey.mv_size, key)) {
+    LOG(ERROR) << "failed to deserialize key";
+    return false;
+  }
+
+  return true;
+}
+
+template <typename K, typename V, typename KCoder, typename VCoder>
+bool LmdbDataset<K, V, KCoder, VCoder>::last_key(K* key) {
+  DLOG(INFO) << "LMDB: Last key";
+
+  int retval;
+
+  MDB_txn* iter_txn;
+
+  retval = mdb_txn_begin(env_, NULL, MDB_RDONLY, &iter_txn);
+  CHECK_EQ(MDB_SUCCESS, retval) << "mdb_txn_begin failed "
+      << mdb_strerror(retval);
+
+  MDB_cursor* cursor;
+  retval = mdb_cursor_open(iter_txn, dbi_, &cursor);
+  CHECK_EQ(retval, MDB_SUCCESS) << mdb_strerror(retval);
+  MDB_val mdbkey;
+  MDB_val mdbval;
+  retval = mdb_cursor_get(cursor, &mdbkey, &mdbval, MDB_LAST);
+  CHECK_EQ(retval, MDB_SUCCESS) << mdb_strerror(retval);
+
+  mdb_cursor_close(cursor);
+  mdb_txn_abort(iter_txn);
+
+  if (!KCoder::deserialize(reinterpret_cast<char*>(mdbkey.mv_data),
+      mdbkey.mv_size, key)) {
+    LOG(ERROR) << "failed to deserialize key";
+    return false;
+  }
+
+  return true;
+}
+
+template <typename K, typename V, typename KCoder, typename VCoder>
 bool LmdbDataset<K, V, KCoder, VCoder>::commit() {
   DLOG(INFO) << "LMDB: Commit";
 

diff --git a/src/caffe/test/test_dataset.cpp b/src/caffe/test/test_dataset.cpp
index 7d0f6b3..e5bd1d2 100644 (file)
--- a/src/caffe/test/test_dataset.cpp
+++ b/src/caffe/test/test_dataset.cpp
@@ -277,6 +277,120 @@ TYPED_TEST(DatasetTest, TestKeys) {
   EXPECT_FALSE(this->equals(keys.at(0), keys.at(1)));
 }
 
+TYPED_TEST(DatasetTest, TestFirstKey) {
+  UNPACK_TYPES;
+
+  string name = this->DBName();
+  shared_ptr<Dataset<string, value_type> > dataset =
+      DatasetFactory<string, value_type>(backend);
+  EXPECT_TRUE(dataset->open(name, Dataset<string, value_type>::New));
+
+  value_type value = this->TestValue();
+
+  string key1 = "01";
+  EXPECT_TRUE(dataset->put(key1, value));
+
+  string key2 = "02";
+  EXPECT_TRUE(dataset->put(key2, value));
+
+  string key3 = "03";
+  EXPECT_TRUE(dataset->put(key3, value));
+
+  EXPECT_TRUE(dataset->commit());
+
+  string first_key;
+  dataset->first_key(&first_key);
+
+  EXPECT_TRUE(this->equals(first_key, key1));
+}
+
+TYPED_TEST(DatasetTest, TestLastKey) {
+  UNPACK_TYPES;
+
+  string name = this->DBName();
+  shared_ptr<Dataset<string, value_type> > dataset =
+      DatasetFactory<string, value_type>(backend);
+  EXPECT_TRUE(dataset->open(name, Dataset<string, value_type>::New));
+
+  value_type value = this->TestValue();
+
+  string key1 = "01";
+  EXPECT_TRUE(dataset->put(key1, value));
+
+  string key2 = "02";
+  EXPECT_TRUE(dataset->put(key2, value));
+
+  string key3 = "03";
+  EXPECT_TRUE(dataset->put(key3, value));
+
+  EXPECT_TRUE(dataset->commit());
+
+  string last_key;
+  dataset->last_key(&last_key);
+
+  EXPECT_TRUE(this->equals(last_key, key3));
+}
+
+TYPED_TEST(DatasetTest, TestFirstLastKeys) {
+  UNPACK_TYPES;
+
+  string name = this->DBName();
+  shared_ptr<Dataset<string, value_type> > dataset =
+      DatasetFactory<string, value_type>(backend);
+  EXPECT_TRUE(dataset->open(name, Dataset<string, value_type>::New));
+
+  value_type value = this->TestValue();
+
+  string key1 = "01";
+  EXPECT_TRUE(dataset->put(key1, value));
+
+  string key2 = "02";
+  EXPECT_TRUE(dataset->put(key2, value));
+
+  string key3 = "03";
+  EXPECT_TRUE(dataset->put(key3, value));
+
+  EXPECT_TRUE(dataset->commit());
+
+  string first_key;
+  dataset->first_key(&first_key);
+  string last_key;
+  dataset->last_key(&last_key);
+
+  EXPECT_TRUE(this->equals(first_key, key1));
+  EXPECT_TRUE(this->equals(last_key, key3));
+}
+
+TYPED_TEST(DatasetTest, TestFirstLastKeysUnOrdered) {
+  UNPACK_TYPES;
+
+  string name = this->DBName();
+  shared_ptr<Dataset<string, value_type> > dataset =
+      DatasetFactory<string, value_type>(backend);
+  EXPECT_TRUE(dataset->open(name, Dataset<string, value_type>::New));
+
+  value_type value = this->TestValue();
+
+  string key3 = "03";
+  EXPECT_TRUE(dataset->put(key3, value));
+
+  string key1 = "01";
+  EXPECT_TRUE(dataset->put(key1, value));
+
+  string key2 = "02";
+  EXPECT_TRUE(dataset->put(key2, value));
+
+  EXPECT_TRUE(dataset->commit());
+
+  string first_key;
+  dataset->first_key(&first_key);
+  string last_key;
+  dataset->last_key(&last_key);
+
+  EXPECT_TRUE(this->equals(first_key, key1));
+  EXPECT_TRUE(this->equals(last_key, key3));
+}
+
 TYPED_TEST(DatasetTest, TestKeysNoCommit) {
   UNPACK_TYPES;