mccullocht commented on code in PR #16030:
URL: https://github.com/apache/lucene/pull/16030#discussion_r4040014579
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -320,39 +346,129 @@ private QuantizedByteVectorValues
mergedQuantizedVectorValues(
return new QuantizedFloatVectorValues(vectorValues, quantizer, encoding,
centroid);
}
+ /**
+ * Returns a view that quantizes a single segment's float vectors against
{@code centroid} using
+ * this writer's encoding, without consulting any quantized bytes the
segment may already store.
+ */
+ private QuantizedFloatVectorValues quantizeFromFloats(
+ KnnVectorsReader reader, FieldInfo fieldInfo, float[] centroid) throws
IOException {
+ OptimizedScalarQuantizer quantizer =
+ new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
+ FloatVectorValues vectorValues =
+ fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT16
+ ? new
Float16AsFloatVectorValues(reader.getFloat16VectorValues(fieldInfo.name))
+ : reader.getFloatVectorValues(fieldInfo.name);
+ if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
+ vectorValues = new NormalizedFloatVectorValues(vectorValues);
+ }
+ return new QuantizedFloatVectorValues(vectorValues, quantizer, encoding,
centroid);
+ }
+
@Override
public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState
mergeState)
throws IOException {
- // Don't need access to the random vectors, we can just use the merged
- rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
if (fieldInfo.getVectorEncoding().isFloatingPoint() == false) {
+ rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
return;
}
- final float[] centroid;
+ if (enableCentering) {
+ mergeOneFlatVectorFieldCentered(fieldInfo, mergeState);
+ } else {
+ mergeOneFlatVectorFieldDataBlind(fieldInfo, mergeState);
+ }
+ }
+
+ private void mergeOneFlatVectorFieldCentered(FieldInfo fieldInfo, MergeState
mergeState)
+ throws IOException {
+ // Don't need access to the random vectors, we can just use the merged
+ rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
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);
}
QuantizedByteVectorValues quantizedVectorValues =
- mergedQuantizedVectorValues(fieldInfo, mergeState, centroid);
+ mergedQuantizedVectorValues(fieldInfo, mergeState, mergedCentroid);
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;
+ docsWithField.cardinality() > 0 ?
VectorUtil.dotProduct(mergedCentroid, mergedCentroid) : 0;
writeMeta(
fieldInfo,
segmentWriteState.segmentInfo.maxDoc(),
vectorDataOffset,
vectorDataLength,
- centroid,
+ mergedCentroid,
centroidDp,
docsWithField);
}
+ private void mergeOneFlatVectorFieldDataBlind(FieldInfo fieldInfo,
MergeState mergeState)
+ throws IOException {
+ float[] zeroCentroid = new float[fieldInfo.getVectorDimension()];
+ // Build one merged view where, per contributing segment, either its
existing quantized bytes
+ // are passed through or its float vectors are quantized fresh. Inputs
already quantized to
+ // {@code encoding} against a zero centroid (data-blind segments) are
copied directly; they are
+ // never dequantized and re-quantized, which would only add loss. Segments
with raw floats are
+ // quantized fresh, as their stored bytes live in a different (centered)
quantization space.
+ List<QuantizedByteVectorValuesSub> subs = new ArrayList<>();
+ for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
+ KnnVectorsReader reader = mergeState.knnVectorsReaders[i];
+ if (reader == null) {
+ continue;
+ }
+ QuantizedByteVectorValues values;
+ if (hasRawVectorValues(reader, fieldInfo)) {
+ // Segment stored full-precision floats; quantize them against the
zero centroid.
+ values = quantizeFromFloats(reader, fieldInfo, zeroCentroid);
+ } else {
+ QuantizedByteVectorValues qvv = getQuantizedVectorValues(reader,
fieldInfo.name);
+ if (qvv == null || qvv.size() == 0) {
+ continue;
+ }
+ if (qvv.getScalarEncoding() != encoding) {
+ // Re-quantization from raw floats would be required, which is not
possible when raw
+ // floats were never written.
+ throw new IllegalStateException(
Review Comment:
Changed it so that we disallow merging back into centered vectors if the
float vectors are missing.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -456,8 +610,10 @@ static int mergeAndRecalculateCentroids(
float[] centroid = getCentroid(knnVectorsReader, fieldInfo.name);
totalVectorCount += vectorCount;
// If there aren't centroids, or previously clustered with more than one
cluster
- // or if there are deleted docs, we must recalculate the centroid
- if (centroid == null || mergeState.liveDocs[i] != null) {
+ // or if there are deleted docs, we must recalculate the centroid. An
all-zero centroid
+ // indicates a data-blind segment (no centering was done); it can't be
combined with the
+ // others, so recompute from the (possibly dequantized) vectors.
+ if (centroid == null || isAllZero(centroid) || mergeState.liveDocs[i] !=
null) {
Review Comment:
Removed isAllZero() and replaced it with a mode check; isAllZero was a
vestige of the previous approach.
##########
lucene/CHANGES.txt:
##########
@@ -393,6 +393,10 @@ New Features
translates each term into a SpanTermQuery within a {@link SpanOrQuery}, but
retains only the most frequent terms so it
will not overflow the boolean max clause count. (Alan Woodward, Nishant
Mehta, David Smiley)
+* GITHUB#16029: Scalar quantization option to disable centering and writing of
float vectors. This
+ reduces vector storage costs by 4x or more but also reduces quantization
accuracy.
+ (Trevor McCulloch)
Review Comment:
Done
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -740,6 +936,148 @@ long quantizationOverheadBytesUsed() {
}
}
+ /** In-memory storage for fp32 vectors used in data-blind mode; nothing is
written to disk. */
+ private static class InMemoryFloatFieldWriter extends
FlatFieldVectorsWriter<float[]> {
Review Comment:
Collapsed the implementations using generics.
I don't think the duplication is a huge issue, at least not yet, so I'm not
going to expose this the Lucene99 writers (it would need to be public)
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java:
##########
@@ -85,19 +89,27 @@
* <li><b>vlong</b> the length of the vector data in the .veq file
* <li><b>vint</b> the number of vectors
* <li><b>vint</b> the wire number for ScalarEncoding
- * <li><b>[float]</b> the centroid
- * <li><b>float</b> the centroid square magnitude
+ * <li><b>[float]</b> the centroid (omitted when the metadata version
indicates data-blind mode)
+ * <li><b>float</b> the centroid square magnitude (omitted when the metadata
version indicates
+ * data-blind mode)
* <li>The sparse vector information, if required, mapping vector ordinal to
doc ID
* </ul>
*
+ * <p>{@code enableCentering} manifests in the version: when true we write
version 0 and when false
Review Comment:
Yes, this is to keep this backwards compatible.
This is definitely unconventional -- typically this kind of encoding change
is a one way door even if it's just a VERSION bump within the codec. TBH this
better matches my priors -- a consumer of Lucene could upgrade, wait for new
version of Lucene to pass their rollback window, then allow disable centering.
IIRC in a previous review iteration I was asked to push everything back into
Lucene104 so I'd prefer to leave this as-is. If the changes get much more
involved we should create a new codec. I think next time I'd consider being
more mechanical: a PR that simply branches the existing format, then a second
PR that makes changes to make review tractable.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloatVectorValues.java:
##########
@@ -93,7 +93,7 @@ abstract class OffHeapScalarQuantizedFloatVectorValues
extends FloatVectorValues
this.byteBuffer = ByteBuffer.allocate(docPackedLength);
this.vectorValue = new float[dimension];
this.byteValue = byteBuffer.array();
- this.unpackedByteVectorValue = new byte[dimension];
+ this.unpackedByteVectorValue = new
byte[encoding.getDiscreteDimensions(dimension)];
Review Comment:
This is an independent change -- it would occasionally compute the wrong
size and tests would fail.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java:
##########
@@ -44,6 +44,10 @@
* quantized vectors in the index.
* <li>Transforming the half-byte quantized query vectors in such a way that
the comparison with
* single bit vectors can be done with bit arithmetic.
+ * <li>Data blind mode: vectors are quantized without centering and float
vectors are discarded.
Review Comment:
Switched the parameter to an `enum Mode` that flips between the existing
centered approach, and data blind with/without source vectors. The argument for
allowing you to keep float vectors is low fidelity (1, 2, maybe 4) bit
quantization where you may wish to use the float vectors to re-rank the results
of the ANN search.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -320,39 +346,129 @@ private QuantizedByteVectorValues
mergedQuantizedVectorValues(
return new QuantizedFloatVectorValues(vectorValues, quantizer, encoding,
centroid);
}
+ /**
+ * Returns a view that quantizes a single segment's float vectors against
{@code centroid} using
+ * this writer's encoding, without consulting any quantized bytes the
segment may already store.
+ */
+ private QuantizedFloatVectorValues quantizeFromFloats(
+ KnnVectorsReader reader, FieldInfo fieldInfo, float[] centroid) throws
IOException {
+ OptimizedScalarQuantizer quantizer =
+ new OptimizedScalarQuantizer(fieldInfo.getVectorSimilarityFunction());
+ FloatVectorValues vectorValues =
+ fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT16
+ ? new
Float16AsFloatVectorValues(reader.getFloat16VectorValues(fieldInfo.name))
+ : reader.getFloatVectorValues(fieldInfo.name);
+ if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
+ vectorValues = new NormalizedFloatVectorValues(vectorValues);
+ }
+ return new QuantizedFloatVectorValues(vectorValues, quantizer, encoding,
centroid);
+ }
+
@Override
public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState
mergeState)
throws IOException {
- // Don't need access to the random vectors, we can just use the merged
- rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
if (fieldInfo.getVectorEncoding().isFloatingPoint() == false) {
+ rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
return;
}
- final float[] centroid;
+ if (enableCentering) {
+ mergeOneFlatVectorFieldCentered(fieldInfo, mergeState);
+ } else {
+ mergeOneFlatVectorFieldDataBlind(fieldInfo, mergeState);
+ }
+ }
+
+ private void mergeOneFlatVectorFieldCentered(FieldInfo fieldInfo, MergeState
mergeState)
+ throws IOException {
+ // Don't need access to the random vectors, we can just use the merged
+ rawVectorDelegate.mergeOneFlatVectorField(fieldInfo, mergeState);
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);
}
QuantizedByteVectorValues quantizedVectorValues =
- mergedQuantizedVectorValues(fieldInfo, mergeState, centroid);
+ mergedQuantizedVectorValues(fieldInfo, mergeState, mergedCentroid);
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;
+ docsWithField.cardinality() > 0 ?
VectorUtil.dotProduct(mergedCentroid, mergedCentroid) : 0;
writeMeta(
fieldInfo,
segmentWriteState.segmentInfo.maxDoc(),
vectorDataOffset,
vectorDataLength,
- centroid,
+ mergedCentroid,
centroidDp,
docsWithField);
}
+ private void mergeOneFlatVectorFieldDataBlind(FieldInfo fieldInfo,
MergeState mergeState)
+ throws IOException {
+ float[] zeroCentroid = new float[fieldInfo.getVectorDimension()];
+ // Build one merged view where, per contributing segment, either its
existing quantized bytes
+ // are passed through or its float vectors are quantized fresh. Inputs
already quantized to
+ // {@code encoding} against a zero centroid (data-blind segments) are
copied directly; they are
+ // never dequantized and re-quantized, which would only add loss. Segments
with raw floats are
+ // quantized fresh, as their stored bytes live in a different (centered)
quantization space.
+ List<QuantizedByteVectorValuesSub> subs = new ArrayList<>();
+ for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
+ KnnVectorsReader reader = mergeState.knnVectorsReaders[i];
+ if (reader == null) {
+ continue;
+ }
+ QuantizedByteVectorValues values;
+ if (hasRawVectorValues(reader, fieldInfo)) {
+ // Segment stored full-precision floats; quantize them against the
zero centroid.
+ values = quantizeFromFloats(reader, fieldInfo, zeroCentroid);
+ } else {
+ QuantizedByteVectorValues qvv = getQuantizedVectorValues(reader,
fieldInfo.name);
+ if (qvv == null || qvv.size() == 0) {
+ continue;
+ }
+ if (qvv.getScalarEncoding() != encoding) {
+ // Re-quantization from raw floats would be required, which is not
possible when raw
+ // floats were never written.
+ throw new IllegalStateException(
+ "Cannot merge field \""
+ + fieldInfo.name
+ + "\" from data-blind segment with encoding "
+ + qvv.getScalarEncoding()
+ + " into data-blind format with encoding "
+ + encoding
+ + ": re-quantization requires raw float vectors");
+ }
+ float[] centroid = getCentroid(reader, fieldInfo.name);
+ if (centroid != null && isAllZero(centroid)) {
+ // Quantized-only segment whose bytes already match the output
format (encoding and zero
+ // centroid): copy them directly.
+ values = qvv;
+ } else {
+ // Bytes were produced against a non-zero (or unknown) centroid, so
they cannot be passed
+ // through into the zero-centroid output; re-quantize from floats.
+ values = quantizeFromFloats(reader, fieldInfo, zeroCentroid);
+ }
+ }
+ subs.add(new QuantizedByteVectorValuesSub(mergeState.docMaps[i],
values));
+ }
+ long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES);
+ MergedQuantizedByteVectorValues mergedQBVV =
+ MergedQuantizedByteVectorValues.merge(mergeState, zeroCentroid,
encoding, subs);
+ DocsWithFieldSet docsWithField = writeVectorData(vectorData, mergedQBVV);
+ long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset;
+ // centroidDp is 0 (zero centroid); the data-blind metadata omits it and
the centroid.
+ writeMeta(
+ fieldInfo,
+ segmentWriteState.segmentInfo.maxDoc(),
+ vectorDataOffset,
+ vectorDataLength,
+ zeroCentroid,
Review Comment:
Done
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -420,6 +536,44 @@ static float[] getCentroid(KnnVectorsReader vectorsReader,
String fieldName) {
return null;
}
+ static QuantizedByteVectorValues getQuantizedVectorValues(
+ KnnVectorsReader vectorsReader, String fieldName) throws IOException {
+ vectorsReader = vectorsReader.unwrapReaderForField(fieldName);
+ if (vectorsReader instanceof Lucene104ScalarQuantizedVectorsReader reader)
{
+ return reader.getQuantizedVectorValues(fieldName);
+ }
+ return null;
+ }
+
+ /**
+ * Returns whether the segment stores full-precision vectors for this field,
or null when the
Review Comment:
Done
--
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]
