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 <tangchenjie1210@gmail.com>

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_;
             }
 

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;
     }
 

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());

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());

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_);

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 {}",

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<Schema>();
+    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
+    schema->set_primary_field_id(pk_fid);
+    Plan plan(schema);
+    plan.plan_node_ = std::make_unique<VectorPlanNode>();
+
+    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<CSearchResult>(&sr),
+        reinterpret_cast<CSearchPlan>(&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<Schema>();