fix: break out of MoveCursorForSingleChunk once the batch is full

Project: milvus-io/milvus | Language: C++ | Concepts: error-handling, testing-strategy, data-structures | Difficulty: mid
Link: https://github.com/milvus-io/milvus/pull/52262

The problem

The function MoveCursorForSingleChunk was supposed to advance a cursor through a batch of rows in a segment, stopping when it had processed the requested batch size. However, it did not exit the loop when the batch was full. Instead, it kept iterating to the end of the segment, leaving the cursor at the segment's end rather than at the intended batch boundary. This caused both correctness issues (future reads would start at the wrong place) and performance problems (unnecessary looping over potentially millions of rows).

Why it's non-trivial

A naive fix might just check the batch size, but if you forget to actually break or return from the loop, the condition will be true for every subsequent iteration, and the loop will continue anyway. This is easy to overlook, especially since sibling functions already have the correct logic, and the bug only appears under specific conditions (single-chunk, non-chunked/growing segments).

The approach

The fix is a one-line change: add a return; statement to exit the function as soon as the batch is full.

if (++processed_rows >= batch_size) {
    current_chunk_id = chunk_id;
    current_chunk_pos = i + 1;
+   return;
}

What this does:

• As soon as processed_rows reaches batch_size, it updates the cursor position and immediately returns from the function, preventing further iteration.
• This matches the intended behavior and the implementation in similar functions.

The PR also adds two new unit tests to verify the fix and prevent regressions:

1. SegmentChunkReader.MoveCursorForSingleChunkStopsAtBatchBoundary
   This test creates a segment with more than one chunk, sets up a reader, and advances the cursor by exactly one batch. It asserts:

   • The cursor advances by exactly batch_size.
   • The cursor does not land at the segment end.
   • A second call continues from where the first left off.

2. SegmentChunkReader.GrowingConjunctSkipKeepsNextBatchCorrect
   This test simulates a real-world expression evaluation scenario where the bug would cause incorrect results. It checks that after skipping a batch, the next batch starts at the correct offset, and all rows are processed as expected.

From the review thread

No notable review discussion on this PR.

Takeaways

• Always exit loops explicitly when a condition is met, especially in cursor or batch-processing code. Failing to do so can cause subtle correctness and performance bugs.
• Write regression tests that reproduce the exact failure mode, including both direct and end-to-end scenarios. This ensures the bug stays fixed and documents the intended behavior.
• Compare similar code paths (e.g., sibling functions) for consistency. If one path handles a case correctly, but another doesn't, it's a clue that something is missing.