From 23af87fc49d7ffaa70584458cacb1528a5fb1f77 Mon Sep 17 00:00:00 2001 From: Sergio Date: Wed, 15 Oct 2014 09:45:09 -0700 Subject: [PATCH] Added first_key and last_key to dataset --- include/caffe/dataset.hpp | 2 + include/caffe/leveldb_dataset.hpp | 2 + include/caffe/lmdb_dataset.hpp | 2 + src/caffe/leveldb_dataset.cpp | 24 ++++++++ src/caffe/lmdb_dataset.cpp | 64 +++++++++++++++++++++ src/caffe/test/test_dataset.cpp | 114 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 208 insertions(+) diff --git a/include/caffe/dataset.hpp b/include/caffe/dataset.hpp index efe3ffd..90deb31 100644 --- 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 --- a/include/caffe/leveldb_dataset.hpp +++ b/include/caffe/leveldb_dataset.hpp @@ -29,6 +29,8 @@ class LeveldbDataset : public Dataset { 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 --- a/include/caffe/lmdb_dataset.hpp +++ b/include/caffe/lmdb_dataset.hpp @@ -33,6 +33,8 @@ class LmdbDataset : public Dataset { 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 --- a/src/caffe/leveldb_dataset.cpp +++ b/src/caffe/leveldb_dataset.cpp @@ -107,6 +107,30 @@ bool LeveldbDataset::get(const K& key, V* value) { } template +bool LeveldbDataset::first_key(K* key) { + DLOG(INFO) << "LevelDB: First key"; + + CHECK_NOTNULL(db_.get()); + shared_ptr 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 +bool LeveldbDataset::last_key(K* key) { + DLOG(INFO) << "LevelDB: Last key"; + + CHECK_NOTNULL(db_.get()); + shared_ptr 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 bool LeveldbDataset::commit() { DLOG(INFO) << "LevelDB: Commit"; diff --git a/src/caffe/lmdb_dataset.cpp b/src/caffe/lmdb_dataset.cpp index ca96843..23852f0 100644 --- a/src/caffe/lmdb_dataset.cpp +++ b/src/caffe/lmdb_dataset.cpp @@ -155,6 +155,70 @@ bool LmdbDataset::get(const K& key, V* value) { } template +bool LmdbDataset::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(mdbkey.mv_data), + mdbkey.mv_size, key)) { + LOG(ERROR) << "failed to deserialize key"; + return false; + } + + return true; +} + +template +bool LmdbDataset::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(mdbkey.mv_data), + mdbkey.mv_size, key)) { + LOG(ERROR) << "failed to deserialize key"; + return false; + } + + return true; +} + +template bool LmdbDataset::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 --- 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 = + DatasetFactory(backend); + EXPECT_TRUE(dataset->open(name, Dataset::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 = + DatasetFactory(backend); + EXPECT_TRUE(dataset->open(name, Dataset::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 = + DatasetFactory(backend); + EXPECT_TRUE(dataset->open(name, Dataset::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 = + DatasetFactory(backend); + EXPECT_TRUE(dataset->open(name, Dataset::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; -- 2.7.4