#include <stdio.h>
#include <string>
#include <vector>
#include "db/dbformat.h"
#include "db/memtable.h"
#include "db/write_batch_internal.h"
#include "rocksdb/db.h"
#include "rocksdb/env.h"
#include "rocksdb/iterator.h"
#include "rocksdb/table.h"
#include "rocksdb/slice_transform.h"
#include "table/block.h"
#include "table/block_builder.h"
#include "table/format.h"
#include "table/block_hash_index.h"
#include "util/random.h"
#include "util/testharness.h"
#include "util/testutil.h"
namespace
rocksdb {
static
std::string RandomString(Random* rnd,
int
len) {
std::string r;
test::RandomString(rnd, len, &r);
return
r;
}
std::string GenerateKey(
int
primary_key,
int
secondary_key,
int
padding_size,
Random *rnd) {
char
buf[50];
char
*p = &buf[0];
snprintf(buf,
sizeof
(buf),
"%6d%4d"
, primary_key, secondary_key);
std::string k(p);
if
(padding_size) {
k += RandomString(rnd, padding_size);
}
return
k;
}
void
GenerateRandomKVs(std::vector<std::string> *keys,
std::vector<std::string> *values,
const
int
from,
const
int
len,
const
int
step = 1,
const
int
padding_size = 0,
const
int
keys_share_prefix = 1) {
Random rnd(302);
for
(
int
i = from; i < from + len; i += step) {
for
(
int
j = 0; j < keys_share_prefix; ++j) {
keys->emplace_back(GenerateKey(i, j, padding_size, &rnd));
values->emplace_back(RandomString(&rnd, 100));
}
}
}
class
BlockTest {};
TEST(BlockTest, SimpleTest) {
Random rnd(301);
Options options = Options();
std::unique_ptr<InternalKeyComparator> ic;
ic.reset(
new
test::PlainInternalKeyComparator(options.comparator));
std::vector<std::string> keys;
std::vector<std::string> values;
BlockBuilder builder(options, ic.get());
int
num_records = 100000;
GenerateRandomKVs(&keys, &values, 0, num_records);
for
(
int
i = 0; i < num_records; i++) {
builder.Add(keys[i], values[i]);
}
Slice rawblock = builder.Finish();
BlockContents contents;
contents.data = rawblock;
contents.cachable =
false
;
contents.heap_allocated =
false
;
Block reader(contents);
int
count = 0;
Iterator* iter = reader.NewIterator(options.comparator);
for
(iter->SeekToFirst();iter->Valid(); count++, iter->Next()) {
Slice k = iter->key();
Slice v = iter->value();
ASSERT_EQ(k.ToString().compare(keys[count]), 0);
ASSERT_EQ(v.ToString().compare(values[count]), 0);
}
delete
iter;
iter = reader.NewIterator(options.comparator);
for
(
int
i = 0; i < num_records; i++) {
int
index = rnd.Uniform(num_records);
Slice k(keys[index]);
iter->Seek(k);
ASSERT_TRUE(iter->Valid());
Slice v = iter->value();
ASSERT_EQ(v.ToString().compare(values[index]), 0);
}
delete
iter;
}
BlockContents GetBlockContents(std::unique_ptr<BlockBuilder> *builder,
const
std::vector<std::string> &keys,
const
std::vector<std::string> &values,
const
int
prefix_group_size = 1) {
builder->reset(
new
BlockBuilder(1
, BytewiseComparator()));
for
(
size_t
i = 0; i < keys.size(); ++i) {
(*builder)->Add(keys[i], values[i]);
}
Slice rawblock = (*builder)->Finish();
BlockContents contents;
contents.data = rawblock;
contents.cachable =
false
;
contents.heap_allocated =
false
;
return
contents;
}
void
CheckBlockContents(BlockContents contents,
const
int
max_key,
const
std::vector<std::string> &keys,
const
std::vector<std::string> &values) {
const
size_t
prefix_size = 6;
Block reader1(contents);
Block reader2(contents);
std::unique_ptr<
const
SliceTransform> prefix_extractor(
NewFixedPrefixTransform(prefix_size));
{
auto
iter1 = reader1.NewIterator(
nullptr
);
auto
iter2 = reader1.NewIterator(
nullptr
);
reader1.SetBlockHashIndex(CreateBlockHashIndexOnTheFly(
iter1, iter2, keys.size(), BytewiseComparator(),
prefix_extractor.get()));
delete
iter1;
delete
iter2;
}
std::unique_ptr<Iterator> hash_iter(
reader1.NewIterator(BytewiseComparator()));
std::unique_ptr<Iterator> regular_iter(
reader2.NewIterator(BytewiseComparator()));
for
(
size_t
i = 0; i < keys.size(); i++) {
hash_iter->Seek(keys[i]);
ASSERT_OK(hash_iter->status());
ASSERT_TRUE(hash_iter->Valid());
Slice v = hash_iter->value();
ASSERT_EQ(v.ToString().compare(values[i]), 0);
}
for
(
int
i = 1; i < max_key - 1; i += 2) {
auto
key = GenerateKey(i, 0, 0,
nullptr
);
hash_iter->Seek(key);
ASSERT_TRUE(!hash_iter->Valid());
regular_iter->Seek(key);
ASSERT_TRUE(regular_iter->Valid());
}
}
TEST(BlockTest, SimpleIndexHash) {
const
int
kMaxKey = 100000;
std::vector<std::string> keys;
std::vector<std::string> values;
GenerateRandomKVs(&keys, &values, 0
,
kMaxKey
, 2
,
8
);
std::unique_ptr<BlockBuilder> builder;
auto
contents = GetBlockContents(&builder, keys, values);
CheckBlockContents(contents, kMaxKey, keys, values);
}
TEST(BlockTest, IndexHashWithSharedPrefix) {
const
int
kMaxKey = 100000;
const
int
kPrefixGroup = 5;
std::vector<std::string> keys;
std::vector<std::string> values;
GenerateRandomKVs(&keys, &values, 0,
kMaxKey,
2,
10,
kPrefixGroup);
std::unique_ptr<BlockBuilder> builder;
auto
contents = GetBlockContents(&builder, keys, values, kPrefixGroup);
CheckBlockContents(contents, kMaxKey, keys, values);
}
}
int
main(
int
argc,
char
** argv) {
return
rocksdb::test::RunAllTests();
}