diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.cc b/src/core/mixed_reducer/mixed_streamer_reducer.cc index cf20a1b..6d8b31c 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.cc +++ b/src/core/mixed_reducer/mixed_streamer_reducer.cc @@ -146,7 +146,19 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { std::vector 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++) { diff --git a/src/db/index/column/vector_column/vector_column_indexer.h b/src/db/index/column/vector_column/vector_column_indexer.h index f3b2286..3c7ac7c 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.h +++ b/src/db/index/column/vector_column/vector_column_indexer.h @@ -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; diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index ff5204a..7c28832 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #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 row_id_filter = - std::make_shared(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 row_id_filter; + if (!delete_row_id_bitmap.isEmpty()) { + row_id_filter = std::make_shared(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(field->index_params()); + using FetchIndexersFn = + std::vector (Segment::*)(const std::string &) + const; + auto collect_merge_indexers = [&](FetchIndexersFn fetch) { + std::vector 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(vector_index_path, *field); - s = vector_indexer->Open({true, true}); - CHECK_RETURN_STATUS(s); - - std::vector 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( - vector_index_path, *field_without_quantize); - s = vector_indexer->Open({true, true}); - CHECK_RETURN_STATUS(s); - - std::vector 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( - 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 &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(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 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(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, diff --git a/src/db/index/segment/segment_helper.h b/src/db/index/segment/segment_helper.h index 96b8ee8..3de46eb 100644 --- a/src/db/index/segment/segment_helper.h +++ b/src/db/index/segment/segment_helper.h @@ -217,6 +217,18 @@ class SegmentHelper { uint64_t max_doc_id, uint32_t doc_count, int concurrency, std::vector *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 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 diff --git a/tests/db/index/segment/segment_helper_test.cc b/tests/db/index/segment/segment_helper_test.cc index ca9902f..d290f82 100644 --- a/tests/db/index/segment/segment_helper_test.cc +++ b/tests/db/index/segment/segment_helper_test.cc @@ -13,10 +13,12 @@ // limitations under the License. #include "db/index/segment/segment_helper.h" +#include #include #include #include #include +#include #include #include #include @@ -28,6 +30,9 @@ #include #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 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(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 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(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 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(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 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(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 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(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( - [&](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(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( - [&](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(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()); -} \ No newline at end of file +} + +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(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 { + 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( + indexer->field_schema().index_params()); + return params ? params->type() : IndexType::FLAT; + } + + static QuantizeType QuantizeTypeOf(const IndexParams::Ptr ¶ms) { + auto vp = std::dynamic_pointer_cast(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(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 MakeFp32QueryVector(uint64_t doc_id_value, + uint32_t dim) { + return std::vector(dim, static_cast(doc_id_value) + 0.1f); + } + + static std::vector RunSearch( + const VectorColumnIndexer::Ptr &indexer, const std::vector &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(results.value().get()); + EXPECT_NE(vec_res, nullptr); + if (vec_res == nullptr) return {}; + std::vector 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 MergeTopKIds( + std::vector> per_seg, uint32_t topk) { + std::vector 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 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(); + p->set_is_linear(true); + return p; + } + case IndexType::IVF: { + auto p = std::make_shared(); + p->set_is_linear(true); + return p; + } + case IndexType::HNSW_RABITQ: { + auto p = std::make_shared(); + p->set_is_linear(true); + return p; + } + case IndexType::FLAT: + default: + return std::make_shared(); + } + } +}; + +// 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 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 query_doc_values{0, kDocsPerSeg, + kSegCount * kDocsPerSeg - 1}; + std::vector> groundtruth; + groundtruth.reserve(query_doc_values.size()); + auto flat_qp = std::make_shared(); + for (uint64_t qv : query_doc_values) { + auto qvec = MakeFp32QueryVector(qv, kDim); + std::vector> 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(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(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(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(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( +// 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( + [](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); +}