singa-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From wang...@apache.org
Subject [1/2] incubator-singa git commit: SINGA-214 Add LMDBReader and LMDBWriter for LMDB
Date Fri, 01 Jul 2016 08:35:27 GMT
Repository: incubator-singa
Updated Branches:
  refs/heads/dev cde7dcf6b -> f0bc22889


SINGA-214 Add LMDBReader and LMDBWriter for LMDB

Add LMDBReader and LMDBWriter for LMDB. The classes are modified from db_lmdb.cpp and db_lmdb.hpp
in Caffe.
Need to swith "USE_LMDB" ON when using these two classes.

(Jul 1) Add SeekToFirst to all IO classes;
LMDB does not support checking empty and duplicate keys, Check empty duplicate keys in LMDBWriter.


Project: http://git-wip-us.apache.org/repos/asf/incubator-singa/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-singa/commit/bcda737c
Tree: http://git-wip-us.apache.org/repos/asf/incubator-singa/tree/bcda737c
Diff: http://git-wip-us.apache.org/repos/asf/incubator-singa/diff/bcda737c

Branch: refs/heads/dev
Commit: bcda737c99ce0ca1a63c8f0488fae138e54b5d0e
Parents: f026b60
Author: XiangruiCAI <caixr91@gmail.com>
Authored: Thu Jun 30 21:25:04 2016 +0800
Committer: XiangruiCAI <caixr91@gmail.com>
Committed: Fri Jul 1 14:55:12 2016 +0800

----------------------------------------------------------------------
 include/singa/io/reader.h      |  63 ++++++++++++++++-
 include/singa/io/writer.h      |  46 +++++++++++-
 src/io/binfile_reader.cc       |  16 ++++-
 src/io/lmdb_reader.cc          | 118 +++++++++++++++++++++++++++++++
 src/io/lmdb_writer.cc          | 133 +++++++++++++++++++++++++++++++++++
 src/io/textfile_reader.cc      |   6 ++
 test/singa/test_binfile_rw.cc  |  38 ++++++++++
 test/singa/test_lmdb_rw.cc     | 136 ++++++++++++++++++++++++++++++++++++
 test/singa/test_textfile_rw.cc |  39 +++++++++++
 9 files changed, 589 insertions(+), 6 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/include/singa/io/reader.h
----------------------------------------------------------------------
diff --git a/include/singa/io/reader.h b/include/singa/io/reader.h
index bd1b3fe..66d7e37 100644
--- a/include/singa/io/reader.h
+++ b/include/singa/io/reader.h
@@ -22,6 +22,13 @@
 #include <cstring>
 #include <fstream>
 #include <string>
+#include "singa/singa_config.h"
+
+#ifdef USE_LMDB
+#include <lmdb.h>
+#include <sys/stat.h>
+#include <vector>
+#endif  // USE_LMDB
 
 namespace singa {
 namespace io {
@@ -46,12 +53,17 @@ class Reader {
   virtual void Close() = 0;
 
   /// Read a tuple.
-  /// return true if read successfully, otherwise false.
+  /// return true if read successfully;
+  /// return flase if coming to the end of the file;
+  /// LOG(FATAL) if error happens.
   virtual bool Read(std::string* key, std::string* value) = 0;
 
   /// Iterate through all tuples to get the num of all tuples.
   /// return num of tuples
   virtual int Count() = 0;
+
+  /// Seek to the first tuple when the cursor arrives to the end of the file
+  virtual void SeekToFirst() = 0;
 };
 
 /// Binfilereader reads tuples from binary file with key-value pairs.
@@ -68,6 +80,8 @@ class BinFileReader : public Reader {
   bool Read(std::string* key, std::string* value) override;
   /// \copydoc Count()
   int Count() override;
+  /// \copydoc SeekToFirst()
+  void SeekToFirst() override;
   /// return path to binary file
   inline std::string path() { return path_; }
 
@@ -110,6 +124,8 @@ class TextFileReader : public Reader {
   bool Read(std::string* key, std::string* value) override;
   /// \copydoc Count()
   int Count() override;
+  /// \copydoc SeekToFirst()
+  void SeekToFirst() override;
   /// return path to text file
   inline std::string path() { return path_; }
 
@@ -121,6 +137,51 @@ class TextFileReader : public Reader {
   /// current line number
   int lineNo_ = 0;
 };
+
+#ifdef USE_LMDB
+/// LMDBReader reads tuples from LMDB.
+class LMDBReader : public Reader {
+ public:
+  ~LMDBReader() { Close(); }
+  /// \copydoc Open(const std::string& path)
+  bool Open(const std::string& path) override;
+  /// \copydoc Close()
+  void Close() override;
+  /// \copydoc Read(std::string* key, std::string* value)
+  bool Read(std::string* key, std::string* value) override;
+  /// \copydoc Count()
+  int Count() override;
+  /// \copydoc SeekToFirst()
+  void SeekToFirst() override;
+  /// Return path to text file
+  inline std::string path() { return path_; }
+  /// Return valid, to indicate SeekToFirst();
+  inline bool valid() { return valid_; }
+
+ protected:
+  /// Seek to a certain position: MDB_FIRST, MDB_NEXT
+  void Seek(MDB_cursor_op op);
+  inline void MDB_CHECK(int mdb_status);
+
+ private:
+  /// file to be read
+  std::string path_ = "";
+  /// lmdb env variable
+  MDB_env* mdb_env_ = nullptr;
+  /// lmdb db instance
+  MDB_dbi mdb_dbi_;
+  /// lmdb transaction
+  MDB_txn* mdb_txn_ = nullptr;
+  /// lmdb cursor
+  MDB_cursor* mdb_cursor_ = nullptr;
+  /// lmdb key-value pair
+  MDB_val mdb_key_, mdb_value_;
+  /// whether the pair is found
+  bool valid_;
+  /// whether the cursor is at the first place
+  bool first_;
+};
+#endif  // USE_LMDB
 }  // namespace io
 }  // namespace singa
 

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/include/singa/io/writer.h
----------------------------------------------------------------------
diff --git a/include/singa/io/writer.h b/include/singa/io/writer.h
index f20a22b..bd4043a 100644
--- a/include/singa/io/writer.h
+++ b/include/singa/io/writer.h
@@ -19,9 +19,16 @@
 #ifndef SINGA_IO_WRITER_H_
 #define SINGA_IO_WRITER_H_
 
-#include <string>
 #include <cstring>
 #include <fstream>
+#include <string>
+#include "singa/singa_config.h"
+
+#ifdef USE_LMDB
+#include <lmdb.h>
+#include <sys/stat.h>
+#include <vector>
+#endif  // USE_LMDB
 
 namespace singa {
 namespace io {
@@ -73,7 +80,7 @@ class BinFileWriter : public Writer {
   /// \copydoc Open(const std::string &path, Mode mode)
   bool Open(const std::string &path, Mode mode) override;
   /// \copydoc Open(const std::string& path), user defines capacity
-  bool Open(const std::string& path, Mode mode, int capacity);
+  bool Open(const std::string &path, Mode mode, int capacity);
   /// \copydoc Close()
   void Close() override;
   /// \copydoc Write(const std::string& key, const std::string& value) override;
@@ -100,7 +107,7 @@ class BinFileWriter : public Writer {
   /// bytes in buf_
   int bufsize_ = 0;
   /// magic word
-  const char kMagicWord[2]= {'s', 'g'};
+  const char kMagicWord[2] = {'s', 'g'};
 };
 
 /// TextFileWriter write training/validation/test tuples in CSV file.
@@ -125,6 +132,39 @@ class TextFileWriter : public Writer {
   /// ofstream
   std::ofstream fdat_;
 };
+
+#ifdef USE_LMDB
+/// LMDBWriter write training/validation/test tuples into LMDB.
+class LMDBWriter : public Writer {
+ public:
+  ~LMDBWriter() { Close(); }
+  /// \copydoc Open(const std::string &path, Mode mode)
+  bool Open(const std::string &path, Mode mode) override;
+  /// \copydoc Close()
+  void Close() override;
+  /// \copydoc Write(const std::string& key, const std::string& value) override;
+  bool Write(const std::string &key, const std::string &value) override;
+  /// \copydoc Flush()
+  void Flush() override;
+  /// return path to text file
+  inline std::string path() { return path_; }
+
+ protected:
+  void DoubleMapSize();
+  inline void MDB_CHECK(int mdb_status);
+
+ private:
+  /// file to be written
+  std::string path_ = "";
+  /// kCreate or kAppend
+  Mode mode_;
+  /// lmdb env variable
+  MDB_env *mdb_env_ = nullptr;
+  /// buffer for key-value pairs
+  std::vector<string> keys, values;
+};
+#endif  // USE_LMDB
+
 }  // namespace io
 }  // namespace singa
 

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/src/io/binfile_reader.cc
----------------------------------------------------------------------
diff --git a/src/io/binfile_reader.cc b/src/io/binfile_reader.cc
index d54eeb5..77e34d8 100644
--- a/src/io/binfile_reader.cc
+++ b/src/io/binfile_reader.cc
@@ -49,11 +49,15 @@ bool BinFileReader::Read(std::string* key, std::string* value) {
   offset_ += smagic;
 
   if (magic[0] == kMagicWord[0] && magic[1] == kMagicWord[1]) {
-    if (magic[2] != 0 && magic[2] != 1) return false;
+    if (magic[2] != 0 && magic[2] != 1)
+      LOG(FATAL) << "File format error: magic word does not match!";
     if (magic[2] == 1)
       if (!ReadField(key)) return false;
     if (!ReadField(value)) return false;
   }
+  else {
+    LOG(FATAL) << "File format error: magic word does not match!";
+  }
   return true;
 }
 
@@ -82,6 +86,14 @@ int BinFileReader::Count() {
   return count;
 }
 
+void BinFileReader::SeekToFirst() {
+  bufsize_ = 0;
+  offset_ = 0;
+  fdat_.clear();
+  fdat_.seekg(0);
+  CHECK(fdat_.is_open()) << "Cannot create file " << path_;
+}
+
 bool BinFileReader::OpenFile() {
   buf_ = new char[capacity_];
   fdat_.open(path_, std::ios::in | std::ios::binary);
@@ -112,7 +124,7 @@ bool BinFileReader::PrepareNextField(int size) {
     } else {
       fdat_.read(buf_ + bufsize_, capacity_ - bufsize_);
       bufsize_ += fdat_.gcount();
-      if (size > bufsize_) return false;
+      CHECK_LE(size, bufsize_) << "Field size is too large: " << size;
     }
   }
   return true;

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/src/io/lmdb_reader.cc
----------------------------------------------------------------------
diff --git a/src/io/lmdb_reader.cc b/src/io/lmdb_reader.cc
new file mode 100644
index 0000000..7f78080
--- /dev/null
+++ b/src/io/lmdb_reader.cc
@@ -0,0 +1,118 @@
+/**
+ * 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.
+ */
+
+#include "singa/io/reader.h"
+#include "singa/utils/logging.h"
+#ifdef USE_LMDB
+
+namespace singa {
+namespace io {
+bool LMDBReader::Open(const std::string& path) {
+  path_ = path;
+  MDB_CHECK(mdb_env_create(&mdb_env_));
+  int flags = MDB_RDONLY | MDB_NOTLS;
+  int rc = mdb_env_open(mdb_env_, path_.c_str(), flags, 0664);
+#ifndef ALLOW_LMDB_NOLOCK
+  MDB_CHECK(rc);
+#else
+  if (rc == EACCES) {
+    LOG(WARNING) << "Permission denied. Trying with MDB_NOLOCK ...";
+    // Close and re-open environment handle
+    mdb_env_close(mdb_env_);
+    MDB_CHECK(mdb_env_create(&mdb_env_));
+    // Try again with MDB_NOLOCK
+    flags |= MDB_NOLOCK;
+    MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
+  } else {
+    MDB_CHECK(rc);
+  }
+#endif
+  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
+  MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
+  MDB_CHECK(mdb_cursor_open(mdb_txn_, mdb_dbi_, &mdb_cursor_));
+  SeekToFirst();
+  return true;
+}
+
+void LMDBReader::Close() {
+  if (mdb_env_ != nullptr) {
+    mdb_cursor_close(mdb_cursor_);
+    mdb_txn_abort(mdb_txn_);
+    mdb_dbi_close(mdb_env_, mdb_dbi_);
+    mdb_env_close(mdb_env_);
+    mdb_env_ = nullptr;
+    mdb_txn_ = nullptr;
+    mdb_cursor_ = nullptr;
+  }
+}
+
+bool LMDBReader::Read(std::string* key, std::string* value) {
+  if (first_ != true)
+    Seek(MDB_NEXT);
+  if (valid_ == false) return false;
+  *key = string(static_cast<const char*>(mdb_key_.mv_data), mdb_key_.mv_size);
+  *value =
+      string(static_cast<const char*>(mdb_value_.mv_data), mdb_value_.mv_size);
+  first_ = false;
+  return true;
+}
+
+int LMDBReader::Count() {
+  MDB_env* env;
+  MDB_dbi dbi;
+  MDB_txn* txn;
+  MDB_cursor* cursor;
+  int flags = MDB_RDONLY | MDB_NOTLS | MDB_NOLOCK;
+  MDB_CHECK(mdb_env_create(&env));
+  MDB_CHECK(mdb_env_open(env, path_.c_str(), flags, 0664));
+  MDB_CHECK(mdb_txn_begin(env, NULL, MDB_RDONLY, &txn));
+  MDB_CHECK(mdb_dbi_open(txn, NULL, 0, &dbi));
+  MDB_CHECK(mdb_cursor_open(txn, dbi, &cursor));
+  int status = MDB_SUCCESS;
+  int count = 0;
+  MDB_val key, value;
+  while (true) {
+    status = mdb_cursor_get(cursor, &key, &value, MDB_NEXT);
+    if (status == MDB_NOTFOUND) break;
+    count++;
+  }
+  mdb_cursor_close(cursor);
+  mdb_txn_abort(txn);
+  mdb_dbi_close(env, dbi);
+  mdb_env_close(env);
+  return count;
+}
+
+void LMDBReader::SeekToFirst() { Seek(MDB_FIRST); first_ = true; }
+
+void LMDBReader::Seek(MDB_cursor_op op) {
+  int mdb_status = mdb_cursor_get(mdb_cursor_, &mdb_key_, &mdb_value_, op);
+  if (mdb_status == MDB_NOTFOUND) {
+    valid_ = false;
+  } else {
+    MDB_CHECK(mdb_status);
+    valid_ = true;
+  }
+}
+
+inline void LMDBReader::MDB_CHECK(int mdb_status) {
+  CHECK_EQ(mdb_status, MDB_SUCCESS) << mdb_strerror(mdb_status);
+}
+}  // namespace io
+}  // namespace singa
+#endif  // USE_LMDB

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/src/io/lmdb_writer.cc
----------------------------------------------------------------------
diff --git a/src/io/lmdb_writer.cc b/src/io/lmdb_writer.cc
new file mode 100644
index 0000000..e89894b
--- /dev/null
+++ b/src/io/lmdb_writer.cc
@@ -0,0 +1,133 @@
+/**
+ * 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.
+ */
+
+#include "singa/io/writer.h"
+#include "singa/utils/logging.h"
+#ifdef USE_LMDB
+
+namespace singa {
+namespace io {
+bool LMDBWriter::Open(const std::string& path, Mode mode) {
+  path_ = path;
+  mode_ = mode;
+  MDB_CHECK(mdb_env_create(&mdb_env_));
+  if (mode_ != kCreate && mode_ != kAppend) {
+    LOG(FATAL) << "unknown mode to open LMDB" << mode_;
+    return false;
+  }
+  if (mode_ == kCreate)
+    // It will fail if there is a dir at "path"
+    CHECK_EQ(mkdir(path.c_str(), 0744), 0) << "mkdir " << path << " failed";
+  int flags = 0;
+  int rc = mdb_env_open(mdb_env_, path.c_str(), flags, 0664);
+#ifndef ALLOW_LMDB_NOLOCK
+  MDB_CHECK(rc);
+#else
+  if (rc == EACCES) {
+    LOG(WARNING) << "Permission denied. Trying with MDB_NOLOCK ...";
+    // Close and re-open environment handle
+    mdb_env_close(mdb_env_);
+    MDB_CHECK(mdb_env_create(&mdb_env_));
+    // Try again with MDB_NOLOCK
+    flags |= MDB_NOLOCK;
+    MDB_CHECK(mdb_env_open(mdb_env_, path.c_str(), flags, 0664));
+  } else
+    MDB_CHECK(rc);
+#endif
+  return true;
+}
+
+void LMDBWriter::Close() {
+  Flush();
+  if (mdb_env_ != nullptr) {
+    mdb_env_close(mdb_env_);
+    mdb_env_ = nullptr;
+  }
+}
+
+bool LMDBWriter::Write(const std::string& key, const std::string& value) {
+  CHECK_NE(key, "") << "Key is an empty string!";
+  keys.push_back(key);
+  values.push_back(value);
+  return true;
+}
+
+// Flush is to "commit to DB"
+void LMDBWriter::Flush() {
+  if (keys.size() == 0) return;
+  MDB_dbi mdb_dbi;
+  MDB_val mdb_key, mdb_data;
+  MDB_txn* mdb_txn;
+
+  // Initialize MDB variables
+  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn));
+  MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi));
+
+  for (size_t i = 0; i < keys.size(); i++) {
+    mdb_key.mv_size = keys[i].size();
+    mdb_key.mv_data = const_cast<char*>(keys[i].data());
+    mdb_data.mv_size = values[i].size();
+    mdb_data.mv_data = const_cast<char*>(values[i].data());
+
+    // Add data to the transaction
+    int put_rc = mdb_put(mdb_txn, mdb_dbi, &mdb_key, &mdb_data, 0);
+    CHECK_NE(put_rc, MDB_KEYEXIST) << "Key already exist: " << keys[i];
+    if (put_rc == MDB_MAP_FULL) {
+      // Out of memory - double the map size and retry
+      mdb_txn_abort(mdb_txn);
+      mdb_dbi_close(mdb_env_, mdb_dbi);
+      DoubleMapSize();
+      Flush();
+      return;
+    }
+    // May have failed for some other reason
+    MDB_CHECK(put_rc);
+  }
+
+  // Commit the transaction
+  int commit_rc = mdb_txn_commit(mdb_txn);
+  if (commit_rc == MDB_MAP_FULL) {
+    // Out of memory - double the map size and retry
+    mdb_dbi_close(mdb_env_, mdb_dbi);
+    DoubleMapSize();
+    Flush();
+    return;
+  }
+  // May have failed for some other reason
+  MDB_CHECK(commit_rc);
+
+  // Cleanup after successful commit
+  mdb_dbi_close(mdb_env_, mdb_dbi);
+  keys.clear();
+  values.clear();
+}
+
+void LMDBWriter::DoubleMapSize() {
+  struct MDB_envinfo current_info;
+  MDB_CHECK(mdb_env_info(mdb_env_, &current_info));
+  size_t new_size = current_info.me_mapsize * 2;
+  LOG(INFO) << "Doubling LMDB map size to " << (new_size >> 20) <<
"MB ...";
+  MDB_CHECK(mdb_env_set_mapsize(mdb_env_, new_size));
+}
+
+inline void LMDBWriter::MDB_CHECK(int mdb_status) {
+  CHECK_EQ(mdb_status, MDB_SUCCESS) << mdb_strerror(mdb_status);
+}
+}  // namespace io
+}  // namespace singa
+#endif  // USE_LMDB

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/src/io/textfile_reader.cc
----------------------------------------------------------------------
diff --git a/src/io/textfile_reader.cc b/src/io/textfile_reader.cc
index 7612241..714aa51 100644
--- a/src/io/textfile_reader.cc
+++ b/src/io/textfile_reader.cc
@@ -59,5 +59,11 @@ int TextFileReader::Count() {
   return count;
 }
 
+void TextFileReader::SeekToFirst() {
+  CHECK(fdat_ != nullptr);
+  lineNo_ = 0;
+  fdat_.clear();
+  fdat_.seekg(0);
+}
 }  // namespace io
 }  // namespace singa

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/test/singa/test_binfile_rw.cc
----------------------------------------------------------------------
diff --git a/test/singa/test_binfile_rw.cc b/test/singa/test_binfile_rw.cc
index ddee4f1..53c29fa 100644
--- a/test/singa/test_binfile_rw.cc
+++ b/test/singa/test_binfile_rw.cc
@@ -91,5 +91,43 @@ TEST(BinFileReader, Read) {
   EXPECT_STREQ("\nThis is another test for binfile io.", value.c_str());
 
   reader.Close();
+}
+
+TEST(BinFileReader, SeekToFirst) {
+  BinFileReader reader;
+  bool ret;
+  ret = reader.Open(path_bin);
+  EXPECT_EQ(true, ret);
+
+  int cnt = reader.Count();
+  EXPECT_EQ(4, cnt);
+
+  std::string key, value;
+  reader.Read(&key, &value);
+  EXPECT_STREQ("", key.c_str());
+  EXPECT_STREQ("\nThis is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("", key.c_str());
+  EXPECT_STREQ("\nThis is a test for binfile io.", value.c_str());
+
+  reader.SeekToFirst();
+  reader.Read(&key, &value);
+  EXPECT_STREQ("", key.c_str());
+  EXPECT_STREQ("\nThis is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("", key.c_str());
+  EXPECT_STREQ("\nThis is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("\nThis is another test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("2", key.c_str());
+  EXPECT_STREQ("\nThis is another test for binfile io.", value.c_str());
+
+  reader.Close();
   remove(path_bin);
 }

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/test/singa/test_lmdb_rw.cc
----------------------------------------------------------------------
diff --git a/test/singa/test_lmdb_rw.cc b/test/singa/test_lmdb_rw.cc
new file mode 100644
index 0000000..0d4025a
--- /dev/null
+++ b/test/singa/test_lmdb_rw.cc
@@ -0,0 +1,136 @@
+/************************************************************
+*
+* 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.
+*
+*************************************************************/
+
+#include "../include/singa/io/reader.h"
+#include "../include/singa/io/writer.h"
+#include "gtest/gtest.h"
+#ifdef USE_LMDB
+
+const char* path_lmdb = "./test_lmdb";
+using singa::io::LMDBReader;
+using singa::io::LMDBWriter;
+TEST(LMDBWriter, Create) {
+  LMDBWriter writer;
+  bool ret;
+  ret = writer.Open(path_lmdb, singa::io::kCreate);
+  EXPECT_EQ(true, ret);
+
+  std::string key = "1";
+  std::string value = "This is the first test for lmdb io.";
+  ret = writer.Write(key, value);
+  EXPECT_EQ(true, ret);
+
+  key = "2";
+  value = "This is the second test for lmdb io.";
+  ret = writer.Write(key, value);
+  EXPECT_EQ(true, ret);
+
+  writer.Flush();
+  writer.Close();
+}
+
+TEST(LMDBWriter, Append) {
+  LMDBWriter writer;
+  bool ret;
+  ret = writer.Open(path_lmdb, singa::io::kAppend);
+  EXPECT_EQ(true, ret);
+
+  std::string key = "3";
+  std::string value = "This is the third test for lmdb io.";
+  ret = writer.Write(key, value);
+  EXPECT_EQ(true, ret);
+
+  key = "4";
+  value = "This is the fourth test for lmdb io.";
+  ret = writer.Write(key, value);
+  EXPECT_EQ(true, ret);
+
+  writer.Flush();
+  writer.Close();
+}
+
+TEST(LMDBReader, Read) {
+  LMDBReader reader;
+  bool ret;
+  ret = reader.Open(path_lmdb);
+  EXPECT_EQ(true, ret);
+
+  int cnt = reader.Count();
+  EXPECT_EQ(4, cnt);
+
+  std::string key, value;
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("This is the first test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("2", key.c_str());
+  EXPECT_STREQ("This is the second test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("3", key.c_str());
+  EXPECT_STREQ("This is the third test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("4", key.c_str());
+  EXPECT_STREQ("This is the fourth test for lmdb io.", value.c_str());
+
+  reader.Close();
+}
+
+TEST(LMDBReader, SeekToFirst) {
+  LMDBReader reader;
+  bool ret;
+  ret = reader.Open(path_lmdb);
+  EXPECT_EQ(true, ret);
+
+  int cnt = reader.Count();
+  EXPECT_EQ(4, cnt);
+
+  std::string key, value;
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("This is the first test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("2", key.c_str());
+  EXPECT_STREQ("This is the second test for lmdb io.", value.c_str());
+
+  reader.SeekToFirst();
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("This is the first test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("2", key.c_str());
+  EXPECT_STREQ("This is the second test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("3", key.c_str());
+  EXPECT_STREQ("This is the third test for lmdb io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("4", key.c_str());
+  EXPECT_STREQ("This is the fourth test for lmdb io.", value.c_str());
+
+  reader.Close();
+}
+#endif  // USE_LMDB

http://git-wip-us.apache.org/repos/asf/incubator-singa/blob/bcda737c/test/singa/test_textfile_rw.cc
----------------------------------------------------------------------
diff --git a/test/singa/test_textfile_rw.cc b/test/singa/test_textfile_rw.cc
index 7494f46..c436478 100644
--- a/test/singa/test_textfile_rw.cc
+++ b/test/singa/test_textfile_rw.cc
@@ -63,6 +63,7 @@ TEST(TextFileWriter, Append) {
   writer.Flush();
   writer.Close();
 }
+
 TEST(TextFileReader, Read) {
   TextFileReader reader;
   bool ret;
@@ -90,5 +91,43 @@ TEST(TextFileReader, Read) {
   EXPECT_STREQ("This is another test for binfile io.", value.c_str());
 
   reader.Close();
+}
+
+TEST(TextFileReader, SeekToFirst) {
+  TextFileReader reader;
+  bool ret;
+  ret = reader.Open(path_csv);
+  EXPECT_EQ(true, ret);
+
+  int cnt = reader.Count();
+  EXPECT_EQ(4, cnt);
+
+  std::string key, value;
+  reader.Read(&key, &value);
+  EXPECT_STREQ("0", key.c_str());
+  EXPECT_STREQ("This is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("This is a test for binfile io.", value.c_str());
+
+  reader.SeekToFirst();
+  reader.Read(&key, &value);
+  EXPECT_STREQ("0", key.c_str());
+  EXPECT_STREQ("This is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("1", key.c_str());
+  EXPECT_STREQ("This is a test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("2", key.c_str());
+  EXPECT_STREQ("This is another test for binfile io.", value.c_str());
+
+  reader.Read(&key, &value);
+  EXPECT_STREQ("3", key.c_str());
+  EXPECT_STREQ("This is another test for binfile io.", value.c_str());
+
+  reader.Close();
   remove(path_csv);
 }


Mime
View raw message