Fix MergedByteVectorValues internal ordinal tracking (#15553)

Move internal ordinal tracking in `MergedByteVectorValues` from `vectorValue` -> `nextDoc` to allow loading only a subset of vectors during iteration.

Author: finnroblin

lucene/CHANGES.txt

@@ -232,6 +232,9 @@ Bug Fixes
 * GITHUB#15554: Fix tessellator failure by preferring the shared vertex that is the leftmost vertex of the hole
   (Ignacio Vera)
 
+* GITHUB#15553: Fix MergedByteVectorValues lastOrd behavior to enable partitioning of vector values
+(Finn Roblin, Dooyong Kim)
+
 Other
 ---------------------
 * GITHUB#15237: Fix SmartChinese to only deserialize dictionary data from classpath with a native array

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsWriter.java

@@ -121,8 +121,12 @@ private void finishMerge(MergeState mergeState) throws IOException {
     }
   }
 
-  /** Tracks state of one sub-reader that we are merging */
-  private static class FloatVectorValuesSub extends DocIDMerger.Sub {
+  /**
+   * Tracks state of one sub-reader of float vectors that we are merging.
+   *
+   * @lucene.internal
+   */
+  static class FloatVectorValuesSub extends DocIDMerger.Sub {
 
     final FloatVectorValues values;
     final KnnVectorValues.DocIndexIterator iterator;
@@ -144,7 +148,12 @@ public int index() {
     }
   }
 
-  private static class ByteVectorValuesSub extends DocIDMerger.Sub {
+  /**
+   * Tracks state of one sub-reader of byte vectors that we are merging.
+   *
+   * @lucene.internal
+   */
+  static class ByteVectorValuesSub extends DocIDMerger.Sub {
 
     final ByteVectorValues values;
     final KnnVectorValues.DocIndexIterator iterator;
@@ -303,6 +312,11 @@ public static ByteVectorValues mergeByteVectorValues(FieldInfo fieldInfo, MergeS
           mergeState);
     }
 
+    /**
+     * Unified view over several segments containing float vector values.
+     *
+     * @lucene.internal
+     */
     static class MergedFloat32VectorValues extends FloatVectorValues {
       private final List<FloatVectorValuesSub> subs;
       private final DocIDMerger<FloatVectorValuesSub> docIdMerger;
@@ -311,7 +325,8 @@ static class MergedFloat32VectorValues extends FloatVectorValues {
       private int lastOrd = -1;
       FloatVectorValuesSub current;
 
-      private MergedFloat32VectorValues(List<FloatVectorValuesSub> subs, MergeState mergeState)
+      // package-private for testing
+      MergedFloat32VectorValues(List<FloatVectorValuesSub> subs, MergeState mergeState)
           throws IOException {
         this.subs = subs;
         docIdMerger = DocIDMerger.of(subs, mergeState.needsIndexSort);
@@ -401,6 +416,11 @@ public FloatVectorValues copy() {
       }
     }
 
+    /**
+     * Unified view over several segments containing byte vector values.
+     *
+     * @lucene.internal
+     */
     static class MergedByteVectorValues extends ByteVectorValues {
       private final List<ByteVectorValuesSub> subs;
       private final DocIDMerger<ByteVectorValuesSub> docIdMerger;
@@ -410,7 +430,8 @@ static class MergedByteVectorValues extends ByteVectorValues {
       private int docId = -1;
       ByteVectorValuesSub current;
 
-      private MergedByteVectorValues(List<ByteVectorValuesSub> subs, MergeState mergeState)
+      // package-private for testing
+      MergedByteVectorValues(List<ByteVectorValuesSub> subs, MergeState mergeState)
           throws IOException {
         this.subs = subs;
         docIdMerger = DocIDMerger.of(subs, mergeState.needsIndexSort);
@@ -423,11 +444,9 @@ private MergedByteVectorValues(List<ByteVectorValuesSub> subs, MergeState mergeS
 
       @Override
       public byte[] vectorValue(int ord) throws IOException {
-        if (ord != lastOrd + 1) {
+        if (ord != lastOrd) {
           throw new IllegalStateException(
               "only supports forward iteration: ord=" + ord + ", lastOrd=" + lastOrd);
-        } else {
-          lastOrd = ord;
         }
         return current.values.vectorValue(current.index());
       }
@@ -455,6 +474,7 @@ public int nextDoc() throws IOException {
               index = NO_MORE_DOCS;
             } else {
               docId = current.mappedDocID;
+              ++lastOrd;
               ++index;
             }
             return docId;

lucene/core/src/test/org/apache/lucene/codecs/TestMergedVectorValues.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.codecs;
+
+import java.io.IOException;
+import java.util.List;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.MergeState;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+/** Tests for merged vector values to ensure lastOrd is properly incremented during iteration. */
+public class TestMergedVectorValues extends LuceneTestCase {
+
+  /**
+   * Test that skipping vectors in MergedByteVectorValues via nextDoc() and then loading a
+   * subsequent vector via vectorValue() works correctly.
+   */
+  public void testSkipsInMergedByteVectorValues() throws IOException {
+    // Data
+    List<byte[]> vectors = List.of(new byte[] {0}, new byte[] {1});
+
+    // Setup
+    KnnVectorsWriter.ByteVectorValuesSub sub =
+        new KnnVectorsWriter.ByteVectorValuesSub(x -> x, ByteVectorValues.fromBytes(vectors, 1));
+    MergeState state =
+        new MergeState(
+            null, null, null, null, null, null, null, null, null, null, null, null, null, null,
+            null, false);
+
+    // Run the test
+    ByteVectorValues values =
+        new KnnVectorsWriter.MergedVectorValues.MergedByteVectorValues(List.of(sub), state);
+
+    // Skip doc 0 and load doc 1
+    values.iterator().nextDoc(); // doc 0
+    values.iterator().nextDoc(); // doc 1
+
+    // Read vector for doc 1
+    assertArrayEquals(vectors.get(1), values.vectorValue(1));
+  }
+
+  /**
+   * Test that skipping vectors in MergedFloat32VectorValues via nextDoc() and then loading a
+   * subsequent vector via vectorValue() works correctly.
+   */
+  public void testSkipsInMergedFloat32VectorValues() throws IOException {
+    // Data
+    List<float[]> vectors = List.of(new float[] {0.0f}, new float[] {1.0f});
+
+    // Setup
+    KnnVectorsWriter.FloatVectorValuesSub sub =
+        new KnnVectorsWriter.FloatVectorValuesSub(x -> x, FloatVectorValues.fromFloats(vectors, 1));
+    MergeState state =
+        new MergeState(
+            null, null, null, null, null, null, null, null, null, null, null, null, null, null,
+            null, false);
+
+    // Run the test
+    FloatVectorValues values =
+        new KnnVectorsWriter.MergedVectorValues.MergedFloat32VectorValues(List.of(sub), state);
+
+    // Skip doc 0 and load doc 1
+    values.iterator().nextDoc(); // doc 0
+    values.iterator().nextDoc(); // doc 1
+
+    // Read vector for doc 1
+    assertArrayEquals(vectors.get(1), values.vectorValue(1), 0.0f);
+  }
+}