SlowCompositeCodecReaderWrapper must copy its sub-vector values to ma…

…intain thread-safety (#14092)
apache · Dec 31, 2024 · 68051f1 · 68051f1
1 parent 525b963
commit 68051f1
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
@@ -301,7 +301,12 @@ public void checkIntegrity() throws IOException {
     }
   }
 
-  private record DocValuesSub<T extends KnnVectorValues>(T sub, int docStart, int ordStart) {}
+  private record DocValuesSub<T extends KnnVectorValues>(T sub, int docStart, int ordStart) {
+    @SuppressWarnings("unchecked")
+    DocValuesSub<T> copy() throws IOException {
+      return new DocValuesSub<T>((T) (sub.copy()), docStart, ordStart);
+    }
+  }
 
   private static class MergedDocIterator<T extends KnnVectorValues>
       extends KnnVectorValues.DocIndexIterator {
@@ -850,15 +855,15 @@ public FloatVectorValues getFloatVectorValues(String field) throws IOException {
     class MergedFloatVectorValues extends FloatVectorValues {
       final int dimension;
       final int size;
-      final DocValuesSub<?>[] subs;
+      final List<DocValuesSub<FloatVectorValues>> subs;
       final MergedDocIterator<FloatVectorValues> iter;
       final int[] starts;
       int lastSubIndex;
 
       MergedFloatVectorValues(int dimension, int size, List<DocValuesSub<FloatVectorValues>> subs) {
         this.dimension = dimension;
         this.size = size;
-        this.subs = subs.toArray(new DocValuesSub<?>[0]);
+        this.subs = subs;
         iter = new MergedDocIterator<>(subs);
         // [0, start(1), ..., size] - we want the extra element
         // to avoid checking for out-of-array bounds
@@ -888,8 +893,8 @@ public int size() {
       @Override
       public FloatVectorValues copy() throws IOException {
         List<DocValuesSub<FloatVectorValues>> subsCopy = new ArrayList<>();
-        for (Object sub : subs) {
-          subsCopy.add((DocValuesSub<FloatVectorValues>) sub);
+        for (DocValuesSub<FloatVectorValues> sub : subs) {
+          subsCopy.add(sub.copy());
         }
         return new MergedFloatVectorValues(dimension, size, subsCopy);
       }
@@ -900,9 +905,9 @@ public float[] vectorValue(int ord) throws IOException {
         // We need to implement fully random-access API here in order to support callers like
         // SortingCodecReader that rely on it.
         lastSubIndex = findSub(ord, lastSubIndex, starts);
-        assert subs[lastSubIndex].sub != null;
-        return ((FloatVectorValues) subs[lastSubIndex].sub)
-            .vectorValue(ord - subs[lastSubIndex].ordStart);
+        DocValuesSub<FloatVectorValues> sub = subs.get(lastSubIndex);
+        assert sub.sub != null;
+        return (sub.sub).vectorValue(ord - sub.ordStart);
       }
     }
 
@@ -929,15 +934,15 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
     class MergedByteVectorValues extends ByteVectorValues {
       final int dimension;
       final int size;
-      final DocValuesSub<?>[] subs;
+      final List<DocValuesSub<ByteVectorValues>> subs;
       final MergedDocIterator<ByteVectorValues> iter;
       final int[] starts;
       int lastSubIndex;
 
       MergedByteVectorValues(int dimension, int size, List<DocValuesSub<ByteVectorValues>> subs) {
         this.dimension = dimension;
         this.size = size;
-        this.subs = subs.toArray(new DocValuesSub<?>[0]);
+        this.subs = subs;
         iter = new MergedDocIterator<>(subs);
         // [0, start(1), ..., size] - we want the extra element
         // to avoid checking for out-of-array bounds
@@ -970,16 +975,16 @@ public byte[] vectorValue(int ord) throws IOException {
         // SortingCodecReader that rely on it.  We maintain lastSubIndex since we expect some
         // repetition.
         lastSubIndex = findSub(ord, lastSubIndex, starts);
-        return ((ByteVectorValues) subs[lastSubIndex].sub)
-            .vectorValue(ord - subs[lastSubIndex].ordStart);
+        DocValuesSub<ByteVectorValues> sub = subs.get(lastSubIndex);
+        return sub.sub.vectorValue(ord - sub.ordStart);
       }
 
       @SuppressWarnings("unchecked")
       @Override
       public ByteVectorValues copy() throws IOException {
         List<DocValuesSub<ByteVectorValues>> newSubs = new ArrayList<>();
-        for (Object sub : subs) {
-          newSubs.add((DocValuesSub<ByteVectorValues>) sub);
+        for (DocValuesSub<ByteVectorValues> sub : subs) {
+          newSubs.add(sub.copy());
         }
         return new MergedByteVectorValues(dimension, size, newSubs);
       }