Skip to content

Commit

Permalink
Address feedback
Browse files Browse the repository at this point in the history
  • Loading branch information
jaykorean committed Feb 28, 2024
1 parent ead3501 commit 15ba79f
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 114 deletions.
11 changes: 7 additions & 4 deletions db/db_impl/db_impl.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3724,12 +3724,15 @@ ArenaWrappedDBIter* DBImpl::NewIteratorImpl(
std::unique_ptr<MultiCfIterator> DBImpl::NewMultiCfIterator(
const ReadOptions& _read_options,
const std::vector<ColumnFamilyHandle*>& column_families) {
assert(column_families.size() > 0);
if (column_families.size() == 0) {
return std::make_unique<EmptyMultiCfIterator>(
Status::InvalidArgument("No column families were provided"));
}
const Comparator* first_comparator = column_families[0]->GetComparator();

for (size_t i = 1; i < column_families.size(); ++i) {
if (first_comparator->GetId().compare(
column_families[i]->GetComparator()->GetId()) != 0) {
const Comparator* cf_comparator = column_families[i]->GetComparator();
if (first_comparator != cf_comparator &&
first_comparator->GetId().compare(cf_comparator->GetId()) != 0) {
return std::make_unique<EmptyMultiCfIterator>(Status::InvalidArgument(
"Different comparators are being used across CFs"));
}
Expand Down
8 changes: 4 additions & 4 deletions db/multi_cf_iterator_impl.cc
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,14 @@ namespace ROCKSDB_NAMESPACE {

void MultiCfIteratorImpl::SeekToFirst() {
Reset();
assert(cfhs_.size() == iterators_.size());
int i = 0;
for (auto& iter : iterators_) {
auto& cfh = cfhs_.at(i);
for (auto& cfh_iter_pair : cfh_iter_pairs_) {
auto& cfh = cfh_iter_pair.first;
auto& iter = cfh_iter_pair.second;
iter->SeekToFirst();
if (iter->Valid()) {
assert(iter->status().ok());
min_heap_.push(MultiCfIteratorInfo{iter, cfh, i});
min_heap_.push(MultiCfIteratorInfo{iter.get(), cfh, i});
} else {
considerStatus(iter->status());
}
Expand Down
24 changes: 9 additions & 15 deletions db/multi_cf_iterator_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,26 +17,20 @@ namespace ROCKSDB_NAMESPACE {
// order provided in the `column_families` parameter.
class MultiCfIteratorImpl : public MultiCfIterator {
public:
MultiCfIteratorImpl() {}
MultiCfIteratorImpl(const Comparator* comparator,
const std::vector<ColumnFamilyHandle*>& column_families,
const std::vector<Iterator*>& child_iterators)
: comparator_(comparator),
min_heap_(MultiCfMinHeapItemComparator(comparator_)) {
assert(column_families.size() > 0);
assert(column_families.size() == child_iterators.size());

cfhs_.reserve(column_families.size());
iterators_.reserve(column_families.size());
assert(column_families.size() > 0 &&
column_families.size() == child_iterators.size());
cfh_iter_pairs_.reserve(column_families.size());
for (size_t i = 0; i < column_families.size(); ++i) {
cfhs_.push_back(column_families[i]);
iterators_.push_back(child_iterators[i]);
cfh_iter_pairs_.emplace_back(
column_families[i], std::unique_ptr<Iterator>(child_iterators[i]));
}
}
~MultiCfIteratorImpl() override {
for (auto iter : iterators_) {
delete iter;
}
status_.PermitUncheckedError();
}

Expand All @@ -45,8 +39,8 @@ class MultiCfIteratorImpl : public MultiCfIterator {
MultiCfIteratorImpl& operator=(const MultiCfIteratorImpl&) = delete;

private:
std::vector<ColumnFamilyHandle*> cfhs_;
std::vector<Iterator*> iterators_;
std::vector<std::pair<ColumnFamilyHandle*, std::unique_ptr<Iterator>>>
cfh_iter_pairs_;
ReadOptions read_options_;
Status status_;

Expand All @@ -60,11 +54,11 @@ class MultiCfIteratorImpl : public MultiCfIterator {

class MultiCfMinHeapItemComparator {
public:
MultiCfMinHeapItemComparator() {}
explicit MultiCfMinHeapItemComparator(const Comparator* comparator)
: comparator_(comparator) {}

bool operator()(MultiCfIteratorInfo a, MultiCfIteratorInfo b) const {
bool operator()(const MultiCfIteratorInfo& a,
const MultiCfIteratorInfo& b) const {
assert(a.iterator);
assert(b.iterator);
assert(a.iterator->Valid());
Expand Down
156 changes: 65 additions & 91 deletions db/multi_cf_iterator_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -13,25 +13,55 @@ class MultiCfIteratorTest : public DBTestBase {
public:
MultiCfIteratorTest()
: DBTestBase("multi_cf_iterator_test", /*env_do_fsync=*/true) {}
};

TEST_F(MultiCfIteratorTest, SimpleValues) {
Options options = GetDefaultOptions();
auto verify = [&](const std::vector<ColumnFamilyHandle*>& cfhs,
const std::vector<Slice>& expected_keys,
const std::vector<Slice>& expected_values) {
void verifyMultiCfIterator(
const std::vector<ColumnFamilyHandle*>& cfhs,
const std::vector<Slice>& expected_keys,
const std::optional<std::vector<Slice>>& expected_values = std::nullopt,
const std::optional<std::vector<WideColumns>>& expected_wide_columns =
std::nullopt,
const std::optional<std::vector<AttributeGroups>>&
expected_attribute_groups = std::nullopt) {
int i = 0;
std::unique_ptr<MultiCfIterator> iter =
db_->NewMultiCfIterator(ReadOptions(), cfhs);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
ASSERT_EQ(expected_values[i], iter->value());
if (expected_values.has_value()) {
ASSERT_EQ(expected_values.value()[i], iter->value());
}
if (expected_wide_columns.has_value()) {
ASSERT_EQ(expected_wide_columns.value()[i], iter->columns());
}
if (expected_attribute_groups.has_value()) {
ASSERT_EQ(expected_attribute_groups.value()[i],
iter->attribute_groups());
}
if (expected_values.has_value()) {
ASSERT_EQ(expected_values.value()[i], iter->value());
}
++i;
}
ASSERT_EQ(expected_keys.size(), i);
ASSERT_OK(iter->status());
}

void verifyExpectedKeys(ColumnFamilyHandle* cfh,
const std::vector<Slice>& expected_keys) {
int i = 0;
Iterator* iter = db_->NewIterator(ReadOptions(), cfh);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
++i;
}
ASSERT_EQ(expected_keys.size(), i);
};
ASSERT_OK(iter->status());
delete iter;
}
};

TEST_F(MultiCfIteratorTest, SimpleValues) {
Options options = GetDefaultOptions();
{
// Case 1: Unique key per CF
CreateAndReopenWithCF({"cf_1", "cf_2", "cf_3"}, options);
Expand All @@ -48,14 +78,14 @@ TEST_F(MultiCfIteratorTest, SimpleValues) {
// Test for iteration over CF default->1->2->3
std::vector<ColumnFamilyHandle*> cfhs_order_0_1_2_3 = {
handles_[0], handles_[1], handles_[2], handles_[3]};
verify(cfhs_order_0_1_2_3, expected_keys, expected_values);
verifyMultiCfIterator(cfhs_order_0_1_2_3, expected_keys, expected_values);

// Test for iteration over CF 3->1->default_cf->2
std::vector<ColumnFamilyHandle*> cfhs_order_3_1_0_2 = {
handles_[3], handles_[1], handles_[0], handles_[2]};
// Iteration order and the return values should be the same since keys are
// unique per CF
verify(cfhs_order_3_1_0_2, expected_keys, expected_values);
verifyMultiCfIterator(cfhs_order_3_1_0_2, expected_keys, expected_values);
}
{
// Case 2: Same key in multiple CFs
Expand All @@ -78,13 +108,13 @@ TEST_F(MultiCfIteratorTest, SimpleValues) {
handles_[0], handles_[1], handles_[2], handles_[3]};
std::vector<Slice> expected_values = {"key_1_cf_0_val", "key_2_cf_1_val",
"key_3_cf_0_val"};
verify(cfhs_order_0_1_2_3, expected_keys, expected_values);
verifyMultiCfIterator(cfhs_order_0_1_2_3, expected_keys, expected_values);

// Test for iteration over CFs 3->2->default_cf->1
std::vector<ColumnFamilyHandle*> cfhs_order_3_2_0_1 = {
handles_[3], handles_[2], handles_[0], handles_[1]};
expected_values = {"key_1_cf_3_val", "key_2_cf_2_val", "key_3_cf_3_val"};
verify(cfhs_order_3_2_0_1, expected_keys, expected_values);
verifyMultiCfIterator(cfhs_order_3_2_0_1, expected_keys, expected_values);
}
}

Expand Down Expand Up @@ -141,23 +171,6 @@ TEST_F(MultiCfIteratorTest, WideColumns) {
ASSERT_OK(db_->PutEntity(WriteOptions(), key_3, key_3_attribute_groups));
ASSERT_OK(db_->PutEntity(WriteOptions(), key_4, key_4_attribute_groups));

auto verify = [&](const std::vector<ColumnFamilyHandle*>& cfhs,
const std::vector<Slice>& expected_keys,
const std::vector<Slice>& expected_values,
const std::vector<WideColumns>& expected_wide_columns) {
int i = 0;
std::unique_ptr<MultiCfIterator> iter =
db_->NewMultiCfIterator(ReadOptions(), cfhs);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
ASSERT_EQ(expected_values[i], iter->value());
ASSERT_EQ(expected_wide_columns[i], iter->columns());
++i;
}
ASSERT_OK(iter->status());
ASSERT_EQ(expected_keys.size(), i);
};

// Test for iteration over CF default->1->2->3
std::vector<ColumnFamilyHandle*> cfhs_order_0_1_2_3 = {
handles_[0], handles_[1], handles_[2], handles_[3]};
Expand All @@ -176,11 +189,11 @@ TEST_F(MultiCfIteratorTest, WideColumns) {
{{"cf_1_col_name_1", "cf_1_col_val_1_key_2"}},
{{"cf_1_col_name_1", "cf_1_col_val_1_key_3"}},
{{"cf_0_col_name_1", "cf_0_col_val_1_key_4"}}};
verify(cfhs_order_0_1_2_3, expected_keys, expected_values,
expected_wide_columns);
verifyMultiCfIterator(cfhs_order_0_1_2_3, expected_keys, expected_values,
expected_wide_columns);
}

TEST_F(MultiCfIteratorTest, DifferentComparatorsinMultiCFs) {
TEST_F(MultiCfIteratorTest, DifferentComparatorsInMultiCFs) {
// This test creates two column families with two different comparators.
// Attempting to create the MultiCFIterator should fail.
Options options = GetDefaultOptions();
Expand All @@ -198,20 +211,8 @@ TEST_F(MultiCfIteratorTest, DifferentComparatorsinMultiCFs) {
ASSERT_OK(Put(1, "key_2", "value_2"));
ASSERT_OK(Put(1, "key_3", "value_3"));

auto verify = [&](ColumnFamilyHandle* cfh,
const std::vector<Slice>& expected_keys) {
int i = 0;
Iterator* iter = db_->NewIterator(ReadOptions(), cfh);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
++i;
}
ASSERT_EQ(expected_keys.size(), i);
ASSERT_OK(iter->status());
delete iter;
};
verify(handles_[0], {"key_1", "key_2", "key_3"});
verify(handles_[1], {"key_3", "key_2", "key_1"});
verifyExpectedKeys(handles_[0], {"key_1", "key_2", "key_3"});
verifyExpectedKeys(handles_[1], {"key_3", "key_2", "key_1"});

std::unique_ptr<MultiCfIterator> iter =
db_->NewMultiCfIterator(ReadOptions(), handles_);
Expand All @@ -221,18 +222,22 @@ TEST_F(MultiCfIteratorTest, DifferentComparatorsinMultiCFs) {

TEST_F(MultiCfIteratorTest, CustomComparatorsInMultiCFs) {
// This test creates two column families with the same custom test
// comparators (but instantiated differently). Attempting to create the
// comparators (but instantiated independently). Attempting to create the
// MultiCFIterator should not fail.
Options options = GetDefaultOptions();
options.create_if_missing = true;
DestroyAndReopen(options);
Comparator* comparator_1 = new test::SimpleSuffixReverseComparator();
Comparator* comparator_2 = new test::SimpleSuffixReverseComparator();
static auto comparator_1 =
std::make_unique<test::SimpleSuffixReverseComparator>(
test::SimpleSuffixReverseComparator());
static auto comparator_2 =
std::make_unique<test::SimpleSuffixReverseComparator>(
test::SimpleSuffixReverseComparator());
ASSERT_NE(comparator_1, comparator_2);

options.comparator = comparator_1;
options.comparator = comparator_1.get();
CreateColumnFamilies({"cf_1"}, options);
options.comparator = comparator_2;
options.comparator = comparator_2.get();
CreateColumnFamilies({"cf_2"}, options);

ASSERT_OK(Put(0, "key_001_001", "value_0_3"));
Expand All @@ -248,22 +253,12 @@ TEST_F(MultiCfIteratorTest, CustomComparatorsInMultiCFs) {
ASSERT_OK(Put(1, "key_003_005", "value_1_5"));
ASSERT_OK(Put(1, "key_003_006", "value_1_4"));

auto verify = [&](ColumnFamilyHandle* cfh,
const std::vector<Slice>& expected_keys) {
int i = 0;
Iterator* iter = db_->NewIterator(ReadOptions(), cfh);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
++i;
}
ASSERT_TRUE(iter->status().ok());
ASSERT_EQ(expected_keys.size(), i);
delete iter;
};
verify(handles_[0], {"key_001_003", "key_001_002", "key_001_001",
"key_002_003", "key_002_002", "key_002_001"});
verify(handles_[1], {"key_001_003", "key_001_002", "key_001_001",
"key_003_006", "key_003_005", "key_003_004"});
verifyExpectedKeys(
handles_[0], {"key_001_003", "key_001_002", "key_001_001", "key_002_003",
"key_002_002", "key_002_001"});
verifyExpectedKeys(
handles_[1], {"key_001_003", "key_001_002", "key_001_001", "key_003_006",
"key_003_005", "key_003_004"});

std::vector<Slice> expected_keys = {
"key_001_003", "key_001_002", "key_001_001", "key_002_003", "key_002_002",
Expand All @@ -280,12 +275,6 @@ TEST_F(MultiCfIteratorTest, CustomComparatorsInMultiCFs) {
++i;
}
ASSERT_OK(iter->status());
if (comparator_1) {
delete comparator_1;
}
if (comparator_2) {
delete comparator_2;
}
}

TEST_F(MultiCfIteratorTest, DISABLED_IterateAttributeGroups) {
Expand Down Expand Up @@ -340,31 +329,16 @@ TEST_F(MultiCfIteratorTest, DISABLED_IterateAttributeGroups) {
ASSERT_OK(db_->PutEntity(WriteOptions(), key_3, key_3_attribute_groups));
ASSERT_OK(db_->PutEntity(WriteOptions(), key_4, key_4_attribute_groups));

auto verify =
[&](const std::vector<ColumnFamilyHandle*>& cfhs,
const std::vector<Slice>& expected_keys,
const std::vector<AttributeGroups>& expected_attribute_groups) {
int i = 0;
std::unique_ptr<MultiCfIterator> iter =
db_->NewMultiCfIterator(ReadOptions(), cfhs);
for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
ASSERT_EQ(expected_keys[i], iter->key());
ASSERT_EQ(expected_attribute_groups[i], iter->attribute_groups());
++i;
}
ASSERT_OK(iter->status());
// TODO - Check when implementation is done
// ASSERT_EQ(expected_keys.size(), i);
};

// Test for iteration over CF default->1->2->3
std::vector<ColumnFamilyHandle*> cfhs_order_0_1_2_3 = {
handles_[0], handles_[1], handles_[2], handles_[3]};
std::vector<Slice> expected_keys = {key_1, key_2, key_3, key_4};
std::vector<AttributeGroups> expected_attribute_groups = {
key_1_attribute_groups, key_2_attribute_groups, key_3_attribute_groups,
key_4_attribute_groups};
verify(cfhs_order_0_1_2_3, expected_keys, expected_attribute_groups);
verifyMultiCfIterator(
cfhs_order_0_1_2_3, expected_keys, std::nullopt /* expected_values */,
std::nullopt /* expected_wide_columns */, expected_attribute_groups);
}

} // namespace ROCKSDB_NAMESPACE
Expand Down

0 comments on commit 15ba79f

Please sign in to comment.