ExampleProgram.cc

#ifndef ROCKSDB_LITE

#include "rocksdb/db.h"
#include "rocksdb/options.h"
#include "rocksdb/slice.h"
#include <iterator>
#include <chrono>
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <set>
#include <sstream>
#include <string>
#include <vector>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <pthread.h>
#include <cstdlib>
#include "rocksdb/db.h"
#include "rocksdb/merge_operator.h"
#include "rocksdb/options.h"
#include "rocksdb/slice.h"
#include "../utilities/merge_operators.h"
#include "rocksdb/utilities/write_batch_with_index.h"

#define CHECK_STATUS(where, s)                         \
  do {                                                 \
    if (!s.ok()) {                                     \
      std::cout << where << s.ToString() << std::endl; \
      abort();                                         \
    }                                                  \
  } while (0)

using namespace rocksdb;

std::string kDBPath = "/tmp/rocksdb_bulkcopy_example";
const char *cfName1 = "zlabs_cf_1";
const char *cfName2 = "zlabs_cf_2";

DB *db;
WriteOptions writeOptions;
ReadOptions readoptions;

class CSVRow
{
    public:
        std::string const& operator[](std::size_t index) const
        {
            return m_data[index];
        }
        std::size_t size() const
        {
            return m_data.size();
        }
        void readNextRow(std::istream& str)
        {
            std::string         line;
            std::getline(str, line);

            std::stringstream   lineStream(line);
            std::string         cell;

            m_data.clear();
            while(std::getline(lineStream, cell, ','))
            {
                m_data.push_back(cell);
            }
            // This checks for a trailing comma with no data after it.
            if (!lineStream && cell.empty())
            {
                // If there was a trailing comma then add an empty element.
                m_data.push_back("");
            }
        }
    // private:
        std::vector<std::string>    m_data;
};

std::istream& operator>>(std::istream& str, CSVRow& data)
{
    data.readNextRow(str);
    return str;
}

class CSVIterator
{
    public:
        typedef std::input_iterator_tag     iterator_category;
        typedef CSVRow                      value_type;
        typedef std::size_t                 difference_type;
        typedef CSVRow*                     pointer;
        typedef CSVRow&                     reference;

        CSVIterator(std::istream& str)  :m_str(str.good()?&str:NULL) { ++(*this); }
        CSVIterator()                   :m_str(NULL) {}

        // Pre Increment
        CSVIterator& operator++()               {if (m_str) { if (!((*m_str) >> m_row)){m_str = NULL;}}return *this;}
        // Post increment
        CSVIterator operator++(int)             {CSVIterator    tmp(*this);++(*this);return tmp;}
        CSVRow const& operator*()   const       {return m_row;}
        CSVRow const* operator->()  const       {return &m_row;}

        bool operator==(CSVIterator const& rhs) {return ((this == &rhs) || ((this->m_str == NULL) && (rhs.m_str == NULL)));}
        bool operator!=(CSVIterator const& rhs) {return !((*this) == rhs);}
    private:
        std::istream*       m_str;
        CSVRow              m_row;
};


int main() {

  // open DB
  Options options;
  options.create_if_missing = true;
  options.merge_operator = rocksdb::MergeOperators::CreateStringAppendOperator();

  options.max_successive_merges = 10;

  bool dbAlreadyExists = false;

  Status s;
  std::ifstream file("PathToFolder/10MillionData.txt");

  if (stat(kDBPath.c_str(), new struct stat()) == -1) {
    s = DB::Open(options, kDBPath, &db);

    assert(s.ok());
    CHECK_STATUS("FirstDB::OPEN", s);

    ColumnFamilyHandle *cf;
    s = db->CreateColumnFamily(ColumnFamilyOptions(options), cfName1, &cf);
    assert(s.ok());

    s = db->CreateColumnFamily(ColumnFamilyOptions(options), cfName2, &cf);
    assert(s.ok());

    delete cf;
    delete db;
  }
  else
    dbAlreadyExists = true;

  std::vector<ColumnFamilyDescriptor> column_families;

  column_families.push_back(ColumnFamilyDescriptor(
      kDefaultColumnFamilyName, ColumnFamilyOptions(options)));
  // open the new one, too
  column_families.push_back(
      ColumnFamilyDescriptor(cfName1, ColumnFamilyOptions(options)));
  column_families.push_back(
      ColumnFamilyDescriptor(cfName2, ColumnFamilyOptions(options)));
  std::vector<ColumnFamilyHandle *> handles;
  s = DB::Open(options, kDBPath, column_families, &handles, &db);
  CHECK_STATUS("Second DB::Open", s);

  using std::cout;
  using std::endl;

  if (!dbAlreadyExists) {
    WriteBatchWithIndex *wbwi = new WriteBatchWithIndex(BytewiseComparator(), 512, false);
    Status s;

    cout << "Inserting data into DB\n";
    /////////////////////////////////////////////////////
    for(CSVIterator loop(file); loop != CSVIterator(); ++loop) {
      wbwi->Merge(handles[1], Slice((*loop)[0]), Slice((*loop)[1]));
      wbwi->Merge(handles[2], Slice((*loop)[1]), Slice((*loop)[0]));
    }
    /////////////////////////////////////////////////////

    s = db->Write(writeOptions, wbwi->GetWriteBatch());

    wbwi->Clear();

    CHECK_STATUS("insertDataToDB Write", s);
    cout << "Insert Complete\n";

    usleep(60000000);// will sleep for 60s, Gives time to check if memory is freed
  }

  for (auto handle : handles) {
    delete handle;
  }
  delete db;
}

#endif  // ROCKSDB_LITE