fix: support vector output with order_by

Project: milvus-io/milvus | Language: C++ | Concepts: query-optimization, data-structures, algorithmic-optimization | Difficulty: mid
Link: https://github.com/milvus-io/milvus/pull/52498

The problem

When executing an ORDER BY query that requested vector fields in the output, Milvus would fail. The query planner's ProjectNode tried to materialize vector columns before the TopK operation, but it only supports scalar columns. As a result, valid queries that requested vector outputs with ORDER BY would hit an internal error during projection.

Why it's non-trivial

A naive fix might try to treat all output columns the same, but vector columns require special handling: they cannot be materialized early in the pipeline, and must be fetched after sorting (late materialization). Additionally, the system must preserve field IDs and vector dimensions through the result pipeline to ensure correct output and ordering.

The approach

1. Identify vector fields that require late materialization

The code in PlanProto.cpp was updated to treat vector fields the same as variable-width fields (like strings), deferring their materialization until after TopK:

-    // Collect non-sort output fields and check for variable-width types.
+    // Collect non-sort output fields and check whether any must be late
+    // materialized. ProjectNode only materializes scalar columns, so vector
+    // outputs must be deferred just like variable-width outputs.

This comment clarifies that vector outputs, like variable-width outputs, must be deferred.

The code then checks for both variable-width and vector types:

-            if (IsVariableDataType(schema->GetFieldType(fid))) {
-                has_variable_width = true;
+            auto field_type = schema->GetFieldType(fid);
+            if (IsVariableDataType(field_type) ||
+                IsVectorDataType(field_type)) {
+                requires_late_materialization = true;
             }

Here, IsVectorDataType(field_type) is used to detect vector fields, setting a flag to defer them.

2. Defer all non-sort output fields if any require late materialization

If any non-sort output field is variable-width or a vector, all such fields are deferred:

-    std::vector<FieldId> deferred_field_ids;
-    if (has_variable_width) {
-        // Two-project mode: defer ALL non-sort output fields until after TopK.
-        deferred_field_ids = non_sort_output_fields;
-    } else {
+    std::vector<FieldId> deferred_field_ids;
+    if (requires_late_materialization) {
+        // Two-project mode: defer ALL non-sort output fields until after TopK.
+        deferred_field_ids = non_sort_output_fields;
+    } else {

This ensures the pipeline only materializes fields it can handle, and defers the rest for later.

3. Add a regression test to verify vector output with order_by

A new test, OrderByWithDenseVectorOutput, was added to verify that vector output fields are handled correctly:

+TEST_P(QueryOrderByTest, OrderByWithDenseVectorOutput) {
+    auto pk_fid = field_map_[string_field];
+    auto double_fid = field_map_[double_field];
+    auto vector_fid = field_map_[vector_field];
+    ...
+    proto::plan::PlanNode plan_node;
+    plan_node.add_output_field_ids(pk_fid.get());
+    plan_node.add_output_field_ids(vector_fid.get());
+    ...
+    auto parser = milvus::query::ProtoParser(schema_);
+    auto plan = parser.CreateRetrievePlan(plan_node);
+    EXPECT_EQ(plan->plan_node_->deferred_field_ids_,
+              std::vector<FieldId>{vector_fid});
+    EXPECT_EQ(plan->plan_node_->pipeline_field_ids_,
+              (std::vector<FieldId>{pk_fid, double_fid, SegmentOffsetFieldID}));
+    ...
+    auto results = segment_->Retrieve(
+        nullptr, plan.get(), MAX_TIMESTAMP, DEFAULT_MAX_OUTPUT_SIZE, false);
+    ...
+    EXPECT_EQ(vector_data.field_id(), vector_fid.get());
+    EXPECT_EQ(vector_data.vectors().dim(), 16);
+    ASSERT_EQ(vector_data.vectors().float_vector().data_size(), row_count * 16);
+    ...
+    // Verify each deferred vector row carries the payload of the ORIGINAL row
+    // selected by the sort: map returned PK -> original segment offset, then
+    // compare the vector payload element-wise against the generated data.
+    ...
+    EXPECT_TRUE(permuted);
+}

This test checks that:

• The vector field is deferred.
• The output includes the correct vector data with the right dimensions.
• The returned vectors match the original data, even after sorting.

From the review thread

No notable review discussion on this PR.

Takeaways

• Late materialization is necessary for output fields that cannot be handled early in the query pipeline (e.g., vectors, variable-width types).
• Query planners must track and preserve metadata (like field IDs and dimensions) for all output fields, especially when deferring their retrieval.
• Regression tests for query planners should verify not only that queries succeed, but that output fields are correctly matched, permuted, and validated against the original data.