aozturk
11/23/2013 - 8:52 PM

RocksComparator.cpp

class TwoPartComparator : public rocksdb::Comparator {
public:
    // Three-way comparison function:
    //   if a < b: negative result
    //   if a > b: positive result
    //   else: zero result
    int Compare(const rocksdb::Slice& a, const rocksdb::Slice& b) const {
        int a1, a2, b1, b2;
        ParseKey(a, &a1, &a2);
        ParseKey(b, &b1, &b2);
        if (a1 < b1) return -1;
        if (a1 > b1) return +1;
        if (a2 < b2) return -1;
        if (a2 > b2) return +1;
        return 0;
    }

    void ParseKey(const rocksdb::Slice& a, int* a1, int* a2) const {
    	std::string parts = a.ToString();
    	int index = parts.find_first_of(":");
    	*a1 = atoi(parts.substr(0, index).c_str());
    	*a2 = atoi(parts.substr(index+1, parts.size()).c_str());
    }

    // Ignore the following methods for now:
    const char* Name() const { return "TwoPartComparator"; }
    void FindShortestSeparator(std::string*, const rocksdb::Slice&) const { }
    void FindShortSuccessor(std::string*) const { }
  };

int main() {
    rocksdb::DB* db;
    rocksdb::Options options;
    options.create_if_missing = true;
    TwoPartComparator cmp;
    options.comparator = &cmp;

    // open a database with custom comparator
    rocksdb::Status status = rocksdb::DB::Open(options, "/tmp/testdb", &db);
    assert(status.ok());

    // populate the database
    rocksdb::Slice key1 = "1:3";
    rocksdb::Slice key2 = "2:3";
    rocksdb::Slice key3 = "2:1";
    std::string val1 = "one";
    std::string val2 = "two";
    std::string val3 = "three";
    db->Put(rocksdb::WriteOptions(), key1, val1);
    db->Put(rocksdb::WriteOptions(), key2, val2);
    db->Put(rocksdb::WriteOptions(), key3, val3);

    //iterate the database
    rocksdb::Iterator* it = db->NewIterator(rocksdb::ReadOptions());
    for (it->SeekToFirst(); it->Valid(); it->Next()) {
        cout << it->key().ToString() << ": " << it->value().ToString() << endl;
    }
    delete it;

    // close the database
    delete db;

    return 0;
}