From 3521a0633a35019ebd9ad37c03f8e122d0a44b3f Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Fri, 12 Jun 2026 11:10:26 +0800 Subject: [PATCH] fix: preserve element indices schema for empty element-level search results (#50417) issue: https://github.com/milvus-io/milvus/issues/50010 ref: https://github.com/milvus-io/milvus/issues/42148 Signed-off-by: SpadeA --- .../src/exec/operator/VectorSearchNode.cpp | 5 ++- .../core/src/query/ExecPlanNodeVisitor.cpp | 8 ++-- internal/core/src/query/SearchOnGrowing.cpp | 2 + internal/core/src/query/SearchOnSealed.cpp | 3 ++ internal/core/src/query/Utils.h | 1 + .../src/segcore/search_result_export_c.cpp | 3 +- .../segcore/search_result_export_c_test.cpp | 43 +++++++++++++++++++ 7 files changed, 58 insertions(+), 7 deletions(-) diff --git a/internal/core/src/exec/operator/VectorSearchNode.cpp b/internal/core/src/exec/operator/VectorSearchNode.cpp index a51b936edc..e87d15a67d 100644 --- a/internal/core/src/exec/operator/VectorSearchNode.cpp +++ b/internal/core/src/exec/operator/VectorSearchNode.cpp @@ -44,11 +44,12 @@ namespace milvus { namespace exec { static milvus::SearchResult -empty_search_result(int64_t num_queries) { +empty_search_result(int64_t num_queries, bool element_level = false) { milvus::SearchResult final_result; final_result.total_nq_ = num_queries; final_result.unity_topK_ = 0; // no result final_result.total_data_cnt_ = 0; + final_result.element_level_ = element_level; return final_result; } @@ -146,7 +147,7 @@ PhyVectorSearchNode::GetOutput() { query_context_->set_active_element_count(element_bitset.size()); if (element_bitset.empty()) { query_context_->set_search_result( - empty_search_result(num_queries)); + empty_search_result(num_queries, ph.element_level_)); return input_; } diff --git a/internal/core/src/query/ExecPlanNodeVisitor.cpp b/internal/core/src/query/ExecPlanNodeVisitor.cpp index b45a050973..cf449d76de 100644 --- a/internal/core/src/query/ExecPlanNodeVisitor.cpp +++ b/internal/core/src/query/ExecPlanNodeVisitor.cpp @@ -35,11 +35,12 @@ namespace milvus::query { static SearchResult -empty_search_result(int64_t num_queries) { +empty_search_result(int64_t num_queries, bool element_level = false) { SearchResult final_result; final_result.total_nq_ = num_queries; final_result.unity_topK_ = 0; // no result final_result.total_data_cnt_ = 0; + final_result.element_level_ = element_level; return final_result; } @@ -476,8 +477,9 @@ ExecPlanNodeVisitor::visit(VectorPlanNode& node) { // PreExecute: skip all calculation if (active_count == 0) { - search_result_opt_ = - empty_search_result(placeholder_group_->at(0).num_of_queries_); + const auto& placeholder = placeholder_group_->at(0); + search_result_opt_ = empty_search_result(placeholder.num_of_queries_, + placeholder.element_level_); return; } diff --git a/internal/core/src/query/SearchOnGrowing.cpp b/internal/core/src/query/SearchOnGrowing.cpp index f1d9d85f10..71ee14b143 100644 --- a/internal/core/src/query/SearchOnGrowing.cpp +++ b/internal/core/src/query/SearchOnGrowing.cpp @@ -185,6 +185,7 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, const bool is_element_level_search = data_type == DataType::VECTOR_ARRAY && info.array_offsets_ != nullptr; + search_result.element_level_ = is_element_level_search; const auto has_offset_mapping = offset_mapping.IsEnabled() && !is_element_level_search; @@ -370,6 +371,7 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, info.array_offsets_.get()); search_result.seg_offsets_ = std::move(seg_offsets); search_result.element_indices_ = std::move(elem_indicies); + search_result.element_level_ = true; } else { if (has_offset_mapping) { offset_mapping.TransformOffsets(final_qr.mutable_offsets()); diff --git a/internal/core/src/query/SearchOnSealed.cpp b/internal/core/src/query/SearchOnSealed.cpp index d24909896e..2d606aaabd 100644 --- a/internal/core/src/query/SearchOnSealed.cpp +++ b/internal/core/src/query/SearchOnSealed.cpp @@ -98,6 +98,7 @@ SearchOnSealedIndex(const Schema& schema, const auto& offset_mapping = vec_index->GetOffsetMapping(); const bool is_element_level_search = search_info.array_offsets_ != nullptr; + search_result.element_level_ = is_element_level_search; TargetBitmap transformed_bitset; BitsetView search_bitset = bitset; const auto has_offset_mapping = @@ -201,6 +202,7 @@ SearchOnSealedColumn(const Schema& schema, bool is_element_level_search = field.get_data_type() == DataType::VECTOR_ARRAY && search_info.array_offsets_ != nullptr; + result.element_level_ = is_element_level_search; TargetBitmap transformed_bitset; BitsetView search_bitview = bitview; const auto has_offset_mapping = @@ -339,6 +341,7 @@ SearchOnSealedColumn(const Schema& schema, search_info.array_offsets_.get()); result.seg_offsets_ = std::move(seg_offsets); result.element_indices_ = std::move(elem_indicies); + result.element_level_ = true; } else { if (has_offset_mapping) { offset_mapping.TransformOffsets(final_qr.mutable_offsets()); diff --git a/internal/core/src/query/Utils.h b/internal/core/src/query/Utils.h index b6f40b310b..37c4453f56 100644 --- a/internal/core/src/query/Utils.h +++ b/internal/core/src/query/Utils.h @@ -85,6 +85,7 @@ FinalizeVectorSearchOffsets(SearchResult& result, ApplyElementIDMapping(result.seg_offsets_, *array_offsets); result.seg_offsets_ = std::move(doc_offsets); result.element_indices_ = std::move(elem_indices); + result.element_level_ = true; } else { if (offset_mapping.IsEnabled()) { offset_mapping.TransformOffsets(result.seg_offsets_); diff --git a/internal/core/src/segcore/search_result_export_c.cpp b/internal/core/src/segcore/search_result_export_c.cpp index d0681b85e2..16d049dceb 100644 --- a/internal/core/src/segcore/search_result_export_c.cpp +++ b/internal/core/src/segcore/search_result_export_c.cpp @@ -605,8 +605,7 @@ BuildSearchResultBatch( // Build $element_indices column when element-level search is active. // element_indices_ is int32 and aligned with seg_offsets_ after compaction. - if (search_result->element_level_ && - !search_result->element_indices_.empty()) { + if (search_result->element_level_) { AssertInfo( search_result->element_indices_.size() == total_valid, "element_indices_ size {} does not match seg_offsets_ size {}", diff --git a/internal/core/src/segcore/search_result_export_c_test.cpp b/internal/core/src/segcore/search_result_export_c_test.cpp index 985fb707f7..08d0580669 100644 --- a/internal/core/src/segcore/search_result_export_c_test.cpp +++ b/internal/core/src/segcore/search_result_export_c_test.cpp @@ -351,6 +351,49 @@ TEST(SearchResultExport, EXPECT_EQ((*batch_result)->num_columns(), 3); } +TEST( + SearchResultExport, + ExportSearchResultAsArrowRecordBatch_EmptyElementLevelResultHasElementIndices) { + auto schema = std::make_shared(); + auto pk_fid = schema->AddDebugField("pk", DataType::INT64); + schema->set_primary_field_id(pk_fid); + Plan plan(schema); + plan.plan_node_ = std::make_unique(); + + SearchResult sr; + sr.total_nq_ = 1; + sr.unity_topK_ = 0; + sr.element_level_ = true; + + ArrowSchema stream_schema{}; + ArrowArray stream_array{}; + int64_t* stream_chunk_sizes = nullptr; + int64_t stream_num_chunks = 0; + auto status = ExportSearchResultAsArrowRecordBatch( + reinterpret_cast(&sr), + reinterpret_cast(&plan), + nullptr, + 0, + &stream_schema, + &stream_array, + &stream_chunk_sizes, + &stream_num_chunks, + nullptr); + [[maybe_unused]] auto stream_chunk_sizes_guard = + AdoptChunkSizes(stream_chunk_sizes); + ASSERT_EQ(status.error_code, 0) << status.error_msg; + + auto batch_result = + ImportExportedRecordBatch(&stream_array, &stream_schema); + ASSERT_TRUE(batch_result.ok()) << batch_result.status().ToString(); + ASSERT_NE(*batch_result, nullptr); + EXPECT_EQ((*batch_result)->num_rows(), 0); + ASSERT_EQ((*batch_result)->num_columns(), 4); + EXPECT_EQ((*batch_result)->schema()->field(3)->name(), "$element_indices"); + EXPECT_TRUE( + (*batch_result)->schema()->field(3)->type()->Equals(arrow::int32())); +} + TEST(SearchResultExport, ExportSearchResultAsArrowRecordBatch_MultiFieldGroupByColumns) { auto schema = std::make_shared();