john-mlika commented on code in PR #16696:
URL: https://github.com/apache/lucene/pull/16696#discussion_r4077889055
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -432,21 +459,168 @@ private void mergeOneFlatVectorFieldCentered(FieldInfo
fieldInfo, MergeState mer
segmentWriteState.infoStream.message(
QUANTIZED_VECTOR_COMPONENT, "Vectors' count:" + vectorCount);
}
- QuantizedByteVectorValues quantizedVectorValues =
- mergedQuantizedVectorValues(fieldInfo, mergeState, mergedCentroid);
+ // Only asymmetric encodings have query-side records, and only FLOAT32
fields score graphs with
+ // them. FLOAT16 fields use fp16 vectors. 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()
+ && fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT32
+ && 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(mergedCentroid, mergedCentroid) : 0;
- writeMeta(
+ 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;
+ }
+ }
+
+ /**
+ * 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",
segmentWriteState.context)) {
+ 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]);
+ // the query side, packed as Lucene104ScalarQuantizedVectorsReader
packs it
+ transposeHalfByte(queryQuantized, queryPacked);
+ queryData.writeBytes(queryPacked, queryPacked.length);
+ writeCorrections(queryData, corrections[1]);
+ docsWithField.add(docV);
+ }
+ CodecUtil.writeFooter(queryData);
+ } catch (Throwable t) {
+ if (queryDataName != null) {
+ IOUtils.deleteFilesSuppressingExceptions(t,
segmentWriteState.directory, queryDataName);
+ }
+ throw t;
+ }
+ return new PreparedQueryData(
+ segmentWriteState.directory,
+ segmentWriteState.context,
fieldInfo,
- segmentWriteState.segmentInfo.maxDoc(),
- vectorDataOffset,
- vectorDataLength,
- mergedCentroid,
- centroidDp,
- docsWithField);
+ (Lucene104ScalarQuantizedVectorScorer) vectorsScorer,
+ 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) {
+ // Accept only this format's reader type. A wrapper around that
reader must use the
+ // fallback, even over the correct segment.
Review Comment:
i was concerned about wrappers whose ordinals don't match the reader they
wrap, but i guess those shouldn't unwrap to that reader in the first place
added an unwrap through KnnVectorsReader#unwrapReaderForField before the
type check as well as a related test
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
