cache vector values iterators in VectorFieldSources

msokolov · Sep 16, 2024 · abc1713 · abc1713
1 parent 4feecf8
commit abc1713
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/...ies/src/java/org/apache/lucene/queries/function/valuesource/ByteKnnVectorFieldSource.java b/...ies/src/java/org/apache/lucene/queries/function/valuesource/ByteKnnVectorFieldSource.java
@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.Objects;
 import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.VectorEncoding;
@@ -63,19 +64,20 @@ protected DocIdSetIterator getVectorIterator() {
     }
 
     return new VectorFieldFunction(this) {
+      KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
 
       @Override
       public byte[] byteVectorVal(int doc) throws IOException {
         if (exists(doc)) {
-          return vectorValues.vectorValue(vectorValues.iterator().index());
+          return vectorValues.vectorValue(iterator.index());
         } else {
           return null;
         }
       }
 
       @Override
       protected DocIdSetIterator getVectorIterator() {
-        return vectorValues.iterator();
+        return iterator;
       }
     };
   }

diff --git a/...es/src/java/org/apache/lucene/queries/function/valuesource/FloatKnnVectorFieldSource.java b/...es/src/java/org/apache/lucene/queries/function/valuesource/FloatKnnVectorFieldSource.java
@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.Objects;
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.VectorEncoding;
@@ -62,19 +63,20 @@ protected DocIdSetIterator getVectorIterator() {
     }
 
     return new VectorFieldFunction(this) {
+      KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
 
       @Override
       public float[] floatVectorVal(int doc) throws IOException {
         if (exists(doc)) {
-          return vectorValues.vectorValue(vectorValues.iterator().index());
+          return vectorValues.vectorValue(iterator.index());
         } else {
           return null;
         }
       }
 
       @Override
       protected DocIdSetIterator getVectorIterator() {
-        return vectorValues.iterator();
+        return iterator;
       }
     };
   }