You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by th...@apache.org on 2026/09/30 12:35:43 UTC
(lucene) branch branch_10x updated: Stop reading merged float vectors back to build the HNSW merge scorer (#16696) (#16750)
This is an automated email from the ASF dual-hosted git repository.
thecoop pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new adfd05f831b Stop reading merged float vectors back to build the HNSW merge scorer (#16696) (#16750)
adfd05f831b is described below
commit adfd05f831b2114302fa4fcf47fa0f728f1b1925
Author: Simon Cooper <si...@elastic.co>
AuthorDate: Wed Sep 30 13:35:38 2026 +0100
Stop reading merged float vectors back to build the HNSW merge scorer (#16696) (#16750)
Produce the merge scorer's query-side records while writing the merged vectors,
instead of reading the merged floats back to re-quantize them from scratch.
Add a method to FlatVectorsWriter that allows the HNSW writer to access the
query records directly and use them during merging.
---
lucene/CHANGES.txt | 9 +
.../lucene/codecs/hnsw/FlatVectorsWriter.java | 58 ++
.../Lucene104ScalarQuantizedVectorsFormat.java | 28 +
.../Lucene104ScalarQuantizedVectorsReader.java | 77 +-
.../Lucene104ScalarQuantizedVectorsWriter.java | 293 ++++++--
.../codecs/lucene99/Lucene99HnswVectorsWriter.java | 68 +-
.../TestLucene104ScalarQuantizedMergeScorer.java | 797 +++++++++++++++++++++
.../apache/lucene/index/TestHnswMergeAbort.java | 44 ++
.../quantization/TestOptimizedScalarQuantizer.java | 69 ++
9 files changed, 1345 insertions(+), 98 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 795a4c2c142..8419048a817 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -8,6 +8,11 @@ http://s.apache.org/luceneversions
API Changes
---------------------
+* GITHUB#16696: FlatVectorsWriter gains mergeOneFlatVectorFieldForMergeScorer and the
+ FlatVectorsWriter.MergeScorerData handle it returns, so that a flat writer can prepare what a
+ merge scorer over the merged vectors needs while it writes them. The default prepares nothing;
+ a writer that wraps another one should forward the new method too. (John Maurice)
+
* GITHUB#14399: Add a default LeafFieldComparator#disableSkipping method so document skipping can be
disabled per segment. (Serhiy Bzhezytskyy)
@@ -190,6 +195,10 @@ Optimizations
* GITHUB#16520: Support two-phase filters in MaxScoreBulkScorer bitset path. (Ke Wei)
+* GITHUB#16696: Merging an HNSW graph over Lucene104HnswScalarQuantizedVectorsFormat, with an
+ asymmetric encoding of a float32 field, no longer reads the merged float vectors back to build
+ the merge scorer. (John Maurice)
+
* GITHUB#16669: Use intoBitSet to speed up ReqExclBulkScorer with dense exclusions. (Ke Wei)
Bug Fixes
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsWriter.java
index 4b0ff173065..937bcf51dd2 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsWriter.java
@@ -17,11 +17,14 @@
package org.apache.lucene.codecs.hnsw;
+import java.io.Closeable;
import java.io.IOException;
+import java.util.function.IntPredicate;
import org.apache.lucene.codecs.KnnVectorsWriter;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.MergeState;
import org.apache.lucene.util.IORunnable;
+import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier;
/**
* Vectors' writer for a field that allows additional indexing logic to be implemented by the caller
@@ -63,4 +66,59 @@ public abstract class FlatVectorsWriter extends KnnVectorsWriter {
public abstract void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState mergeState)
throws IOException;
+
+ /**
+ * Merges one field and, while writing the merged vectors, optionally prepares what a merge scorer
+ * needs, so that the merged float vectors need not be read back.
+ *
+ * <p>A wrapping writer should forward this method along with {@link #mergeOneFlatVectorField}.
+ * Forwarding only the latter stays correct, but the wrapped writer then prepares nothing.
+ * Forwarding helps only when the wrapping format's reader over the merged segment is the wrapped
+ * writer's own reader. Otherwise {@link MergeScorerData#scorerSupplier} rejects the prepared data
+ * and the merge uses its fallback.
+ *
+ * @param fieldInfo field to merge
+ * @param mergeState merge state
+ * @param needsMergeScorer whether to prepare scorer data for a vector count. The writer tests
+ * this predicate against the count before deletions, which can exceed the number of vectors
+ * written.
+ * @return prepared data that the caller must close, or {@code null} if none was prepared, which
+ * is what the default implementation returns and what any writer returns for a field it
+ * cannot prepare for. The caller then builds the scorer itself.
+ * @throws IOException if an I/O error occurs
+ */
+ public MergeScorerData mergeOneFlatVectorFieldForMergeScorer(
+ FieldInfo fieldInfo, MergeState mergeState, IntPredicate needsMergeScorer)
+ throws IOException {
+ mergeOneFlatVectorField(fieldInfo, mergeState);
+ return null;
+ }
+
+ /**
+ * Data prepared while merging a field for scoring its merged vectors. The caller must close it
+ * whether or not it obtains a supplier; closing releases only what {@link #scorerSupplier} did
+ * not hand over. A single merge thread uses each instance, so implementations need not be
+ * thread-safe.
+ */
+ public interface MergeScorerData extends Closeable {
+
+ /**
+ * Creates a scorer supplier for {@code mergedReader}. The caller must close the supplier to
+ * release the prepared data.
+ *
+ * <p>The reader is matched to the preparing writer's own reader by type; any other reader is
+ * rejected.
+ *
+ * <p>This method may be called at most once, after the preparing writer is finished and closed
+ * and {@code mergedReader} is open on the segment it wrote. A second call, or a call after
+ * {@link #close}, throws {@link IllegalStateException}.
+ *
+ * @param mergedReader a reader over the segment the writer wrote
+ * @return a scorer supplier over the merged vectors, or {@code null} if this data cannot serve
+ * that reader
+ * @throws IOException if an I/O error occurs
+ */
+ CloseableRandomVectorScorerSupplier scorerSupplier(FlatVectorsReader mergedReader)
+ throws IOException;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java
index 1ff69527da3..c49d8499733 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java
@@ -24,7 +24,10 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
import org.apache.lucene.index.SegmentReadState;
import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.store.NoReuseHint;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer.QuantizationResult;
import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
/**
@@ -105,6 +108,31 @@ public class Lucene104ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
static final String VECTOR_DATA_EXTENSION = "veq";
static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
+ /**
+ * Writes the corrective terms that follow a quantized record. Index-side records and query-side
+ * records share this layout, so this is the one definition of it.
+ */
+ static void writeCorrections(
+ IndexOutput output, OptimizedScalarQuantizer.QuantizationResult corrections)
+ throws IOException {
+ output.writeInt(Float.floatToIntBits(corrections.lowerInterval()));
+ output.writeInt(Float.floatToIntBits(corrections.upperInterval()));
+ output.writeInt(Float.floatToIntBits(corrections.additionalCorrection()));
+ output.writeInt(corrections.quantizedComponentSum());
+ }
+
+ /**
+ * Packs a quantized query-side record from {@code scratch} into {@code dest} and writes it with
+ * its corrections. The writer and the reader produce query-side records through this method.
+ */
+ static void writeQueryRecord(
+ IndexOutput out, byte[] scratch, byte[] dest, QuantizationResult corrections)
+ throws IOException {
+ OptimizedScalarQuantizer.transposeHalfByte(scratch, dest);
+ out.writeBytes(dest, dest.length);
+ writeCorrections(out, corrections);
+ }
+
private static final FlatVectorsFormat rawVectorFormat =
new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java
index 2a7ed807aeb..9fd6d94e6c4 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java
@@ -17,10 +17,10 @@
package org.apache.lucene.codecs.lucene104;
import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION;
+import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.writeQueryRecord;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readSimilarityFunction;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
-import static org.apache.lucene.util.quantization.OptimizedScalarQuantizer.transposeHalfByte;
import java.io.IOException;
import java.util.HashMap;
@@ -48,6 +48,7 @@ import org.apache.lucene.search.KnnCollector;
import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.DataAccessHint;
+import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FileDataHint;
import org.apache.lucene.store.FileTypeHint;
import org.apache.lucene.store.IOContext;
@@ -506,36 +507,68 @@ public class Lucene104ScalarQuantizedVectorsReader extends FlatVectorsReader
}
throw t;
}
- IndexInput quantizedScoreDataInput =
- segmentWriteState.directory.openInput(
- tempScoreQuantizedVectorName, queryDataContext(segmentWriteState));
+ assert docsWithField.cardinality() == vectorValues.size();
+ return mergeScorerSupplier(
+ fieldInfo,
+ vectorValues,
+ vectorScorer,
+ segmentWriteState.directory,
+ queryDataContext(segmentWriteState),
+ tempScoreQuantizedVectorName);
+ }
+
+ /**
+ * Builds a supplier that scores query-side records in {@code queryDataName} against the segment's
+ * own {@code indexVectors}. The query records use the layout written by {@link
+ * #writeBinarizedQueryData}.
+ *
+ * <p>This method takes ownership of the file. The returned supplier deletes it when closed, and
+ * this method deletes it if supplier construction fails.
+ */
+ static CloseableRandomVectorScorerSupplier mergeScorerSupplier(
+ FieldInfo fieldInfo,
+ QuantizedByteVectorValues indexVectors,
+ Lucene104ScalarQuantizedVectorScorer vectorScorer,
+ Directory directory,
+ IOContext context,
+ String queryDataName)
+ throws IOException {
+ IndexInput queryData = null;
try {
- OffHeapScalarQuantizedVectorValues scoreVectorValues =
+ queryData = directory.openInput(queryDataName, context);
+ // Query-side and index-side data contain one record per field vector.
+ OffHeapScalarQuantizedVectorValues queryVectors =
new OffHeapScalarQuantizedVectorValues.DenseOffHeapVectorValues(
true,
fieldInfo.getVectorDimension(),
- docsWithField.cardinality(),
- vectorValues.getCentroid(),
- vectorValues.getCentroidDP(),
- quantizer,
- fi.scalarEncoding,
+ indexVectors.size(),
+ indexVectors.getCentroid(),
+ indexVectors.getCentroidDP(),
+ new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction()),
+ indexVectors.getScalarEncoding(),
fieldInfo.getVectorSimilarityFunction(),
vectorScorer,
- quantizedScoreDataInput);
+ queryData);
RandomVectorScorerSupplier scorerSupplier =
vectorScorer.getRandomVectorScorerSupplier(
- fieldInfo.getVectorSimilarityFunction(), scoreVectorValues, vectorValues);
- final String finalTempScoreQuantizedVectorName = tempScoreQuantizedVectorName;
+ fieldInfo.getVectorSimilarityFunction(), queryVectors, indexVectors);
+ final IndexInput queryDataInput = queryData;
return CloseableRandomVectorScorerSupplier.create(
scorerSupplier,
- vectorValues.size(),
+ indexVectors.size(),
() -> {
- IOUtils.close(quantizedScoreDataInput);
- IOUtils.deleteFilesIgnoringExceptions(
- segmentWriteState.directory, finalTempScoreQuantizedVectorName);
+ try {
+ IOUtils.close(queryDataInput);
+ } finally {
+ IOUtils.deleteFilesIgnoringExceptions(directory, queryDataName);
+ }
});
} catch (Throwable t) {
- IOUtils.closeWhileSuppressingExceptions(t, quantizedScoreDataInput);
+ try {
+ IOUtils.closeWhileSuppressingExceptions(t, queryData);
+ } finally {
+ IOUtils.deleteFilesIgnoringExceptions(directory, queryDataName);
+ }
throw t;
}
}
@@ -564,13 +597,7 @@ public class Lucene104ScalarQuantizedVectorsReader extends FlatVectorsReader
encoding.getQueryBits(),
quantizedByteVectorValues.getCentroid());
docsWithField.add(docV);
- // pack and store the 4bit query vector
- transposeHalfByte(quantizationScratch, toQuery);
- binarizedQueryData.writeBytes(toQuery, toQuery.length);
- binarizedQueryData.writeInt(Float.floatToIntBits(r.lowerInterval()));
- binarizedQueryData.writeInt(Float.floatToIntBits(r.upperInterval()));
- binarizedQueryData.writeInt(Float.floatToIntBits(r.additionalCorrection()));
- binarizedQueryData.writeInt(r.quantizedComponentSum());
+ writeQueryRecord(binarizedQueryData, quantizationScratch, toQuery, r);
}
return docsWithField;
}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java
index 0311fd39456..27d16d92cf9 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java
@@ -18,6 +18,8 @@ package org.apache.lucene.codecs.lucene104;
import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.DIRECT_MONOTONIC_BLOCK_SHIFT;
import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.QUANTIZED_VECTOR_COMPONENT;
+import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.writeCorrections;
+import static org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat.writeQueryRecord;
import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance;
@@ -28,9 +30,11 @@ import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.function.IntPredicate;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
@@ -47,10 +51,16 @@ import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.internal.hppc.FloatArrayList;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.store.DataAccessHint;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FileDataHint;
+import org.apache.lucene.store.FileTypeHint;
+import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier;
import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
@@ -69,6 +79,7 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
private final IndexOutput meta, vectorData;
private final ScalarEncoding encoding;
private final FlatVectorsWriter rawVectorDelegate;
+ private final Lucene104ScalarQuantizedVectorScorer vectorScorer;
private boolean finished;
/** Sole constructor */
@@ -79,6 +90,7 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
Lucene104ScalarQuantizedVectorScorer vectorsScorer)
throws IOException {
super(vectorsScorer);
+ this.vectorScorer = vectorsScorer;
this.encoding = encoding;
this.segmentWriteState = state;
String metaFileName =
@@ -197,17 +209,9 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
float[] v = fieldData.getVectors().get(i);
OptimizedScalarQuantizer.QuantizationResult corrections =
scalarQuantizer.scalarQuantize(v, scratch, encoding.getBits(), clusterCenter);
- switch (encoding) {
- case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(scratch, vector);
- case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(scratch, vector);
- case DIBIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.transposeDibit(scratch, vector);
- case UNSIGNED_BYTE, SEVEN_BIT -> {}
- }
+ packIndexRecord(encoding, scratch, vector);
vectorData.writeBytes(vector, vector.length);
- vectorData.writeInt(Float.floatToIntBits(corrections.lowerInterval()));
- vectorData.writeInt(Float.floatToIntBits(corrections.upperInterval()));
- vectorData.writeInt(Float.floatToIntBits(corrections.additionalCorrection()));
- vectorData.writeInt(corrections.quantizedComponentSum());
+ writeCorrections(vectorData, corrections);
}
}
@@ -258,17 +262,9 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
float[] v = fieldData.getVectors().get(ordinal);
OptimizedScalarQuantizer.QuantizationResult corrections =
scalarQuantizer.scalarQuantize(v, scratch, encoding.getBits(), clusterCenter);
- switch (encoding) {
- case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(scratch, vector);
- case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(scratch, vector);
- case DIBIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.transposeDibit(scratch, vector);
- case UNSIGNED_BYTE, SEVEN_BIT -> {}
- }
+ packIndexRecord(encoding, scratch, vector);
vectorData.writeBytes(vector, vector.length);
- vectorData.writeInt(Float.floatToIntBits(corrections.lowerInterval()));
- vectorData.writeInt(Float.floatToIntBits(corrections.upperInterval()));
- vectorData.writeInt(Float.floatToIntBits(corrections.additionalCorrection()));
- vectorData.writeInt(corrections.quantizedComponentSum());
+ writeCorrections(vectorData, corrections);
}
}
@@ -319,46 +315,217 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
}
}
+ /** The merged float vectors of the field, as the quantizer sees them: normalized for COSINE. */
+ private FloatVectorValues mergedFloatVectorValues(FieldInfo fieldInfo, MergeState mergeState)
+ throws IOException {
+ FloatVectorValues vectorValues =
+ MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
+ if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
+ vectorValues = new NormalizedFloatVectorValues(vectorValues);
+ }
+ return vectorValues;
+ }
+
+ private QuantizedByteVectorValues mergedQuantizedVectorValues(
+ FieldInfo fieldInfo, MergeState mergeState, float[] centroid) throws IOException {
+ OptimizedScalarQuantizer quantizer =
+ new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
+ return new QuantizedFloatVectorValues(
+ mergedFloatVectorValues(fieldInfo, mergeState), quantizer, encoding, centroid);
+ }
+
@Override
public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState mergeState)
throws IOException {
+ mergeOneFlatVectorField(fieldInfo, mergeState, vectorCount -> false);
+ }
+
+ /**
+ * {@inheritDoc}
+ *
+ * <p>HNSW merges call this overload rather than {@link #mergeOneFlatVectorField(FieldInfo,
+ * MergeState)}, so subclasses that customize merging must override this method too.
+ */
+ @Override
+ public MergeScorerData mergeOneFlatVectorFieldForMergeScorer(
+ FieldInfo fieldInfo, MergeState mergeState, IntPredicate needsMergeScorer)
+ throws IOException {
+ return mergeOneFlatVectorField(fieldInfo, mergeState, needsMergeScorer);
+ }
+
+ private MergeScorerData mergeOneFlatVectorField(
+ FieldInfo fieldInfo, MergeState mergeState, IntPredicate needsMergeScorer)
+ throws IOException {
// Don't need access to the random vectors, we can just use the merged
rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
if (!fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
- return;
+ return null;
}
- final float[] centroid;
final float[] mergedCentroid = new float[fieldInfo.getVectorDimension()];
int vectorCount = mergeAndRecalculateCentroids(mergeState, fieldInfo, mergedCentroid);
- centroid = mergedCentroid;
if (segmentWriteState.infoStream.isEnabled(QUANTIZED_VECTOR_COMPONENT)) {
segmentWriteState.infoStream.message(
QUANTIZED_VECTOR_COMPONENT, "Vectors' count:" + vectorCount);
}
- FloatVectorValues floatVectorValues =
- MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
- if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
- floatVectorValues = new NormalizedFloatVectorValues(floatVectorValues);
- }
- QuantizedFloatVectorValues quantizedVectorValues =
- new QuantizedFloatVectorValues(
- floatVectorValues,
- new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction()),
- encoding,
- centroid);
+ // Only asymmetric encodings have query-side records. The predicate sees vectorCount before
+ // deletions, so it can request data even when the merged field is too small to build a graph.
+ boolean prepareQueryData = encoding.isAsymmetric() && needsMergeScorer.test(vectorCount);
long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES);
- DocsWithFieldSet docsWithField = writeVectorData(vectorData, quantizedVectorValues);
- long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
- float centroidDp =
- docsWithField.cardinality() > 0 ? VectorUtil.dotProduct(centroid, centroid) : 0;
- writeMeta(
- fieldInfo,
- segmentWriteState.segmentInfo.maxDoc(),
- vectorDataOffset,
- vectorDataLength,
- centroid,
- centroidDp,
- docsWithField);
+ DocsWithFieldSet docsWithField;
+ MergeScorerData mergeScorerData = null;
+ if (prepareQueryData) {
+ docsWithField = new DocsWithFieldSet();
+ mergeScorerData =
+ writeVectorAndQueryData(fieldInfo, mergeState, mergedCentroid, docsWithField);
+ } else {
+ QuantizedByteVectorValues quantizedVectorValues =
+ mergedQuantizedVectorValues(fieldInfo, mergeState, mergedCentroid);
+ docsWithField = writeVectorData(vectorData, quantizedVectorValues);
+ }
+ try {
+ long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
+ float centroidDp =
+ docsWithField.cardinality() > 0
+ ? VectorUtil.dotProduct(mergedCentroid, mergedCentroid)
+ : 0;
+ writeMeta(
+ fieldInfo,
+ segmentWriteState.segmentInfo.maxDoc(),
+ vectorDataOffset,
+ vectorDataLength,
+ mergedCentroid,
+ centroidDp,
+ docsWithField);
+ return mergeScorerData;
+ } catch (Throwable t) {
+ // The handle was never returned, so nobody else can release its records.
+ IOUtils.closeWhileSuppressingExceptions(t, mergeScorerData);
+ throw t;
+ }
+ }
+
+ /**
+ * The quantized query vectors staged for the merge scorer. A graph build reads them by ordinal as
+ * it walks, so they are read at random.
+ */
+ private IOContext queryDataContext() {
+ return segmentWriteState.context.withHints(
+ FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM);
+ }
+
+ /**
+ * Writes merged index-side and query-side records in one pass. Index-side records go to the
+ * quantized vector data file, while query-side records go to a temporary file owned by the
+ * returned handle.
+ *
+ * <p>{@link OptimizedScalarQuantizer#multiScalarQuantize} centers each vector once and quantizes
+ * it at both bit widths.
+ *
+ * @param docsWithField filled with the documents that were written
+ */
+ private MergeScorerData writeVectorAndQueryData(
+ FieldInfo fieldInfo, MergeState mergeState, float[] centroid, DocsWithFieldSet docsWithField)
+ throws IOException {
+ assert encoding.isAsymmetric();
+ OptimizedScalarQuantizer quantizer =
+ new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
+ FloatVectorValues vectorValues = mergedFloatVectorValues(fieldInfo, mergeState);
+ int discretizedDims = encoding.getDiscreteDimensions(vectorValues.dimension());
+ byte[] indexQuantized = new byte[discretizedDims];
+ byte[] queryQuantized = new byte[discretizedDims];
+ byte[] indexPacked =
+ switch (encoding) {
+ case UNSIGNED_BYTE, SEVEN_BIT -> indexQuantized;
+ case PACKED_NIBBLE, SINGLE_BIT_QUERY_NIBBLE, DIBIT_QUERY_NIBBLE ->
+ new byte[encoding.getDocPackedLength(discretizedDims)];
+ };
+ byte[] queryPacked = new byte[encoding.getQueryPackedLength(discretizedDims)];
+ byte[] bits = new byte[] {encoding.getBits(), encoding.getQueryBits()};
+ byte[][] destinations = new byte[][] {indexQuantized, queryQuantized};
+ String queryDataName = null;
+ try (IndexOutput queryData =
+ segmentWriteState.directory.createTempOutput(
+ segmentWriteState.segmentInfo.name, "queries", queryDataContext())) {
+ queryDataName = queryData.getName();
+ KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
+ for (int docV = iterator.nextDoc(); docV != NO_MORE_DOCS; docV = iterator.nextDoc()) {
+ OptimizedScalarQuantizer.QuantizationResult[] corrections =
+ quantizer.multiScalarQuantize(
+ vectorValues.vectorValue(iterator.index()), destinations, bits, centroid);
+ // the index side, packed as QuantizedFloatVectorValues packs it
+ packIndexRecord(encoding, indexQuantized, indexPacked);
+ vectorData.writeBytes(indexPacked, indexPacked.length);
+ writeCorrections(vectorData, corrections[0]);
+ writeQueryRecord(queryData, queryQuantized, queryPacked, corrections[1]);
+ docsWithField.add(docV);
+ }
+ CodecUtil.writeFooter(queryData);
+ } catch (Throwable t) {
+ if (queryDataName != null) {
+ IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, queryDataName);
+ }
+ throw t;
+ }
+ return new PreparedQueryData(
+ segmentWriteState.directory, queryDataContext(), fieldInfo, vectorScorer, queryDataName);
+ }
+
+ /** Owns one field's temporary query-side records until its graph scorer takes them over. */
+ private static final class PreparedQueryData implements MergeScorerData {
+ private final Directory directory;
+ private final IOContext context;
+ private final FieldInfo fieldInfo;
+ private final Lucene104ScalarQuantizedVectorScorer vectorScorer;
+ private final String fileName;
+ private boolean spent;
+
+ PreparedQueryData(
+ Directory directory,
+ IOContext context,
+ FieldInfo fieldInfo,
+ Lucene104ScalarQuantizedVectorScorer vectorScorer,
+ String fileName) {
+ this.directory = directory;
+ this.context = context;
+ this.fieldInfo = fieldInfo;
+ this.vectorScorer = vectorScorer;
+ this.fileName = fileName;
+ }
+
+ @Override
+ public CloseableRandomVectorScorerSupplier scorerSupplier(FlatVectorsReader mergedReader)
+ throws IOException {
+ if (spent) {
+ throw new IllegalStateException("already consumed or closed: " + fileName);
+ }
+ spent = true;
+ boolean handedOver = false;
+ try {
+ QuantizedByteVectorValues indexVectors =
+ mergedReader instanceof Lucene104ScalarQuantizedVectorsReader quantizedReader
+ ? quantizedReader.getQuantizedVectorValues(fieldInfo.name)
+ : null;
+ if (indexVectors == null) {
+ // not this format's reader
+ return null;
+ }
+ handedOver = true;
+ return Lucene104ScalarQuantizedVectorsReader.mergeScorerSupplier(
+ fieldInfo, indexVectors, vectorScorer, directory, context, fileName);
+ } finally {
+ if (handedOver == false) {
+ IOUtils.deleteFilesIgnoringExceptions(directory, fileName);
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ if (spent == false) {
+ spent = true;
+ IOUtils.deleteFilesIgnoringExceptions(directory, fileName);
+ }
+ }
}
static DocsWithFieldSet writeVectorData(
@@ -369,19 +536,34 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
// write vector
byte[] binaryValue = quantizedByteVectorValues.vectorValue(iterator.index());
output.writeBytes(binaryValue, binaryValue.length);
- OptimizedScalarQuantizer.QuantizationResult corrections =
- quantizedByteVectorValues.getCorrectiveTerms(iterator.index());
- output.writeInt(Float.floatToIntBits(corrections.lowerInterval()));
- output.writeInt(Float.floatToIntBits(corrections.upperInterval()));
- output.writeInt(Float.floatToIntBits(corrections.additionalCorrection()));
- output.writeInt(corrections.quantizedComponentSum());
+ writeCorrections(output, quantizedByteVectorValues.getCorrectiveTerms(iterator.index()));
docsWithField.add(docV);
}
return docsWithField;
}
+ /**
+ * Packs a quantized index-side record from {@code scratch} into {@code dest}. {@link
+ * ScalarEncoding#UNSIGNED_BYTE} and {@link ScalarEncoding#SEVEN_BIT} write the quantized bytes as
+ * they are, so for them this is a copy, and callers that pass the same array for both save it.
+ */
+ private static void packIndexRecord(ScalarEncoding encoding, byte[] scratch, byte[] dest) {
+ switch (encoding) {
+ case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(scratch, dest);
+ case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(scratch, dest);
+ case DIBIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.transposeDibit(scratch, dest);
+ case UNSIGNED_BYTE, SEVEN_BIT -> {
+ if (dest != scratch) {
+ System.arraycopy(scratch, 0, dest, 0, scratch.length);
+ }
+ }
+ }
+ }
+
@Override
public void close() throws IOException {
+ // Query-side records are deliberately not released here: this writer is closed before the
+ // caller can open the reader a merge scorer needs, so every handle it handed out is pending.
IOUtils.close(meta, vectorData, rawVectorDelegate);
}
@@ -676,12 +858,7 @@ public class Lucene104ScalarQuantizedVectorsWriter extends FlatVectorsWriter {
corrections =
quantizer.scalarQuantize(
values.vectorValue(ord), quantized, encoding.getBits(), centroid);
- switch (encoding) {
- case PACKED_NIBBLE -> OffHeapScalarQuantizedVectorValues.packNibbles(quantized, packed);
- case SINGLE_BIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.packAsBinary(quantized, packed);
- case DIBIT_QUERY_NIBBLE -> OptimizedScalarQuantizer.transposeDibit(quantized, packed);
- case UNSIGNED_BYTE, SEVEN_BIT -> {}
- }
+ packIndexRecord(encoding, quantized, packed);
}
@Override
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
index 401f69654db..55f03546b47 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
@@ -25,6 +25,7 @@ import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.VERSIO
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
import static org.apache.lucene.util.hnsw.HnswGraphSearcher.expectedVisitedNodes;
+import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
@@ -38,6 +39,7 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.codecs.hnsw.FlatVectorsWriter.MergeScorerData;
import org.apache.lucene.index.DocsWithFieldSet;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexFileNames;
@@ -93,6 +95,13 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
private final int version;
private final List<FieldWriter<?>> fields = new ArrayList<>();
+
+ /**
+ * Tracks prepared merge-scorer data so {@link #close()} can release handles that are never
+ * consumed. Closing a consumed handle is a no-op.
+ */
+ private final List<MergeScorerData> preparedMergeScorers = new ArrayList<>();
+
private boolean finished;
public Lucene99HnswVectorsWriter(
@@ -424,7 +433,12 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
@Override
public IORunnable mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
- flatVectorWriter.mergeOneFlatVectorField(fieldInfo, mergeState);
+ MergeScorerData mergeScorerData =
+ flatVectorWriter.mergeOneFlatVectorFieldForMergeScorer(
+ fieldInfo, mergeState, this::buildsGraph);
+ if (mergeScorerData != null) {
+ preparedMergeScorers.add(mergeScorerData);
+ }
return () -> {
// Bail out before the potentially heavy graph build if the merge was already aborted
mergeState.checkAborted();
@@ -437,22 +451,17 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
case FLOAT32 -> flatVectorsReader.getFloatVectorValues(fieldInfo.name);
};
int totalVectorCount = vectorValues == null ? 0 : vectorValues.size();
- if (totalVectorCount > 0 && shouldCreateGraph(tinySegmentsThreshold, totalVectorCount)) {
- if (flatVectorsReader instanceof QuantizedVectorsReader quantizedVectorsReader
- && fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
- CloseableRandomVectorScorerSupplier scorerSupplier =
- quantizedVectorsReader.getRandomVectorScorerSupplierForMerge(
- fieldInfo,
- segmentWriteState.withHints(
- FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM));
+ if (buildsGraph(totalVectorCount)) {
+ CloseableRandomVectorScorerSupplier mergeScorer =
+ mergeScorerSupplier(fieldInfo, mergeScorerData);
+ if (mergeScorer != null) {
try {
- buildAndWriteGraph(
- fieldInfo, mergeState, vectorValues, scorerSupplier, totalVectorCount);
+ buildAndWriteGraph(fieldInfo, mergeState, vectorValues, mergeScorer, totalVectorCount);
} catch (Throwable t) {
- IOUtils.closeWhileSuppressingExceptions(t, scorerSupplier);
+ IOUtils.closeWhileSuppressingExceptions(t, mergeScorer);
throw t;
}
- IOUtils.close(scorerSupplier);
+ IOUtils.close(mergeScorer);
} else {
RandomVectorScorerSupplier scorerSupplier =
flatVectorsReader
@@ -467,6 +476,28 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
};
}
+ /**
+ * Returns the merge scorer the flat writer prepared, or the one the reader builds from the merged
+ * vectors. Returns {@code null} if neither path supports the field.
+ */
+ private CloseableRandomVectorScorerSupplier mergeScorerSupplier(
+ FieldInfo fieldInfo, MergeScorerData prepared) throws IOException {
+ if (prepared != null) {
+ CloseableRandomVectorScorerSupplier supplier = prepared.scorerSupplier(flatVectorsReader);
+ if (supplier != null) {
+ return supplier;
+ }
+ }
+ if (flatVectorsReader instanceof QuantizedVectorsReader quantizedVectorsReader
+ && fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
+ return quantizedVectorsReader.getRandomVectorScorerSupplierForMerge(
+ fieldInfo,
+ segmentWriteState.withHints(
+ FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM));
+ }
+ return null;
+ }
+
private void buildAndWriteGraph(
FieldInfo fieldInfo,
MergeState mergeState,
@@ -659,10 +690,13 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
@Override
public void close() throws IOException {
+ // Release handles left by aborted merges, fields that built no graph after deletions, and
+ // fields skipped after another graph build failed.
+ Closeable handles = () -> IOUtils.close(preparedMergeScorers);
if (flatWriterClosed) {
- IOUtils.close(meta, vectorIndex, flatVectorsReader);
+ IOUtils.close(meta, vectorIndex, flatVectorsReader, handles);
} else {
- IOUtils.close(meta, vectorIndex, flatVectorWriter, flatVectorsReader);
+ IOUtils.close(meta, vectorIndex, flatVectorWriter, flatVectorsReader, handles);
}
}
@@ -675,6 +709,10 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
throw new IllegalArgumentException("invalid distance function: " + func);
}
+ private boolean buildsGraph(int numVectors) {
+ return numVectors > 0 && shouldCreateGraph(tinySegmentsThreshold, numVectors);
+ }
+
private static boolean shouldCreateGraph(int k, int numNodes) {
if (k <= 0) {
return true;
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedMergeScorer.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedMergeScorer.java
new file mode 100644
index 00000000000..6a3fd5d5b40
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedMergeScorer.java
@@ -0,0 +1,797 @@
+/*
+ * 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.lucene104;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.IntPredicate;
+import java.util.function.LongConsumer;
+import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.KnnVectorsWriter;
+import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
+import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
+import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.codecs.hnsw.FlatVectorsWriter.MergeScorerData;
+import org.apache.lucene.codecs.hnsw.HnswGraphProvider;
+import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FieldInfos;
+import org.apache.lucene.index.FilterCodecReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MergeState;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentInfos;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.index.SerialMergeScheduler;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.FilterIndexInput;
+import org.apache.lucene.store.FilterIndexOutput;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.hnsw.HnswGraph;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Tests writer-produced merge scorer data for asymmetric scalar encodings. The writer creates
+ * query-side records while writing merged vectors, avoiding the reader's second quantization pass.
+ *
+ * <p>The tests pin that the read-back is gone, compare both paths' output, and verify that handoff
+ * files do not outlive their merge. {@code TestHnswMergeAbort#testRollbackDuringQuantizedMerge}
+ * covers aborts after the scorer supplier takes ownership.
+ */
+public class TestLucene104ScalarQuantizedMergeScorer extends LuceneTestCase {
+
+ private static final int DIM = 64;
+
+ private static final int DOCS_PER_SEGMENT = 50;
+
+ /** Segment size that stays below the default HNSW graph threshold when two segments merge. */
+ private static final int TINY_SEGMENT_DOCS = 200;
+
+ private static final int MAX_CONN = 16;
+ private static final int BEAM_WIDTH = 32;
+
+ /** A threshold of 0 makes every merge build a graph, so the merge scorer is always requested. */
+ private static final int ALWAYS_GRAPH = 0;
+
+ private static final FlatVectorsFormat RAW_FORMAT =
+ new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+ private static final Lucene104ScalarQuantizedVectorScorer QUANTIZED_SCORER =
+ new Lucene104ScalarQuantizedVectorScorer(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+
+ private static List<ScalarEncoding> asymmetricEncodings() {
+ List<ScalarEncoding> encodings = new ArrayList<>();
+ for (ScalarEncoding encoding : ScalarEncoding.values()) {
+ if (encoding.isAsymmetric()) {
+ encodings.add(encoding);
+ }
+ }
+ assertFalse("no asymmetric encoding to test", encodings.isEmpty());
+ return encodings;
+ }
+
+ /** The shipped format: the writer prepares the merge scorer's records. */
+ private static KnnVectorsFormat writerPathFormat(ScalarEncoding encoding, int threshold) {
+ return new Lucene104HnswScalarQuantizedVectorsFormat(
+ encoding, MAX_CONN, BEAM_WIDTH, 1, null, threshold);
+ }
+
+ /**
+ * Returns the same on-disk format with writer preparation disabled, forcing the {@code
+ * QuantizedVectorsReader} fallback. Tests use it as a read-back control and to retain fallback
+ * coverage.
+ */
+ private static KnnVectorsFormat readerPathFormat(ScalarEncoding encoding, int threshold) {
+ return new HnswOverFlatFormat(new NoPrepareFlatFormat(encoding), threshold);
+ }
+
+ /**
+ * Verifies that the writer path does not read the merged raw float vectors back; the graph build
+ * still streams the merged quantized records. Opening the merged reader still reads the raw
+ * vector header and checksum, so the limit is one vector rather than zero bytes.
+ *
+ * <p>The reader fallback is a positive control and must stream the entire merged raw vector file.
+ */
+ public void testMergeDoesNotReadMergedVectorsBack() throws IOException {
+ for (ScalarEncoding encoding : asymmetricEncodings()) {
+ for (VectorSimilarityFunction similarity :
+ new VectorSimilarityFunction[] {
+ VectorSimilarityFunction.EUCLIDEAN, VectorSimilarityFunction.COSINE
+ }) {
+ MergeCounts counts =
+ runMerge(writerPathFormat(encoding, ALWAYS_GRAPH), similarity, true, true);
+ assertTrue(
+ "the merge read the merged vectors back: "
+ + counts.mergedRawBytesRead()
+ + " bytes of "
+ + encoding
+ + "/"
+ + similarity,
+ counts.mergedRawBytesRead() < (long) DIM * Float.BYTES);
+ }
+ MergeCounts fallback =
+ runMerge(
+ readerPathFormat(encoding, ALWAYS_GRAPH),
+ VectorSimilarityFunction.EUCLIDEAN,
+ true,
+ true);
+ assertTrue(
+ "the reader fallback did not stream the merged vectors back: "
+ + fallback.mergedRawBytesRead()
+ + " bytes for "
+ + encoding,
+ fallback.mergedRawBytesRead() >= (long) DIM * Float.BYTES * 2 * DOCS_PER_SEGMENT);
+ }
+ }
+
+ /**
+ * Verifies that a merge with no graph never creates a handoff file. This covers an HNSW merge
+ * rejected by the graph-size predicate and a flat merge that supplies no predicate to the
+ * two-argument merge method.
+ *
+ * <p>The file must never be created, not merely cleaned up.
+ */
+ public void testNoHandOffWhenNoGraphIsBuilt() throws IOException {
+ for (ScalarEncoding encoding : asymmetricEncodings()) {
+ MergeCounts belowThreshold =
+ runMerge(
+ writerPathFormat(encoding, Lucene99HnswVectorsFormat.HNSW_GRAPH_THRESHOLD),
+ VectorSimilarityFunction.EUCLIDEAN,
+ false,
+ true,
+ TINY_SEGMENT_DOCS);
+ assertEquals(
+ "a hand-off file was created for a merge that builds no graph: "
+ + belowThreshold.handOffs(),
+ List.of(),
+ belowThreshold.handOffs());
+
+ MergeCounts flatOnly =
+ runMerge(
+ new Lucene104ScalarQuantizedVectorsFormat(encoding),
+ VectorSimilarityFunction.EUCLIDEAN,
+ false,
+ false);
+ assertEquals(
+ "the flat format prepared merge scorer data nobody can ask for: " + flatOnly.handOffs(),
+ List.of(),
+ flatOnly.handOffs());
+ }
+ }
+
+ /**
+ * Verifies that {@link IndexWriter#addIndexes(CodecReader...)} merges readers whose field numbers
+ * differ from their underlying segments. The flat writer must resolve the field by name.
+ */
+ public void testMergeOfRenumberedCodecReaders() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ IndexWriterConfig config =
+ new IndexWriterConfig()
+ .setCodec(TestUtil.alwaysKnnVectorsFormat(writerPathFormat(encoding, ALWAYS_GRAPH)))
+ .setUseCompoundFile(false);
+ config.getMergePolicy().setNoCFSRatio(0.0);
+ try (Directory source = newDirectory();
+ Directory target = newDirectory()) {
+ float[][] vectors = randomVectors(2 * DOCS_PER_SEGMENT, VectorSimilarityFunction.EUCLIDEAN);
+ try (IndexWriter writer = new IndexWriter(source, config)) {
+ for (int i = 0; i < vectors.length; i++) {
+ Document doc = new Document();
+ for (int f = 0; f < 5; f++) {
+ doc.add(new StringField("s" + f, Integer.toString(i % 7), Field.Store.NO));
+ }
+ doc.add(new KnnFloatVectorField("v", vectors[i], VectorSimilarityFunction.EUCLIDEAN));
+ writer.addDocument(doc);
+ if (i == DOCS_PER_SEGMENT - 1) {
+ writer.commit();
+ }
+ }
+ writer.commit();
+ }
+ try (DirectoryReader reader = DirectoryReader.open(source);
+ IndexWriter writer =
+ new IndexWriter(
+ target,
+ new IndexWriterConfig().setCodec(config.getCodec()).setUseCompoundFile(false))) {
+ List<CodecReader> readers = new ArrayList<>();
+ for (LeafReaderContext context : reader.leaves()) {
+ readers.add(renumbered((CodecReader) context.reader()));
+ }
+ assertEquals(2, readers.size());
+ writer.addIndexes(readers.toArray(new CodecReader[0]));
+ writer.forceMerge(1);
+ }
+ assertNoTempFiles(target);
+ try (DirectoryReader reader = DirectoryReader.open(target)) {
+ assertEquals(vectors.length, reader.numDocs());
+ KnnVectorsReader vectorsReader = fieldVectorsReader(reader, "v");
+ HnswGraph graph = ((HnswGraphProvider) vectorsReader).getGraph("v");
+ assertNotNull("the merged segment has no graph", graph);
+ assertEquals(vectors.length, graph.size());
+ }
+ }
+ }
+
+ /**
+ * Verifies that each handoff supplies at most one scorer. Closing it after transfer is a no-op,
+ * and later supplier requests fail.
+ */
+ public void testHandOffIsSingleUse() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ List<MergeScorerData> handles = new ArrayList<>();
+ runMerge(
+ new HnswOverFlatFormat(new CapturingFlatFormat(encoding, handles), ALWAYS_GRAPH),
+ VectorSimilarityFunction.EUCLIDEAN,
+ true,
+ true);
+ assertEquals("the merge prepared no hand-off to test", 1, handles.size());
+ MergeScorerData handle = handles.get(0);
+ expectThrows(IllegalStateException.class, () -> handle.scorerSupplier(null));
+ handle.close();
+ expectThrows(IllegalStateException.class, () -> handle.scorerSupplier(null));
+ }
+
+ /**
+ * Verifies that a phase-one failure releases an earlier field's unconsumed handoff. The second
+ * field fails before the first field's deferred graph work runs, so closing the handle must
+ * remove its temporary file.
+ */
+ public void testPhaseOneFailureLeavesNoHandOff() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ try (Directory dir =
+ new FilterDirectory(newDirectory()) {
+ private final AtomicBoolean firstHandOff = new AtomicBoolean(true);
+
+ @Override
+ public IndexOutput createTempOutput(String prefix, String suffix, IOContext context)
+ throws IOException {
+ if ("queries".equals(suffix) && firstHandOff.compareAndSet(true, false) == false) {
+ throw new IOException("simulated failure creating the second hand-off file");
+ }
+ return super.createTempOutput(prefix, suffix, context);
+ }
+ }) {
+ assertMergeFailsWithoutLeftovers(dir, encoding, false);
+ }
+ }
+
+ /**
+ * Verifies that failure while reopening merged quantized vectors releases the handoff before its
+ * scorer supplier consumes it.
+ */
+ public void testHandOffIsReleasedWhenItIsNeverConsumed() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ try (Directory dir =
+ new FilterDirectory(newDirectory()) {
+ private final AtomicBoolean handedOff = new AtomicBoolean();
+
+ @Override
+ public IndexOutput createTempOutput(String prefix, String suffix, IOContext context)
+ throws IOException {
+ IndexOutput out = super.createTempOutput(prefix, suffix, context);
+ if ("queries".equals(suffix)) {
+ handedOff.set(true);
+ }
+ return out;
+ }
+
+ @Override
+ public IndexInput openInput(String name, IOContext context) throws IOException {
+ // the merged quantized vectors are opened when the graph build reopens the segment,
+ // which is before the handle is asked for anything
+ if (handedOff.get()
+ && name.endsWith(
+ "." + Lucene104ScalarQuantizedVectorsFormat.VECTOR_DATA_EXTENSION)) {
+ throw new IOException("simulated failure reopening the merged quantized vectors");
+ }
+ return super.openInput(name, context);
+ }
+ }) {
+ assertMergeFailsWithoutLeftovers(dir, encoding, true);
+ }
+ }
+
+ /** Verifies that the handoff is deleted when its scorer cannot open the temporary file. */
+ public void testPhaseTwoFailureLeavesNoHandOff() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ try (Directory dir =
+ new FilterDirectory(newDirectory()) {
+ @Override
+ public IndexInput openInput(String name, IOContext context) throws IOException {
+ if (name.contains("queries")) {
+ throw new IOException("simulated failure opening the hand-off file");
+ }
+ return super.openInput(name, context);
+ }
+ }) {
+ assertMergeFailsWithoutLeftovers(dir, encoding, true);
+ }
+ }
+
+ /**
+ * Verifies that a metadata-write failure deletes the handoff after its records are written but
+ * before it is returned to {@link Lucene99HnswVectorsWriter}. Only the flat writer knows the
+ * temporary file exists in this interval.
+ */
+ public void testMetaFailureLeavesNoHandOff() throws IOException {
+ ScalarEncoding encoding = randomAsymmetricEncoding();
+ try (Directory dir =
+ new FilterDirectory(newDirectory()) {
+ private final AtomicBoolean handedOff = new AtomicBoolean();
+
+ @Override
+ public IndexOutput createTempOutput(String prefix, String suffix, IOContext context)
+ throws IOException {
+ IndexOutput out = super.createTempOutput(prefix, suffix, context);
+ if ("queries".equals(suffix)) {
+ // Start failing metadata writes only after the handoff file has been created.
+ handedOff.set(true);
+ }
+ return out;
+ }
+
+ @Override
+ public IndexOutput createOutput(String name, IOContext context) throws IOException {
+ IndexOutput out = super.createOutput(name, context);
+ if (name.endsWith("." + Lucene104ScalarQuantizedVectorsFormat.META_EXTENSION)
+ == false) {
+ return out;
+ }
+ return new FilterIndexOutput("failing quantized meta", name, out) {
+ @Override
+ public void writeByte(byte b) throws IOException {
+ failIfHandedOff();
+ super.writeByte(b);
+ }
+
+ @Override
+ public void writeBytes(byte[] b, int offset, int length) throws IOException {
+ failIfHandedOff();
+ super.writeBytes(b, offset, length);
+ }
+
+ private void failIfHandedOff() throws IOException {
+ if (handedOff.get()) {
+ throw new IOException("simulated failure writing the quantized meta");
+ }
+ }
+ };
+ }
+ }) {
+ assertMergeFailsWithoutLeftovers(dir, encoding, true);
+ }
+ }
+
+ /**
+ * Runs a failing merge of one or two vector fields and verifies that the simulated failure is the
+ * one raised and that no temporary files remain.
+ */
+ private void assertMergeFailsWithoutLeftovers(
+ Directory dir, ScalarEncoding encoding, boolean singleField) throws IOException {
+ IndexWriterConfig config =
+ new IndexWriterConfig()
+ .setCodec(TestUtil.alwaysKnnVectorsFormat(writerPathFormat(encoding, ALWAYS_GRAPH)))
+ .setMergeScheduler(new SerialMergeScheduler())
+ .setUseCompoundFile(false);
+ config.getMergePolicy().setNoCFSRatio(0.0);
+ IndexWriter writer = new IndexWriter(dir, config);
+ try {
+ float[][] vectors = randomVectors(2 * DOCS_PER_SEGMENT, VectorSimilarityFunction.EUCLIDEAN);
+ for (int i = 0; i < vectors.length; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("v", vectors[i], VectorSimilarityFunction.EUCLIDEAN));
+ if (singleField == false) {
+ doc.add(
+ new KnnFloatVectorField(
+ "w", vectors[vectors.length - 1 - i], VectorSimilarityFunction.EUCLIDEAN));
+ }
+ writer.addDocument(doc);
+ if (i == DOCS_PER_SEGMENT - 1) {
+ writer.commit();
+ }
+ }
+ writer.commit();
+ Exception failure = expectThrows(Exception.class, () -> writer.forceMerge(1));
+ boolean simulated = false;
+ for (Throwable t = failure; t != null && simulated == false; t = t.getCause()) {
+ simulated = String.valueOf(t.getMessage()).contains("simulated failure");
+ }
+ assertTrue("unexpected merge failure: " + failure, simulated);
+ // Check before rollback. When no tragedy is recorded, rollback sweeps unreferenced temporary
+ // files and could hide a merge cleanup failure.
+ assertNoTempFiles(dir);
+ } finally {
+ // the writer may already be unusable after the failed merge, so this is best effort
+ IOUtils.closeWhileHandlingException(writer::rollback);
+ }
+ }
+
+ private ScalarEncoding randomAsymmetricEncoding() {
+ List<ScalarEncoding> encodings = asymmetricEncodings();
+ return encodings.get(random().nextInt(encodings.size()));
+ }
+
+ private static void assertNoTempFiles(Directory dir) throws IOException {
+ for (String file : dir.listAll()) {
+ assertFalse("a temporary file outlived the merge: " + file, file.endsWith(".tmp"));
+ }
+ }
+
+ /** Returns the same leaf with every field assigned a number unused by its segment. */
+ private static CodecReader renumbered(CodecReader reader) {
+ List<FieldInfo> renumbered = new ArrayList<>();
+ for (FieldInfo info : reader.getFieldInfos()) {
+ renumbered.add(
+ new FieldInfo(
+ info.name,
+ info.number + 100,
+ info.hasTermVectors(),
+ info.omitsNorms(),
+ info.hasPayloads(),
+ info.getIndexOptions(),
+ info.getDocValuesType(),
+ info.docValuesSkipIndexType(),
+ info.getDocValuesGen(),
+ info.attributes(),
+ info.getPointDimensionCount(),
+ info.getPointIndexDimensionCount(),
+ info.getPointNumBytes(),
+ info.getVectorDimension(),
+ info.getVectorEncoding(),
+ info.getVectorSimilarityFunction(),
+ info.isSoftDeletesField(),
+ info.isParentField()));
+ }
+ FieldInfos fieldInfos = new FieldInfos(renumbered.toArray(new FieldInfo[0]));
+ return new FilterCodecReader(reader) {
+ @Override
+ public FieldInfos getFieldInfos() {
+ return fieldInfos;
+ }
+
+ @Override
+ public CacheHelper getCoreCacheHelper() {
+ return null;
+ }
+
+ @Override
+ public CacheHelper getReaderCacheHelper() {
+ return null;
+ }
+ };
+ }
+
+ private record MergeCounts(long mergedRawBytesRead, List<String> handOffs) {}
+
+ private MergeCounts runMerge(
+ KnnVectorsFormat format,
+ VectorSimilarityFunction similarity,
+ boolean expectGraph,
+ boolean hnsw)
+ throws IOException {
+ return runMerge(format, similarity, expectGraph, hnsw, DOCS_PER_SEGMENT);
+ }
+
+ /**
+ * Indexes two segments, merges them, and reports what the merge read and what it created.
+ *
+ * @param expectGraph whether the merge is expected to build a graph
+ * @param hnsw whether the format is an HNSW format at all (a flat format builds no graph and has
+ * no reader to ask about one)
+ */
+ private MergeCounts runMerge(
+ KnnVectorsFormat format,
+ VectorSimilarityFunction similarity,
+ boolean expectGraph,
+ boolean hnsw,
+ int docsPerSegment)
+ throws IOException {
+ float[][] vectors = randomVectors(2 * docsPerSegment, similarity);
+ try (ReadCountingDirectory dir = new ReadCountingDirectory(newDirectory())) {
+ IndexWriterConfig config =
+ new IndexWriterConfig()
+ .setCodec(TestUtil.alwaysKnnVectorsFormat(format))
+ // keep the raw vector file a file of its own rather than a region of a .cfs
+ .setUseCompoundFile(false);
+ config.getMergePolicy().setNoCFSRatio(0.0);
+ try (IndexWriter writer = new IndexWriter(dir, config)) {
+ for (int i = 0; i < vectors.length; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField("v", vectors[i], similarity));
+ writer.addDocument(doc);
+ if (i == docsPerSegment - 1) {
+ writer.commit();
+ }
+ }
+ writer.commit();
+ SegmentInfos beforeMerge = SegmentInfos.readLatestCommit(dir);
+ assertEquals("expected two segments to merge", 2, beforeMerge.size());
+ long sourceRawBytes = 0;
+ for (SegmentCommitInfo info : beforeMerge) {
+ for (String file : info.files()) {
+ if (file.endsWith(".vec")) {
+ sourceRawBytes += dir.fileLength(file);
+ }
+ }
+ }
+ dir.resetCounts();
+
+ writer.forceMerge(1);
+ writer.commit();
+
+ SegmentInfos afterMerge = SegmentInfos.readLatestCommit(dir);
+ assertEquals("expected one segment after the merge", 1, afterMerge.size());
+ SegmentCommitInfo merged = afterMerge.info(0);
+ String mergedRaw = null;
+ for (String file : merged.files()) {
+ assertFalse("a temporary file became a segment file: " + file, file.endsWith(".tmp"));
+ if (file.endsWith(".vec")) {
+ mergedRaw = file;
+ }
+ }
+ assertNotNull("the merged segment has no raw vector file", mergedRaw);
+ assertNoTempFiles(dir);
+
+ long sourceRead = 0;
+ for (SegmentCommitInfo info : beforeMerge) {
+ for (String file : info.files()) {
+ if (file.endsWith(".vec")) {
+ sourceRead += dir.bytesRead(file);
+ }
+ }
+ }
+ // positive control for the counter: the merge did read the source vectors, in full
+ assertTrue(
+ "the byte counter saw no read of the source vectors: " + sourceRead,
+ sourceRead >= sourceRawBytes);
+
+ if (hnsw) {
+ try (DirectoryReader reader = DirectoryReader.open(writer)) {
+ KnnVectorsReader vectorsReader = fieldVectorsReader(reader, "v");
+ HnswGraph graph = ((HnswGraphProvider) vectorsReader).getGraph("v");
+ assertEquals("graph presence", expectGraph, graph != null && graph.size() > 0);
+ }
+ }
+ return new MergeCounts(dir.bytesRead(mergedRaw), dir.handOffsCreated());
+ }
+ }
+ }
+
+ private float[][] randomVectors(int count, VectorSimilarityFunction similarity) {
+ float[][] vectors = new float[count][];
+ for (int i = 0; i < count; i++) {
+ float[] vector = new float[DIM];
+ for (int j = 0; j < DIM; j++) {
+ vector[j] = random().nextFloat() * 2 - 1;
+ }
+ if (similarity == VectorSimilarityFunction.DOT_PRODUCT) {
+ // DOT_PRODUCT requires unit vectors; leave COSINE non-unit to exercise normalization.
+ VectorUtil.l2normalize(vector);
+ } else if (similarity == VectorSimilarityFunction.COSINE) {
+ for (int j = 0; j < DIM; j++) {
+ vector[j] *= 5f;
+ }
+ }
+ vectors[i] = vector;
+ }
+ return vectors;
+ }
+
+ /** Counts every byte read from every file, through clones and slices, and names files created. */
+ private static class ReadCountingDirectory extends FilterDirectory {
+ private final Map<String, Long> bytesRead = new ConcurrentHashMap<>();
+ private final Set<String> created = ConcurrentHashMap.newKeySet();
+
+ ReadCountingDirectory(Directory in) {
+ super(in);
+ }
+
+ @Override
+ public IndexInput openInput(String name, IOContext context) throws IOException {
+ return new CountingIndexInput(
+ super.openInput(name, context), read -> bytesRead.merge(name, read, Long::sum));
+ }
+
+ @Override
+ public IndexOutput createOutput(String name, IOContext context) throws IOException {
+ created.add(name);
+ return super.createOutput(name, context);
+ }
+
+ @Override
+ public IndexOutput createTempOutput(String prefix, String suffix, IOContext context)
+ throws IOException {
+ IndexOutput output = super.createTempOutput(prefix, suffix, context);
+ created.add(output.getName());
+ return output;
+ }
+
+ void resetCounts() {
+ bytesRead.clear();
+ created.clear();
+ }
+
+ long bytesRead(String name) {
+ return bytesRead.getOrDefault(name, 0L);
+ }
+
+ /** Returns query-data files created by either merge-scorer path. */
+ List<String> handOffsCreated() {
+ return created.stream().filter(name -> name.contains("queries")).sorted().toList();
+ }
+ }
+
+ /** Reports every byte read through it, and through everything cloned or sliced off it. */
+ private static class CountingIndexInput extends FilterIndexInput {
+ private final LongConsumer bytesRead;
+
+ CountingIndexInput(IndexInput in, LongConsumer bytesRead) {
+ super(in.toString(), in);
+ this.bytesRead = bytesRead;
+ }
+
+ @Override
+ public IndexInput clone() {
+ return new CountingIndexInput(in.clone(), bytesRead);
+ }
+
+ @Override
+ public IndexInput slice(String sliceDescription, long offset, long length) throws IOException {
+ return new CountingIndexInput(in.slice(sliceDescription, offset, length), bytesRead);
+ }
+
+ @Override
+ public byte readByte() throws IOException {
+ bytesRead.accept(1);
+ return in.readByte();
+ }
+
+ @Override
+ public void readBytes(byte[] b, int offset, int len) throws IOException {
+ bytesRead.accept(len);
+ in.readBytes(b, offset, len);
+ }
+ }
+
+ /**
+ * An HNSW format backed by a caller-supplied flat format. It writes Lucene104 bytes and reuses
+ * the registered name so test segments reopen through SPI with the shipped format.
+ */
+ private static final class HnswOverFlatFormat extends KnnVectorsFormat {
+ private static final KnnVectorsFormat READ_FORMAT =
+ new Lucene104HnswScalarQuantizedVectorsFormat();
+
+ private final FlatVectorsFormat flatVectorsFormat;
+ private final int tinySegmentsThreshold;
+
+ HnswOverFlatFormat(FlatVectorsFormat flatVectorsFormat, int tinySegmentsThreshold) {
+ super(Lucene104HnswScalarQuantizedVectorsFormat.NAME);
+ this.flatVectorsFormat = flatVectorsFormat;
+ this.tinySegmentsThreshold = tinySegmentsThreshold;
+ }
+
+ @Override
+ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+ return new Lucene99HnswVectorsWriter(
+ state,
+ MAX_CONN,
+ BEAM_WIDTH,
+ flatVectorsFormat,
+ flatVectorsFormat.fieldsWriter(state),
+ 1,
+ null,
+ tinySegmentsThreshold);
+ }
+
+ @Override
+ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
+ return READ_FORMAT.fieldsReader(state);
+ }
+
+ @Override
+ public int getMaxDimensions(String fieldName) {
+ return 1024;
+ }
+ }
+
+ /** Lucene104 flat format whose writer uses the default no-preparation merge path. */
+ private static final class NoPrepareFlatFormat extends Lucene104ScalarQuantizedVectorsFormat {
+ private final ScalarEncoding encoding;
+
+ NoPrepareFlatFormat(ScalarEncoding encoding) {
+ super(encoding);
+ this.encoding = encoding;
+ }
+
+ @Override
+ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+ return new Lucene104ScalarQuantizedVectorsWriter(
+ state, encoding, RAW_FORMAT.fieldsWriter(state), QUANTIZED_SCORER) {
+ @Override
+ public MergeScorerData mergeOneFlatVectorFieldForMergeScorer(
+ FieldInfo fieldInfo, MergeState mergeState, IntPredicate needsMergeScorer)
+ throws IOException {
+ mergeOneFlatVectorField(fieldInfo, mergeState);
+ return null;
+ }
+ };
+ }
+ }
+
+ /** The shipped flat format, keeping every hand-off it prepares for the test to inspect. */
+ private static final class CapturingFlatFormat extends Lucene104ScalarQuantizedVectorsFormat {
+ private final ScalarEncoding encoding;
+ private final List<MergeScorerData> handles;
+
+ CapturingFlatFormat(ScalarEncoding encoding, List<MergeScorerData> handles) {
+ super(encoding);
+ this.encoding = encoding;
+ this.handles = handles;
+ }
+
+ @Override
+ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+ return new Lucene104ScalarQuantizedVectorsWriter(
+ state, encoding, RAW_FORMAT.fieldsWriter(state), QUANTIZED_SCORER) {
+ @Override
+ public MergeScorerData mergeOneFlatVectorFieldForMergeScorer(
+ FieldInfo fieldInfo, MergeState mergeState, IntPredicate needsMergeScorer)
+ throws IOException {
+ MergeScorerData handle =
+ super.mergeOneFlatVectorFieldForMergeScorer(fieldInfo, mergeState, needsMergeScorer);
+ if (handle != null) {
+ handles.add(handle);
+ }
+ return handle;
+ }
+ };
+ }
+ }
+
+ /** The per-field reader for {@code field} of the only segment. */
+ private static KnnVectorsReader fieldVectorsReader(DirectoryReader reader, String field) {
+ KnnVectorsReader vectorsReader = ((CodecReader) getOnlyLeafReader(reader)).getVectorReader();
+ if (vectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
+ vectorsReader = fieldsReader.getFieldReader(field);
+ }
+ return vectorsReader;
+ }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java
index ba1e95f05df..0d2bdf487f9 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java
@@ -16,6 +16,7 @@
*/
package org.apache.lucene.index;
+import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
@@ -26,6 +27,7 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.lucene104.Lucene104HnswScalarQuantizedVectorsFormat;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
@@ -37,6 +39,7 @@ import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.InfoStream;
import org.apache.lucene.util.NamedThreadFactory;
import org.apache.lucene.util.ThreadInterruptedException;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
/**
* Tests that aborting a merge (e.g. via {@link IndexWriter#rollback()}) interrupts HNSW graph
@@ -99,6 +102,23 @@ public class TestHnswMergeAbort extends LuceneTestCase {
}
}
+ /**
+ * Verifies that aborting an asymmetric scalar merge releases the query-side temporary file
+ * prepared for its graph scorer.
+ */
+ public void testRollbackDuringQuantizedMerge() throws Exception {
+ doTestRollbackDuringMerge(
+ true,
+ new Lucene104HnswScalarQuantizedVectorsFormat(
+ ScalarEncoding.SINGLE_BIT_QUERY_NIBBLE,
+ 16,
+ BEAM_WIDTH,
+ 1,
+ null,
+ TINY_SEGMENTS_THRESHOLD),
+ "build graph from " + LIVE_DOCS_AFTER_DELETES + " vectors");
+ }
+
private void doTestRollbackDuringMerge(
boolean withDeletes, KnnVectorsFormat format, String expectedBuildStart) throws Exception {
try (Directory dir = newDirectory()) {
@@ -183,9 +203,19 @@ public class TestHnswMergeAbort extends LuceneTestCase {
}
});
merger.start();
+ List<String> midMerge = List.of();
try {
assertTrue(
"HNSW graph construction never started", buildStarted.await(120, TimeUnit.SECONDS));
+ // Capture the handoff while the graph build is parked. If no tragedy is recorded, rollback
+ // can delete an unreferenced temporary file, so the final listing alone cannot prove
+ // cleanup. MockDirectoryWrapper also fails on close if the handoff remains open.
+ midMerge = tempFiles(dir);
+ if (format instanceof Lucene104HnswScalarQuantizedVectorsFormat) {
+ assertFalse(
+ "the quantized merge wrote no hand-off file for the abort to release",
+ midMerge.isEmpty());
+ }
w2.rollback();
assertTrue("merge thread was not released by the abort signal", releasedAfterAbort.get());
synchronized (hnswMessages) {
@@ -210,6 +240,20 @@ public class TestHnswMergeAbort extends LuceneTestCase {
merger.join(TimeUnit.MINUTES.toMillis(5));
assertFalse("merge thread did not terminate", merger.isAlive());
}
+ assertEquals(
+ "a temporary file outlived the aborted merge, which had " + midMerge + " mid-build",
+ List.of(),
+ tempFiles(dir));
+ }
+ }
+
+ private static List<String> tempFiles(Directory dir) throws IOException {
+ List<String> temps = new ArrayList<>();
+ for (String file : dir.listAll()) {
+ if (file.endsWith(".tmp")) {
+ temps.add(file);
+ }
}
+ return temps;
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java b/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java
index 510008a4b31..6edf05b4e12 100644
--- a/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java
+++ b/lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java
@@ -179,6 +179,75 @@ public class TestOptimizedScalarQuantizer extends LuceneTestCase {
}
}
+ /**
+ * Verifies that {@code multiScalarQuantize} matches independent {@code scalarQuantize} calls at
+ * both bit widths. The merge path depends on byte-identical index-side and query-side records.
+ */
+ public void testMultiScalarQuantizeMatchesScalarQuantize() {
+ byte[][] bitPairs = new byte[][] {{1, 4}, {2, 4}, {4, 8}};
+ for (int trial = 0; trial < 25; trial++) {
+ int dims = randomIntBetween(1, 1024);
+ float[] vector = new float[dims];
+ float[] centroid = new float[dims];
+ for (int i = 0; i < dims; ++i) {
+ vector[i] = randomFloat();
+ centroid[i] = randomFloat();
+ }
+ for (VectorSimilarityFunction similarityFunction : VectorSimilarityFunction.values()) {
+ float[] scratch = new float[dims];
+ System.arraycopy(vector, 0, scratch, 0, dims);
+ float[] centroidCopy = new float[dims];
+ System.arraycopy(centroid, 0, centroidCopy, 0, dims);
+ if (similarityFunction == VectorSimilarityFunction.COSINE) {
+ VectorUtil.l2normalize(scratch);
+ VectorUtil.l2normalize(centroidCopy);
+ }
+ OptimizedScalarQuantizer osq = new OptimizedScalarQuantizer(similarityFunction);
+ for (byte[] bits : bitPairs) {
+ byte[] multiLow = new byte[dims];
+ byte[] multiHigh = new byte[dims];
+ OptimizedScalarQuantizer.QuantizationResult[] multi =
+ osq.multiScalarQuantize(
+ scratch.clone(), new byte[][] {multiLow, multiHigh}, bits, centroidCopy);
+ byte[] singleLow = new byte[dims];
+ byte[] singleHigh = new byte[dims];
+ // each single-width call gets its own copy: both center the vector in place
+ OptimizedScalarQuantizer.QuantizationResult low =
+ osq.scalarQuantize(scratch.clone(), singleLow, bits[0], centroidCopy);
+ OptimizedScalarQuantizer.QuantizationResult high =
+ osq.scalarQuantize(scratch.clone(), singleHigh, bits[1], centroidCopy);
+ String where = similarityFunction + " dims=" + dims + " bits=" + bits[0] + "/" + bits[1];
+ assertArrayEquals(where, singleLow, multiLow);
+ assertArrayEquals(where, singleHigh, multiHigh);
+ assertIdenticalResults(where, low, multi[0]);
+ assertIdenticalResults(where, high, multi[1]);
+ }
+ }
+ }
+ }
+
+ private static void assertIdenticalResults(
+ String where,
+ OptimizedScalarQuantizer.QuantizationResult expected,
+ OptimizedScalarQuantizer.QuantizationResult actual) {
+ assertEquals(
+ where + " lowerInterval",
+ Float.floatToRawIntBits(expected.lowerInterval()),
+ Float.floatToRawIntBits(actual.lowerInterval()));
+ assertEquals(
+ where + " upperInterval",
+ Float.floatToRawIntBits(expected.upperInterval()),
+ Float.floatToRawIntBits(actual.upperInterval()));
+ assertEquals(
+ where + " additionalCorrection",
+ Float.floatToRawIntBits(expected.additionalCorrection()),
+ Float.floatToRawIntBits(actual.additionalCorrection()));
+ assertEquals(
+ where + " quantizedComponentSum",
+ expected.quantizedComponentSum(),
+ actual.quantizedComponentSum());
+ }
+
public void testUnpackBinary() {
int dim = randomIntBetween(1, 4096);
QuantizedByteVectorValues.ScalarEncoding encoding =