thecoop commented on code in PR #16696:
URL: https://github.com/apache/lucene/pull/16696#discussion_r4073319980
##########
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,
Review Comment:
Isn't the scorer already this type (or it could be wrapped I suppose...)
--
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]
