perf(segment): reuse first vector index file as merge base during compaction (#440)
Add a fast path that copies the first segment's index file as the merge base and only merges the tail segments into it. Limited to streaming indexes (HNSW, HNSW_RABITQ, FLAT) with matching index_type + quantize_type and no filter; IVF/VAMANA always rebuild (their Merge is dump-then-reopen and would drop the base docs). Also, fix the incorrect concurrency of the compaction task.
This commit is contained in:
Jalin Wang
GitHub
parent
355523a628
commit
95e5ad5105
|
|
@ -146,7 +146,19 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) {
|
|||
|
||||
std::vector<int> read_results(streamers_.size(), -1);
|
||||
// TODO: use id instead of key
|
||||
uint32_t id_offset = 0, next_id = 0;
|
||||
// When merging into a non-empty target (e.g. reusing one input as base),
|
||||
// append new docs after the existing ones instead of overwriting from 0.
|
||||
uint32_t id_offset = 0;
|
||||
uint32_t next_id = 0;
|
||||
if (target_builder_ == nullptr) {
|
||||
if (is_sparse_) {
|
||||
auto provider = target_streamer_->create_sparse_provider();
|
||||
if (provider) next_id = provider->count();
|
||||
} else {
|
||||
auto provider = target_streamer_->create_provider();
|
||||
if (provider) next_id = provider->count();
|
||||
}
|
||||
}
|
||||
|
||||
if (is_sparse_) {
|
||||
for (size_t i = 0; i < num_of_add_threads_; i++) {
|
||||
|
|
|
|||
|
|
@ -98,6 +98,10 @@ class VectorColumnIndexer {
|
|||
return index_file_path_;
|
||||
}
|
||||
|
||||
const FieldSchema &field_schema() const {
|
||||
return field_schema_;
|
||||
}
|
||||
|
||||
size_t doc_count() const {
|
||||
if (index == nullptr) {
|
||||
return -1;
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@
|
|||
#include <arrow/compute/api_vector.h>
|
||||
#include <arrow/type_fwd.h>
|
||||
#include <zvec/ailego/logger/logger.h>
|
||||
#include <zvec/db/config.h>
|
||||
#include <zvec/db/status.h>
|
||||
#include <zvec/db/type.h>
|
||||
#if RABITQ_SUPPORTED
|
||||
|
|
@ -129,10 +130,13 @@ Status SegmentHelper::ExecuteCompactTask(CompactTask &task) {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
// RowIdFilter copies the bitmap so ReduceFts below can reuse it; sharing
|
||||
// lets the FTS reducer skip its own per-doc dense rank table.
|
||||
std::shared_ptr<RowIdFilter> row_id_filter =
|
||||
std::make_shared<RowIdFilter>(delete_row_id_bitmap);
|
||||
// Leave row_id_filter null when there are no deletes or
|
||||
// create_compaction_task with rebuild=false, so downstream merge
|
||||
// can take a faster per-doc path that skips the filter callback entirely.
|
||||
std::shared_ptr<RowIdFilter> row_id_filter;
|
||||
if (!delete_row_id_bitmap.isEmpty()) {
|
||||
row_id_filter = std::make_shared<RowIdFilter>(delete_row_id_bitmap);
|
||||
}
|
||||
|
||||
s = ReduceVectorIndex(schema, input_segments, output_segment_path,
|
||||
row_id_filter, block_id_generator, min_doc_id,
|
||||
|
|
@ -636,40 +640,28 @@ Status SegmentHelper::ReduceVectorIndex(
|
|||
auto vector_index_params =
|
||||
std::dynamic_pointer_cast<VectorIndexParams>(field->index_params());
|
||||
|
||||
using FetchIndexersFn =
|
||||
std::vector<VectorColumnIndexer::Ptr> (Segment::*)(const std::string &)
|
||||
const;
|
||||
auto collect_merge_indexers = [&](FetchIndexersFn fetch) {
|
||||
std::vector<VectorColumnIndexer::Ptr> source_indexers;
|
||||
for (const auto &seg : input_segments) {
|
||||
auto seg_indexers = (seg.get()->*fetch)(field->name());
|
||||
source_indexers.insert(source_indexers.end(), seg_indexers.begin(),
|
||||
seg_indexers.end());
|
||||
}
|
||||
return source_indexers;
|
||||
};
|
||||
|
||||
auto vector_block_id = block_id_generator();
|
||||
if (vector_index_params->quantize_type() == QuantizeType::UNDEFINED) {
|
||||
auto vector_index_path = FileHelper::MakeVectorIndexPath(
|
||||
output_segment_path, field->name(), vector_block_id);
|
||||
|
||||
// only create original vector indexer
|
||||
auto vector_indexer =
|
||||
std::make_shared<VectorColumnIndexer>(vector_index_path, *field);
|
||||
s = vector_indexer->Open({true, true});
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
std::vector<VectorColumnIndexer::Ptr> merge_indexers;
|
||||
for (auto &input_segment : input_segments) {
|
||||
// merge_indexers should be ordered put
|
||||
auto to_merge_indexers =
|
||||
input_segment->get_vector_indexer(field->name());
|
||||
merge_indexers.insert(merge_indexers.end(), to_merge_indexers.begin(),
|
||||
to_merge_indexers.end());
|
||||
}
|
||||
|
||||
vector_column_params::MergeOptions merge_options;
|
||||
if (concurrency == 0) {
|
||||
merge_options.pool = GlobalResource::Instance().optimize_thread_pool();
|
||||
} else {
|
||||
merge_options.write_concurrency = concurrency;
|
||||
}
|
||||
|
||||
s = vector_indexer->Merge(merge_indexers, filter, merge_options);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer->Flush();
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer->Close();
|
||||
s = MergeWithOptionalReuse(
|
||||
vector_index_path, *field,
|
||||
collect_merge_indexers(&Segment::get_vector_indexer), filter,
|
||||
concurrency, nullptr);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
BlockMeta new_block_meta;
|
||||
|
|
@ -689,32 +681,11 @@ Status SegmentHelper::ReduceVectorIndex(
|
|||
field_without_quantize->set_index_params(
|
||||
MakeDefaultVectorIndexParams(vector_index_params->metric_type()));
|
||||
|
||||
// create flat index
|
||||
auto vector_indexer = std::make_shared<VectorColumnIndexer>(
|
||||
vector_index_path, *field_without_quantize);
|
||||
s = vector_indexer->Open({true, true});
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
std::vector<VectorColumnIndexer::Ptr> merge_indexers;
|
||||
for (auto &input_segment : input_segments) {
|
||||
// merge_indexers should be ordered put
|
||||
auto to_merge_indexers =
|
||||
input_segment->get_vector_indexer(field->name());
|
||||
merge_indexers.insert(merge_indexers.end(), to_merge_indexers.begin(),
|
||||
to_merge_indexers.end());
|
||||
}
|
||||
|
||||
vector_column_params::MergeOptions merge_options;
|
||||
if (concurrency == 0) {
|
||||
merge_options.pool = GlobalResource::Instance().optimize_thread_pool();
|
||||
} else {
|
||||
merge_options.write_concurrency = concurrency;
|
||||
}
|
||||
|
||||
s = vector_indexer->Merge(merge_indexers, filter, merge_options);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer->Flush();
|
||||
VectorColumnIndexer::Ptr vector_indexer;
|
||||
s = MergeWithOptionalReuse(
|
||||
vector_index_path, *field_without_quantize,
|
||||
collect_merge_indexers(&Segment::get_vector_indexer), filter,
|
||||
concurrency, &vector_indexer);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
// The training step (for RABITQ) and the subsequent quantize merge both
|
||||
|
|
@ -745,27 +716,10 @@ Status SegmentHelper::ReduceVectorIndex(
|
|||
auto vector_quan_index_path = FileHelper::MakeQuantizeVectorIndexPath(
|
||||
output_segment_path, field->name(), vector_quan_block_id);
|
||||
|
||||
auto vector_indexer_quantize = std::make_shared<VectorColumnIndexer>(
|
||||
vector_quan_index_path, *field_for_quantize);
|
||||
s = vector_indexer_quantize->Open({true, true});
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
merge_indexers.clear();
|
||||
for (auto &input_segment : input_segments) {
|
||||
// merge_indexers should be ordered put
|
||||
auto to_merge_indexers =
|
||||
input_segment->get_quant_vector_indexer(field->name());
|
||||
merge_indexers.insert(merge_indexers.end(), to_merge_indexers.begin(),
|
||||
to_merge_indexers.end());
|
||||
}
|
||||
|
||||
s = vector_indexer_quantize->Merge(merge_indexers, filter, merge_options);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer_quantize->Flush();
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer_quantize->Close();
|
||||
s = MergeWithOptionalReuse(
|
||||
vector_quan_index_path, *field_for_quantize,
|
||||
collect_merge_indexers(&Segment::get_quant_vector_indexer), filter,
|
||||
concurrency, nullptr);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer->Close();
|
||||
|
|
@ -784,6 +738,100 @@ Status SegmentHelper::ReduceVectorIndex(
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
namespace {
|
||||
|
||||
// Only the first indexer's file is reused as the merge base; the remaining
|
||||
// indexers are merged in via Merge(). Reuse is restricted to streaming
|
||||
// indexes (HNSW, FLAT). Builder-rebuild indexes (IVF, VAMANA) and HNSW_RABITQ
|
||||
// fall back to the full-rebuild merge.
|
||||
bool CanReuseFirstIndexer(const std::vector<VectorColumnIndexer::Ptr> &indexers,
|
||||
const FieldSchema &output_field,
|
||||
const IndexFilter::Ptr &filter) {
|
||||
if (filter != nullptr || indexers.empty()) {
|
||||
return false;
|
||||
}
|
||||
if (output_field.index_type() != IndexType::HNSW &&
|
||||
output_field.index_type() != IndexType::FLAT) {
|
||||
return false;
|
||||
}
|
||||
const auto &first_field = indexers.front()->field_schema();
|
||||
if (first_field.index_type() != output_field.index_type()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// When creating a quantized column, the fp32 indexers will be merged into a
|
||||
// quantized vector indexer.
|
||||
auto quantize_type_of = [](const FieldSchema &f) {
|
||||
auto params =
|
||||
std::dynamic_pointer_cast<VectorIndexParams>(f.index_params());
|
||||
return params ? params->quantize_type() : QuantizeType::UNDEFINED;
|
||||
};
|
||||
return quantize_type_of(first_field) == quantize_type_of(output_field);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
Status SegmentHelper::MergeWithOptionalReuse(
|
||||
const std::string &output_index_path, const FieldSchema &index_field,
|
||||
std::vector<VectorColumnIndexer::Ptr> source_indexers,
|
||||
const IndexFilter::Ptr &filter, int concurrency,
|
||||
VectorColumnIndexer::Ptr *merged_indexer) {
|
||||
vector_column_params::MergeOptions merge_options;
|
||||
if (concurrency == 0) {
|
||||
merge_options.pool = GlobalResource::Instance().optimize_thread_pool();
|
||||
merge_options.write_concurrency =
|
||||
GlobalConfig::Instance().optimize_thread_count();
|
||||
} else {
|
||||
merge_options.write_concurrency = concurrency;
|
||||
}
|
||||
auto vector_indexer =
|
||||
std::make_shared<VectorColumnIndexer>(output_index_path, index_field);
|
||||
bool reused_base_index = false;
|
||||
Status s;
|
||||
|
||||
if (CanReuseFirstIndexer(source_indexers, index_field, filter)) {
|
||||
const auto &first_indexer = source_indexers.front();
|
||||
LOG_INFO(
|
||||
"Reusing first indexer as merge base. "
|
||||
"field[%s] src[%s] dst[%s] tail_indexers[%zu]",
|
||||
index_field.name().c_str(), first_indexer->index_file_path().c_str(),
|
||||
output_index_path.c_str(), source_indexers.size() - 1);
|
||||
if (FileHelper::CopyFile(first_indexer->index_file_path(),
|
||||
output_index_path)) {
|
||||
// Open the copied file in-place (create_new=false).
|
||||
s = vector_indexer->Open(vector_column_params::ReadOptions{true, false});
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
source_indexers.erase(source_indexers.begin());
|
||||
s = vector_indexer->Merge(source_indexers, filter, merge_options);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
reused_base_index = true;
|
||||
} else {
|
||||
LOG_WARN("Failed to copy %s to %s; falling back to full rebuild",
|
||||
first_indexer->index_file_path().c_str(),
|
||||
output_index_path.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
if (!reused_base_index) {
|
||||
s = vector_indexer->Open(vector_column_params::ReadOptions{true, true});
|
||||
CHECK_RETURN_STATUS(s);
|
||||
|
||||
s = vector_indexer->Merge(source_indexers, filter, merge_options);
|
||||
CHECK_RETURN_STATUS(s);
|
||||
}
|
||||
|
||||
s = vector_indexer->Flush();
|
||||
CHECK_RETURN_STATUS(s);
|
||||
if (merged_indexer != nullptr) {
|
||||
*merged_indexer = vector_indexer;
|
||||
} else {
|
||||
s = vector_indexer->Close();
|
||||
CHECK_RETURN_STATUS(s);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status SegmentHelper::PrepareQuantizeField(
|
||||
const FieldSchema &field,
|
||||
const core::IndexProvider::Pointer &raw_vector_provider,
|
||||
|
|
|
|||
|
|
@ -217,6 +217,18 @@ class SegmentHelper {
|
|||
uint64_t max_doc_id, uint32_t doc_count, int concurrency,
|
||||
std::vector<BlockMeta> *output_block_metas);
|
||||
|
||||
// Merges `source_indexers` into a new VectorColumnIndexer at
|
||||
// `output_index_path`. When the first indexer is eligible for reuse (see
|
||||
// CanReuseFirstIndexer), its file is copied to the
|
||||
// output path and opened in-place as the merge base. If `merged_indexer` is
|
||||
// non-null, it receives the opened indexer (caller owns Close()); otherwise
|
||||
// the indexer is closed before returning.
|
||||
static Status MergeWithOptionalReuse(
|
||||
const std::string &output_index_path, const FieldSchema &index_field,
|
||||
std::vector<VectorColumnIndexer::Ptr> source_indexers,
|
||||
const IndexFilter::Ptr &filter, int concurrency,
|
||||
VectorColumnIndexer::Ptr *merged_indexer);
|
||||
|
||||
// Returns a FieldSchema clone whose index_params is ready for building the
|
||||
// quantize indexer.
|
||||
// - RABITQ: clones HnswRabitqIndexParams, trains a RabitqConverter against
|
||||
|
|
|
|||
|
|
@ -13,10 +13,12 @@
|
|||
// limitations under the License.
|
||||
|
||||
#include "db/index/segment/segment_helper.h"
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <thread>
|
||||
#include <variant>
|
||||
#include <arrow/array/array_binary.h>
|
||||
|
|
@ -28,6 +30,9 @@
|
|||
#include <gtest/gtest.h>
|
||||
#include "db/common/constants.h"
|
||||
#include "db/common/file_helper.h"
|
||||
#include "db/index/column/vector_column/vector_column_indexer.h"
|
||||
#include "db/index/column/vector_column/vector_column_params.h"
|
||||
#include "db/index/column/vector_column/vector_index_results.h"
|
||||
#include "db/index/common/delete_store.h"
|
||||
#include "db/index/common/id_map.h"
|
||||
#include "db/index/common/meta.h"
|
||||
|
|
@ -35,6 +40,7 @@
|
|||
#include "db/index/segment/segment.h"
|
||||
#include "utils/utils.h"
|
||||
#include "zvec/db/options.h"
|
||||
#include "zvec/db/query_params.h"
|
||||
#include "zvec/db/schema.h"
|
||||
|
||||
using namespace zvec;
|
||||
|
|
@ -72,6 +78,64 @@ class SegmentHelperTest : public testing::Test {
|
|||
}
|
||||
|
||||
protected:
|
||||
VersionManager::Ptr CreateVersionManager(const CollectionSchema &schema) {
|
||||
Version version;
|
||||
version.set_schema(schema);
|
||||
auto vm = VersionManager::Create(col_path, version);
|
||||
if (!vm.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
return vm.value();
|
||||
}
|
||||
|
||||
SegmentOptions WriteOptions() const {
|
||||
return SegmentOptions{false, true, DEFAULT_MAX_BUFFER_SIZE};
|
||||
}
|
||||
|
||||
struct CompactResult {
|
||||
CompactTask compact_task;
|
||||
Segment::Ptr output_segment; // null when filter dropped every doc
|
||||
};
|
||||
|
||||
// Execute a CompactTask end-to-end: build it, run it, move the tmp segment
|
||||
// dir into place, and reopen the output segment in read-only mode.
|
||||
CompactResult RunCompactAndOpen(CollectionSchema::Ptr schema,
|
||||
std::vector<Segment::Ptr> segments,
|
||||
SegmentID output_segment_id,
|
||||
IndexFilter::Ptr filter,
|
||||
const VersionManager::Ptr &version_manager,
|
||||
int concurrency = 1) {
|
||||
const bool forward_use_parquet = false;
|
||||
CompactTask task(col_path, schema, std::move(segments), output_segment_id,
|
||||
std::move(filter), forward_use_parquet, concurrency);
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
EXPECT_NE(segment_task, nullptr);
|
||||
if (segment_task == nullptr) return {task, nullptr};
|
||||
|
||||
auto status = SegmentHelper::Execute(segment_task);
|
||||
EXPECT_TRUE(status.ok()) << status.message();
|
||||
|
||||
auto executed = std::get<CompactTask>(segment_task->task_info());
|
||||
if (executed.output_segment_meta_ == nullptr) {
|
||||
return {executed, nullptr};
|
||||
}
|
||||
|
||||
auto tmp_path =
|
||||
FileHelper::MakeTempSegmentPath(col_path, output_segment_id);
|
||||
auto dst_path = FileHelper::MakeSegmentPath(col_path, output_segment_id);
|
||||
EXPECT_TRUE(FileHelper::MoveDirectory(tmp_path, dst_path));
|
||||
|
||||
SegmentOptions read_options{true, !forward_use_parquet,
|
||||
DEFAULT_MAX_BUFFER_SIZE};
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg_ret =
|
||||
Segment::Open(col_path, *schema, *executed.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, read_options);
|
||||
EXPECT_TRUE(seg_ret.has_value());
|
||||
if (!seg_ret.has_value()) return {executed, nullptr};
|
||||
return {executed, std::move(seg_ret.value())};
|
||||
}
|
||||
|
||||
std::string col_name = "test_segment_helper";
|
||||
std::string col_path = "./test_collection";
|
||||
IDMap::Ptr id_map;
|
||||
|
|
@ -80,84 +144,29 @@ class SegmentHelperTest : public testing::Test {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_General) {
|
||||
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
|
||||
// Create segments
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
seg_options, 0, 1000);
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 1000);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
seg_options, 1000, 1000);
|
||||
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
write_options, 1000, 1000);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
std::cout << "seg2: " << seg2->meta()->to_string_formatted() << std::endl;
|
||||
|
||||
// Prepare segments for compaction
|
||||
std::vector<Segment::Ptr> segments = {seg1, seg2};
|
||||
|
||||
// Create compact task
|
||||
SegmentID output_segment_id = 2;
|
||||
CompactTask task(GetColPath(), schema, segments,
|
||||
output_segment_id, // output_segment_id
|
||||
nullptr, // filter
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, seg3] = RunCompactAndOpen(
|
||||
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
|
||||
|
||||
// Create segment task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
|
||||
// Verify task creation
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
// Execute the task
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
ASSERT_EQ(output_segment_meta->id(), output_segment_id);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
|
||||
// Move segment directory
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
auto new_segment_path =
|
||||
FileHelper::MakeSegmentPath(GetColPath(), output_segment_id);
|
||||
FileHelper::MoveDirectory(tmp_segment_path, new_segment_path);
|
||||
|
||||
seg_options.read_only_ = true;
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg3_ret = Segment::Open(
|
||||
GetColPath(), *schema, *segment_compact_task.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, seg_options);
|
||||
if (!seg3_ret.has_value()) {
|
||||
std::cout << seg3_ret.error().message() << std::endl;
|
||||
ASSERT_TRUE(false);
|
||||
}
|
||||
|
||||
auto seg3 = std::move(seg3_ret.value());
|
||||
ASSERT_NE(seg3, nullptr);
|
||||
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
|
||||
ASSERT_FALSE(
|
||||
compact_task.output_segment_meta_->writing_forward_block().has_value());
|
||||
ASSERT_EQ(seg3->id(), output_segment_id);
|
||||
|
||||
std::cout << seg3->meta()->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
|
||||
|
||||
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
|
||||
|
|
@ -173,84 +182,29 @@ TEST_F(SegmentHelperTest, CompactTask_General) {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_ScalarIndex) {
|
||||
auto schema = test::TestHelper::CreateSchemaWithScalarIndex(false);
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
|
||||
// Create segments
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
seg_options, 0, 1000);
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 1000);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
seg_options, 1000, 1000);
|
||||
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
write_options, 1000, 1000);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
std::cout << "seg2: " << seg2->meta()->to_string_formatted() << std::endl;
|
||||
|
||||
// Prepare segments for compaction
|
||||
std::vector<Segment::Ptr> segments = {seg1, seg2};
|
||||
|
||||
// Create compact task
|
||||
SegmentID output_segment_id = 2;
|
||||
CompactTask task(GetColPath(), schema, segments,
|
||||
output_segment_id, // output_segment_id
|
||||
nullptr, // filter
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, seg3] = RunCompactAndOpen(
|
||||
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
|
||||
|
||||
// Create segment task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
|
||||
// Verify task creation
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
// Execute the task
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
ASSERT_EQ(output_segment_meta->id(), output_segment_id);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
|
||||
// Move segment directory
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
auto new_segment_path =
|
||||
FileHelper::MakeSegmentPath(GetColPath(), output_segment_id);
|
||||
FileHelper::MoveDirectory(tmp_segment_path, new_segment_path);
|
||||
|
||||
seg_options.read_only_ = true;
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg3_ret = Segment::Open(
|
||||
GetColPath(), *schema, *segment_compact_task.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, seg_options);
|
||||
if (!seg3_ret.has_value()) {
|
||||
std::cout << seg3_ret.error().message() << std::endl;
|
||||
ASSERT_TRUE(false);
|
||||
}
|
||||
|
||||
auto seg3 = std::move(seg3_ret.value());
|
||||
ASSERT_NE(seg3, nullptr);
|
||||
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
|
||||
ASSERT_FALSE(
|
||||
compact_task.output_segment_meta_->writing_forward_block().has_value());
|
||||
ASSERT_EQ(seg3->id(), output_segment_id);
|
||||
|
||||
std::cout << seg3->meta()->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
|
||||
|
||||
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
|
||||
|
|
@ -266,84 +220,29 @@ TEST_F(SegmentHelperTest, CompactTask_ScalarIndex) {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_VectorIndex) {
|
||||
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
|
||||
// Create segments
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
seg_options, 0, 1000);
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 1000);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
seg_options, 1000, 1000);
|
||||
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
|
||||
write_options, 1000, 1000);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
std::cout << "seg2: " << seg2->meta()->to_string_formatted() << std::endl;
|
||||
|
||||
// Prepare segments for compaction
|
||||
std::vector<Segment::Ptr> segments = {seg1, seg2};
|
||||
|
||||
// Create compact task
|
||||
SegmentID output_segment_id = 2;
|
||||
CompactTask task(GetColPath(), schema, segments,
|
||||
output_segment_id, // output_segment_id
|
||||
nullptr, // filter
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, seg3] = RunCompactAndOpen(
|
||||
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
|
||||
|
||||
// Create segment task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
|
||||
// Verify task creation
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
// Execute the task
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
ASSERT_EQ(output_segment_meta->id(), output_segment_id);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
|
||||
// Move segment directory
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
auto new_segment_path =
|
||||
FileHelper::MakeSegmentPath(GetColPath(), output_segment_id);
|
||||
FileHelper::MoveDirectory(tmp_segment_path, new_segment_path);
|
||||
|
||||
seg_options.read_only_ = true;
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg3_ret = Segment::Open(
|
||||
GetColPath(), *schema, *segment_compact_task.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, seg_options);
|
||||
if (!seg3_ret.has_value()) {
|
||||
std::cout << seg3_ret.error().message() << std::endl;
|
||||
ASSERT_TRUE(false);
|
||||
}
|
||||
|
||||
auto seg3 = std::move(seg3_ret.value());
|
||||
ASSERT_NE(seg3, nullptr);
|
||||
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
|
||||
ASSERT_FALSE(
|
||||
compact_task.output_segment_meta_->writing_forward_block().has_value());
|
||||
ASSERT_EQ(seg3->id(), output_segment_id);
|
||||
|
||||
std::cout << seg3->meta()->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
|
||||
|
||||
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
|
||||
|
|
@ -359,86 +258,35 @@ TEST_F(SegmentHelperTest, CompactTask_VectorIndex) {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_MultipleSegments) {
|
||||
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
std::vector<Segment::Ptr> input_segs;
|
||||
int seg_count = 10;
|
||||
int doc_count_per_seg = 100;
|
||||
const int seg_count = 10;
|
||||
const int doc_count_per_seg = 100;
|
||||
for (int i = 0; i < seg_count; i++) {
|
||||
auto seg = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, i, i * doc_count_per_seg, id_map, delete_store,
|
||||
version_manager, seg_options, i * doc_count_per_seg, doc_count_per_seg);
|
||||
col_path, *schema, i, i * doc_count_per_seg, id_map, delete_store,
|
||||
version_manager, write_options, i * doc_count_per_seg,
|
||||
doc_count_per_seg);
|
||||
ASSERT_TRUE(seg != nullptr);
|
||||
ASSERT_TRUE(seg->flush().ok());
|
||||
input_segs.push_back(seg);
|
||||
}
|
||||
|
||||
// Create compact task
|
||||
SegmentID output_segment_id = seg_count;
|
||||
CompactTask task(GetColPath(), schema, input_segs,
|
||||
output_segment_id, // output_segment_id
|
||||
nullptr, // filter
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, seg3] = RunCompactAndOpen(
|
||||
schema, input_segs, output_segment_id, nullptr, version_manager);
|
||||
|
||||
// Create segment task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
|
||||
// Verify task creation
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
// Execute the task
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
ASSERT_EQ(output_segment_meta->id(), output_segment_id);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
|
||||
// Move segment directory
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
auto new_segment_path =
|
||||
FileHelper::MakeSegmentPath(GetColPath(), output_segment_id);
|
||||
FileHelper::MoveDirectory(tmp_segment_path, new_segment_path);
|
||||
|
||||
seg_options.read_only_ = true;
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg3_ret = Segment::Open(
|
||||
GetColPath(), *schema, *segment_compact_task.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, seg_options);
|
||||
if (!seg3_ret.has_value()) {
|
||||
std::cout << seg3_ret.error().message() << std::endl;
|
||||
ASSERT_TRUE(false);
|
||||
}
|
||||
|
||||
auto seg3 = std::move(seg3_ret.value());
|
||||
ASSERT_NE(seg3, nullptr);
|
||||
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
|
||||
ASSERT_FALSE(
|
||||
compact_task.output_segment_meta_->writing_forward_block().has_value());
|
||||
ASSERT_EQ(seg3->id(), output_segment_id);
|
||||
|
||||
std::cout << seg3->meta()->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(seg3->doc_count(), seg_count * doc_count_per_seg);
|
||||
|
||||
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
|
||||
auto doc = seg3->Fetch(i);
|
||||
if (doc == nullptr) {
|
||||
std::cout << "doc is null: " << i << std::endl;
|
||||
}
|
||||
ASSERT_NE(doc, nullptr);
|
||||
auto expect_doc = test::TestHelper::CreateDoc(i, *schema);
|
||||
ASSERT_EQ(*doc, expect_doc);
|
||||
|
|
@ -447,78 +295,27 @@ TEST_F(SegmentHelperTest, CompactTask_MultipleSegments) {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_Filter) {
|
||||
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
|
||||
// Create segments
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
seg_options, 0, 1000);
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 1000);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
|
||||
// Create a simple filter
|
||||
auto filter = std::make_shared<EasyIndexFilter>(
|
||||
[&](uint64_t id) -> bool { return id < 10; });
|
||||
// Note: Actual filter configuration would depend on the IndexFilter
|
||||
// implementation
|
||||
[](uint64_t id) -> bool { return id < 10; });
|
||||
|
||||
// Create compact task with filter
|
||||
SegmentID output_segment_id = 1;
|
||||
CompactTask task(GetColPath(), schema, {seg1}, // Single segment with filter
|
||||
output_segment_id, // output_segment_id
|
||||
filter,
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, seg2] = RunCompactAndOpen(
|
||||
schema, {seg1}, output_segment_id, filter, version_manager);
|
||||
|
||||
// Create and execute task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
std::cout << output_segment_meta->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(output_segment_meta->id(), output_segment_id);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
|
||||
// Move segment directory
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
auto new_segment_path =
|
||||
FileHelper::MakeSegmentPath(GetColPath(), output_segment_id);
|
||||
FileHelper::MoveDirectory(tmp_segment_path, new_segment_path);
|
||||
|
||||
seg_options.read_only_ = true;
|
||||
version_manager->set_enable_mmap(!forward_use_parquet);
|
||||
auto seg2_ret = Segment::Open(
|
||||
GetColPath(), *schema, *segment_compact_task.output_segment_meta_, id_map,
|
||||
delete_store, version_manager, seg_options);
|
||||
if (!seg2_ret.has_value()) {
|
||||
std::cout << seg2_ret.error().message() << std::endl;
|
||||
ASSERT_TRUE(false);
|
||||
}
|
||||
|
||||
auto seg2 = std::move(seg2_ret.value());
|
||||
ASSERT_NE(seg2, nullptr);
|
||||
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
|
||||
ASSERT_FALSE(
|
||||
compact_task.output_segment_meta_->writing_forward_block().has_value());
|
||||
ASSERT_EQ(seg2->id(), output_segment_id);
|
||||
|
||||
std::cout << seg2->meta()->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(seg2->doc_count(), seg1->doc_count() - 10);
|
||||
|
||||
ASSERT_TRUE(seg1->destroy().ok());
|
||||
|
|
@ -526,58 +323,26 @@ TEST_F(SegmentHelperTest, CompactTask_Filter) {
|
|||
|
||||
TEST_F(SegmentHelperTest, CompactTask_FilterAll) {
|
||||
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
Version version;
|
||||
version.set_schema(*schema);
|
||||
auto version_manager_tmp = VersionManager::Create(col_path, version);
|
||||
if (!version_manager_tmp.has_value()) {
|
||||
throw std::runtime_error("Failed to create version manager");
|
||||
}
|
||||
|
||||
auto version_manager = version_manager_tmp.value();
|
||||
|
||||
bool forward_use_parquet = false;
|
||||
auto seg_options =
|
||||
SegmentOptions{false, !forward_use_parquet, DEFAULT_MAX_BUFFER_SIZE};
|
||||
|
||||
// Create segments
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
seg_options, 0, 1000);
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 1000);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
|
||||
// Create a simple filter
|
||||
auto filter = std::make_shared<EasyIndexFilter>(
|
||||
[&](uint64_t id) -> bool { return true; });
|
||||
// Note: Actual filter configuration would depend on the IndexFilter
|
||||
// implementation
|
||||
[](uint64_t /*id*/) -> bool { return true; });
|
||||
|
||||
// Create compact task with filter
|
||||
SegmentID output_segment_id = 1;
|
||||
CompactTask task(GetColPath(), schema, {seg1}, // Single segment with filter
|
||||
output_segment_id, // output_segment_id
|
||||
filter,
|
||||
forward_use_parquet, // forward_use_parquet
|
||||
1 // concurrency
|
||||
);
|
||||
auto [compact_task, output_segment] = RunCompactAndOpen(
|
||||
schema, {seg1}, output_segment_id, filter, version_manager);
|
||||
|
||||
// Create and execute task
|
||||
auto segment_task = SegmentTask::CreateCompactTask(task);
|
||||
ASSERT_TRUE(segment_task != nullptr);
|
||||
|
||||
Status status = SegmentHelper::Execute(segment_task);
|
||||
std::cout << "status: " << status.message() << std::endl;
|
||||
ASSERT_TRUE(status.ok());
|
||||
|
||||
auto segment_compact_task = std::get<CompactTask>(segment_task->task_info());
|
||||
// Verify output segment
|
||||
auto output_segment_meta = segment_compact_task.output_segment_meta_;
|
||||
ASSERT_EQ(output_segment_meta, nullptr);
|
||||
|
||||
auto tmp_segment_path =
|
||||
FileHelper::MakeTempSegmentPath(GetColPath(), output_segment_id);
|
||||
ASSERT_FALSE(FileHelper::DirectoryExists(tmp_segment_path));
|
||||
ASSERT_EQ(compact_task.output_segment_meta_, nullptr);
|
||||
ASSERT_EQ(output_segment, nullptr);
|
||||
ASSERT_FALSE(FileHelper::DirectoryExists(
|
||||
FileHelper::MakeTempSegmentPath(col_path, output_segment_id)));
|
||||
}
|
||||
|
||||
TEST_F(SegmentHelperTest, CreateVectorIndexTask_AllFields) {
|
||||
|
|
@ -691,4 +456,351 @@ TEST_F(SegmentHelperTest, CreateVectorIndexTask_SingleField) {
|
|||
<< output_segment_meta->to_string_formatted() << std::endl;
|
||||
ASSERT_EQ(output_segment_meta->id(), 0);
|
||||
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(SegmentHelperTest, CompactTask_VectorIndexThreeSegmentsRegression) {
|
||||
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 300);
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 1, 300, id_map, delete_store, version_manager,
|
||||
write_options, 300, 300);
|
||||
auto seg3 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 2, 600, id_map, delete_store, version_manager,
|
||||
write_options, 600, 300);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg3 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
ASSERT_TRUE(seg3->flush().ok());
|
||||
|
||||
auto [compact_task, output_segment] = RunCompactAndOpen(
|
||||
schema, {seg1, seg2, seg3}, 3, nullptr, version_manager);
|
||||
|
||||
ASSERT_NE(output_segment, nullptr);
|
||||
ASSERT_EQ(output_segment->doc_count(), 900);
|
||||
ASSERT_NE(output_segment->Fetch(0), nullptr);
|
||||
ASSERT_NE(output_segment->Fetch(899), nullptr);
|
||||
}
|
||||
|
||||
TEST_F(SegmentHelperTest,
|
||||
CompactTask_QuantizedVectorIndexThreeSegmentsRegression) {
|
||||
auto schema = test::TestHelper::CreateSchemaWithVectorIndex(
|
||||
false, col_name,
|
||||
std::make_shared<HnswIndexParams>(MetricType::IP, 16, 20,
|
||||
QuantizeType::FP16));
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 300);
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 1, 300, id_map, delete_store, version_manager,
|
||||
write_options, 300, 300);
|
||||
auto seg3 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 2, 600, id_map, delete_store, version_manager,
|
||||
write_options, 600, 300);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg3 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
ASSERT_TRUE(seg3->flush().ok());
|
||||
ASSERT_GT(seg1->get_quant_vector_indexer("dense_fp32").size(), 0u);
|
||||
ASSERT_GT(seg2->get_quant_vector_indexer("dense_fp32").size(), 0u);
|
||||
ASSERT_GT(seg3->get_quant_vector_indexer("dense_fp32").size(), 0u);
|
||||
|
||||
auto [compact_task, output_segment] = RunCompactAndOpen(
|
||||
schema, {seg1, seg2, seg3}, 3, nullptr, version_manager);
|
||||
|
||||
ASSERT_NE(output_segment, nullptr);
|
||||
ASSERT_EQ(output_segment->doc_count(), 900);
|
||||
ASSERT_NE(output_segment->Fetch(0), nullptr);
|
||||
ASSERT_NE(output_segment->Fetch(899), nullptr);
|
||||
ASSERT_GT(output_segment->get_vector_indexer("dense_fp32").size(), 0u);
|
||||
ASSERT_GT(output_segment->get_quant_vector_indexer("dense_fp32").size(), 0u);
|
||||
}
|
||||
|
||||
struct SegmentCompactReuseParam {
|
||||
IndexParams::Ptr vector_index_params;
|
||||
IndexType expected_output_type;
|
||||
};
|
||||
|
||||
class SegmentCompactReuseTest
|
||||
: public SegmentHelperTest,
|
||||
public testing::WithParamInterface<SegmentCompactReuseParam> {
|
||||
protected:
|
||||
// Returns the indexer's underlying VectorIndexParams::type(), defaulting to
|
||||
// FLAT if the params can't be downcast (matches the freshly-inserted state).
|
||||
static IndexType IndexerType(const VectorColumnIndexer::Ptr &indexer) {
|
||||
auto params = std::dynamic_pointer_cast<VectorIndexParams>(
|
||||
indexer->field_schema().index_params());
|
||||
return params ? params->type() : IndexType::FLAT;
|
||||
}
|
||||
|
||||
static QuantizeType QuantizeTypeOf(const IndexParams::Ptr ¶ms) {
|
||||
auto vp = std::dynamic_pointer_cast<VectorIndexParams>(params);
|
||||
return vp ? vp->quantize_type() : QuantizeType::UNDEFINED;
|
||||
}
|
||||
|
||||
// Run CreateVectorIndexTask on `segment` for `column` with `index_params`,
|
||||
// then reload the segment so its in-memory indexer reflects the new index
|
||||
// (matching collection.cc's post-optimize reload path).
|
||||
void OptimizeSegmentToVectorIndex(const Segment::Ptr &segment,
|
||||
const CollectionSchema &schema,
|
||||
const std::string &column,
|
||||
const IndexParams::Ptr &index_params) {
|
||||
CreateVectorIndexTask task(segment, column, index_params, 1);
|
||||
auto segment_task = SegmentTask::CreateCreateVectorIndexTask(task);
|
||||
ASSERT_NE(segment_task, nullptr);
|
||||
ASSERT_TRUE(SegmentHelper::Execute(segment_task).ok());
|
||||
auto executed = std::get<CreateVectorIndexTask>(segment_task->task_info());
|
||||
ASSERT_NE(executed.output_segment_meta_, nullptr);
|
||||
ASSERT_TRUE(
|
||||
segment
|
||||
->reload_vector_index(schema, executed.output_segment_meta_,
|
||||
executed.output_vector_indexers_,
|
||||
executed.output_quant_vector_indexers_)
|
||||
.ok());
|
||||
}
|
||||
|
||||
struct ScoredDoc {
|
||||
uint64_t doc_id;
|
||||
float score;
|
||||
};
|
||||
|
||||
// CreateDoc seeds VECTOR_FP32 with a constant vector of value (doc_id+0.1f).
|
||||
static std::vector<float> MakeFp32QueryVector(uint64_t doc_id_value,
|
||||
uint32_t dim) {
|
||||
return std::vector<float>(dim, static_cast<float>(doc_id_value) + 0.1f);
|
||||
}
|
||||
|
||||
static std::vector<ScoredDoc> RunSearch(
|
||||
const VectorColumnIndexer::Ptr &indexer, const std::vector<float> &qvec,
|
||||
uint32_t topk, const zvec::QueryParams::Ptr &query_params) {
|
||||
vector_column_params::QueryParams qp;
|
||||
qp.topk = topk;
|
||||
qp.filter = nullptr;
|
||||
qp.fetch_vector = false;
|
||||
qp.query_params = query_params;
|
||||
vector_column_params::VectorData data{
|
||||
vector_column_params::DenseVector{qvec.data()}};
|
||||
auto results = indexer->Search(data, qp);
|
||||
EXPECT_TRUE(results.has_value());
|
||||
if (!results.has_value()) return {};
|
||||
auto vec_res = dynamic_cast<VectorIndexResults *>(results.value().get());
|
||||
EXPECT_NE(vec_res, nullptr);
|
||||
if (vec_res == nullptr) return {};
|
||||
std::vector<ScoredDoc> out;
|
||||
for (auto it = vec_res->create_iterator(); it->valid(); it->next()) {
|
||||
out.push_back({it->doc_id(), it->score()});
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// All test instantiations use IP — higher score is better.
|
||||
static std::set<uint64_t> MergeTopKIds(
|
||||
std::vector<std::vector<ScoredDoc>> per_seg, uint32_t topk) {
|
||||
std::vector<ScoredDoc> all;
|
||||
for (auto &v : per_seg)
|
||||
for (auto &d : v) all.push_back(d);
|
||||
std::sort(all.begin(), all.end(),
|
||||
[](const ScoredDoc &a, const ScoredDoc &b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
std::set<uint64_t> ids;
|
||||
for (size_t i = 0; i < all.size() && ids.size() < topk; ++i) {
|
||||
ids.insert(all[i].doc_id);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
static zvec::QueryParams::Ptr MakeIsLinearQueryParam(IndexType type) {
|
||||
switch (type) {
|
||||
case IndexType::HNSW: {
|
||||
auto p = std::make_shared<zvec::HnswQueryParams>();
|
||||
p->set_is_linear(true);
|
||||
return p;
|
||||
}
|
||||
case IndexType::IVF: {
|
||||
auto p = std::make_shared<zvec::IVFQueryParams>();
|
||||
p->set_is_linear(true);
|
||||
return p;
|
||||
}
|
||||
case IndexType::HNSW_RABITQ: {
|
||||
auto p = std::make_shared<zvec::HnswRabitqQueryParams>();
|
||||
p->set_is_linear(true);
|
||||
return p;
|
||||
}
|
||||
case IndexType::FLAT:
|
||||
default:
|
||||
return std::make_shared<zvec::FlatQueryParams>();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Mimic the normal insertion lifecycle: small segments accumulate vectors
|
||||
// in flat storage (no vector index built yet), then compaction merges them
|
||||
// into a single segment whose vector column is built per schema.
|
||||
TEST_P(SegmentCompactReuseTest, OptimizedSegmentsReuseFirstIndexer) {
|
||||
const auto ¶m = GetParam();
|
||||
auto schema = test::TestHelper::CreateSchemaWithVectorIndex(
|
||||
false, col_name, param.vector_index_params);
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
constexpr int kSegCount = 3;
|
||||
constexpr int kDocsPerSeg = 300;
|
||||
constexpr uint32_t kTopK = 10;
|
||||
constexpr uint32_t kDim = 128;
|
||||
std::vector<Segment::Ptr> segs;
|
||||
for (int i = 0; i < kSegCount; i++) {
|
||||
auto seg = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, i, i * kDocsPerSeg, id_map, delete_store,
|
||||
version_manager, write_options, i * kDocsPerSeg, kDocsPerSeg);
|
||||
ASSERT_NE(seg, nullptr);
|
||||
ASSERT_TRUE(seg->flush().ok());
|
||||
segs.push_back(seg);
|
||||
}
|
||||
|
||||
// Capture groundtruth via FlatQuery on each source segment while every
|
||||
// segment is still backed by a flat indexer (before seg[0] is optimized).
|
||||
const std::vector<uint64_t> query_doc_values{0, kDocsPerSeg,
|
||||
kSegCount * kDocsPerSeg - 1};
|
||||
std::vector<std::set<uint64_t>> groundtruth;
|
||||
groundtruth.reserve(query_doc_values.size());
|
||||
auto flat_qp = std::make_shared<zvec::FlatQueryParams>();
|
||||
for (uint64_t qv : query_doc_values) {
|
||||
auto qvec = MakeFp32QueryVector(qv, kDim);
|
||||
std::vector<std::vector<ScoredDoc>> per_seg;
|
||||
per_seg.reserve(segs.size());
|
||||
// Per-segment indexers use block-local doc ids (0..kDocsPerSeg-1).
|
||||
// The compacted output indexer reindexes them sequentially across
|
||||
// segments, so add segment offset to align id spaces before merging.
|
||||
for (size_t s = 0; s < segs.size(); ++s) {
|
||||
auto in_indexers = segs[s]->get_vector_indexer("dense_fp32");
|
||||
ASSERT_FALSE(in_indexers.empty());
|
||||
ASSERT_EQ(IndexerType(in_indexers.front()), IndexType::FLAT);
|
||||
auto local = RunSearch(in_indexers.front(), qvec, kTopK, flat_qp);
|
||||
const uint64_t offset = static_cast<uint64_t>(s) * kDocsPerSeg;
|
||||
for (auto &d : local) d.doc_id += offset;
|
||||
per_seg.push_back(std::move(local));
|
||||
}
|
||||
auto gt = MergeTopKIds(std::move(per_seg), kTopK);
|
||||
ASSERT_EQ(gt.size(), kTopK);
|
||||
groundtruth.push_back(std::move(gt));
|
||||
}
|
||||
|
||||
// Optimize seg[0]'s vector fields to the parametric index type, mimicking
|
||||
// the lifecycle the compact path exercises.
|
||||
for (const auto &vf : schema->vector_fields()) {
|
||||
OptimizeSegmentToVectorIndex(segs[0], *schema, vf->name(),
|
||||
vf->index_params());
|
||||
}
|
||||
|
||||
// For quantized index types (e.g. HNSW_RABITQ) the built index lives in
|
||||
// get_quant_vector_indexer; get_vector_indexer keeps the raw FLAT
|
||||
// indexer. See CompactTask_QuantizedVectorIndexThreeSegmentsRegression.
|
||||
const bool quantized =
|
||||
QuantizeTypeOf(param.vector_index_params) != QuantizeType::UNDEFINED;
|
||||
for (int i = 0; i < kSegCount; i++) {
|
||||
auto in_indexers = quantized
|
||||
? segs[i]->get_quant_vector_indexer("dense_fp32")
|
||||
: segs[i]->get_vector_indexer("dense_fp32");
|
||||
ASSERT_FALSE(in_indexers.empty());
|
||||
ASSERT_EQ(IndexerType(in_indexers.front()),
|
||||
i == 0 ? param.expected_output_type : IndexType::FLAT);
|
||||
}
|
||||
|
||||
auto [compact_task, output_segment] =
|
||||
RunCompactAndOpen(schema, segs, kSegCount, nullptr, version_manager);
|
||||
|
||||
ASSERT_NE(output_segment, nullptr);
|
||||
ASSERT_EQ(output_segment->doc_count(), kSegCount * kDocsPerSeg);
|
||||
ASSERT_NE(output_segment->Fetch(0), nullptr);
|
||||
ASSERT_NE(output_segment->Fetch(kSegCount * kDocsPerSeg - 1), nullptr);
|
||||
|
||||
auto out_indexers =
|
||||
quantized ? output_segment->get_quant_vector_indexer("dense_fp32")
|
||||
: output_segment->get_vector_indexer("dense_fp32");
|
||||
ASSERT_FALSE(out_indexers.empty());
|
||||
EXPECT_EQ(IndexerType(out_indexers.front()), param.expected_output_type);
|
||||
|
||||
// is_linear queries on the merged indexer must reproduce the pre-compact
|
||||
// groundtruth. Quantized indexers are allowed a small recall hit.
|
||||
auto linear_qp = MakeIsLinearQueryParam(param.expected_output_type);
|
||||
const double kMinRecall = quantized ? 0.8 : 1.0;
|
||||
for (size_t qi = 0; qi < query_doc_values.size(); ++qi) {
|
||||
auto qvec = MakeFp32QueryVector(query_doc_values[qi], kDim);
|
||||
auto hits = RunSearch(out_indexers.front(), qvec, kTopK, linear_qp);
|
||||
ASSERT_EQ(hits.size(), kTopK);
|
||||
size_t intersect = 0;
|
||||
for (const auto &h : hits) {
|
||||
if (groundtruth[qi].count(h.doc_id)) intersect++;
|
||||
}
|
||||
double recall = static_cast<double>(intersect) / kTopK;
|
||||
EXPECT_GE(recall, kMinRecall)
|
||||
<< "query[" << qi << "] (value=" << query_doc_values[qi]
|
||||
<< ") recall=" << recall;
|
||||
}
|
||||
}
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(Hnsw, SegmentCompactReuseTest,
|
||||
testing::Values(SegmentCompactReuseParam{
|
||||
std::make_shared<HnswIndexParams>(MetricType::IP,
|
||||
16, 200),
|
||||
IndexType::HNSW}));
|
||||
|
||||
// CreateNormalSchema() only puts the test's vector_index_params on dense_fp32.
|
||||
// The other 4 vector fields are hardcoded — dense_fp16/dense_int8/sparse_fp16
|
||||
// are always FlatIndexParams, and sparse_fp32 gets the
|
||||
// cloned params only if supports_sparse is true (utils.cc:117-124), which
|
||||
// excludes IVF and HNSW_RABITQ — so for IVF it also falls back to FLAT.
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
Ivf, SegmentCompactReuseTest,
|
||||
testing::Values(SegmentCompactReuseParam{
|
||||
std::make_shared<IVFIndexParams>(MetricType::IP, 10, 4, false,
|
||||
QuantizeType::UNDEFINED),
|
||||
IndexType::IVF}));
|
||||
|
||||
// TODO: re-enable when rabitQ merge fixed
|
||||
// #if RABITQ_SUPPORTED
|
||||
// INSTANTIATE_TEST_SUITE_P(HnswRabitq, SegmentCompactReuseTest,
|
||||
// testing::Values(SegmentCompactReuseParam{
|
||||
// std::make_shared<HnswRabitqIndexParams>(
|
||||
// MetricType::IP, 7, 256, 16, 200, 0),
|
||||
// IndexType::HNSW_RABITQ}));
|
||||
// #endif
|
||||
|
||||
TEST_F(SegmentHelperTest, CompactTask_FilterMultiSegmentsRegression) {
|
||||
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
|
||||
auto version_manager = CreateVersionManager(*schema);
|
||||
auto write_options = WriteOptions();
|
||||
|
||||
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
|
||||
write_options, 0, 400);
|
||||
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
|
||||
col_path, *schema, 1, 400, id_map, delete_store, version_manager,
|
||||
write_options, 400, 400);
|
||||
ASSERT_TRUE(seg1 != nullptr);
|
||||
ASSERT_TRUE(seg2 != nullptr);
|
||||
ASSERT_TRUE(seg1->flush().ok());
|
||||
ASSERT_TRUE(seg2->flush().ok());
|
||||
|
||||
auto filter = std::make_shared<EasyIndexFilter>(
|
||||
[](uint64_t id) -> bool { return id < 100 || (id >= 400 && id < 450); });
|
||||
|
||||
auto [compact_task, output_segment] =
|
||||
RunCompactAndOpen(schema, {seg1, seg2}, 2, filter, version_manager);
|
||||
|
||||
ASSERT_NE(output_segment, nullptr);
|
||||
ASSERT_EQ(output_segment->doc_count(), 650);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue