msokolov commented on code in PR #16473:
URL: https://github.com/apache/lucene/pull/16473#discussion_r3797289598
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java:
##########
@@ -111,6 +111,43 @@ public RandomVectorScorer getRandomVectorScorer(
public RandomVectorScorer getRandomVectorScorer(
VectorSimilarityFunction similarityFunction, KnnVectorValues
vectorValues, short[] target)
throws IOException {
+ if (vectorValues instanceof QuantizedByteVectorValues qv) {
+ FlatVectorsScorer.checkDimensions(target.length, qv.dimension());
+ OptimizedScalarQuantizer quantizer = qv.getQuantizer();
+ ScalarEncoding scalarEncoding = qv.getScalarEncoding();
+ byte[] scratch = new
byte[scalarEncoding.getDiscreteDimensions(qv.dimension())];
+ final byte[] targetQuantized;
+ if (scalarEncoding.isAsymmetric() == false) {
+ targetQuantized = scratch;
+ } else {
+ // This is asymmetric quantization, we will pack the vector
+ targetQuantized = new
byte[scalarEncoding.getQueryPackedLength(scratch.length)];
+ }
+ // Inflate the fp16 query to fp32 and normalize there; quantization
operates on fp32.
Review Comment:
Do we want to add a TODO: linking to an issue for implementing quantization
directly over fp16?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsReader.java:
##########
@@ -224,7 +224,26 @@ public RandomVectorScorer getRandomVectorScorer(String
field, byte[] target) thr
@Override
public RandomVectorScorer getRandomVectorScorer(String field, short[]
target) throws IOException {
- return rawVectorsReader.getRandomVectorScorer(field, target);
+ FieldEntry fi = fields.get(field);
Review Comment:
again, can we DRY this up?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java:
##########
@@ -0,0 +1,398 @@
+/*
+ * 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 static
org.apache.lucene.util.quantization.OptimizedScalarQuantizer.deQuantize;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.packed.DirectMonotonicReader;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Reads quantized vector values from the index input and returns float16
vector values after
+ * dequantizing them.
+ *
+ * <p>Used for read-only indexes whose raw float16 vectors have been dropped
to save storage: only
+ * the scalar-quantized bytes remain, so {@link #vectorValue(int)}
reconstructs float16 values by
+ * dequantizing them, with some precision loss.
+ *
+ * @lucene.internal
+ */
+abstract class OffHeapScalarQuantizedFloat16VectorValues extends
Float16VectorValues
+ implements HasIndexSlice {
+
+ final int dimension;
+ final int size;
+ final VectorSimilarityFunction similarityFunction;
+ final FlatVectorsScorer vectorsScorer;
+
+ final IndexInput slice;
+ final short[] vectorValue;
+ final byte[] byteValue;
+ final ByteBuffer byteBuffer;
+ final byte[] unpackedByteVectorValue;
+ final int byteSize;
+ private int lastOrd = -1;
+ final float[] correctiveValues;
+ int quantizedComponentSum;
+ final ScalarEncoding encoding;
+ final float[] centroid;
+
+ OffHeapScalarQuantizedFloat16VectorValues(
+ int dimension,
+ int size,
+ float[] centroid,
+ ScalarEncoding encoding,
+ VectorSimilarityFunction similarityFunction,
+ FlatVectorsScorer vectorsScorer,
+ IndexInput slice) {
+ this.dimension = dimension;
+ this.size = size;
+ this.similarityFunction = similarityFunction;
+ this.vectorsScorer = vectorsScorer;
+ this.slice = slice;
+ this.centroid = centroid;
+ this.correctiveValues = new float[3];
+ this.encoding = encoding;
+ int docPackedLength = encoding.getDocPackedLength(dimension);
+ this.byteSize = docPackedLength + (Float.BYTES * 3) + Integer.BYTES;
Review Comment:
what is this calculation about?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -613,6 +668,75 @@ public long ramBytesUsed() {
}
}
+ private static class Float32FieldWriter extends FieldWriter<float[]> {
+ private final float[] normalized;
+
+ Float32FieldWriter(
+ FieldInfo fieldInfo, FlatFieldVectorsWriter<float[]>
flatFieldVectorsWriter) {
+ super(fieldInfo, flatFieldVectorsWriter);
+ this.normalized = new float[dim];
+ }
+
+ @Override
+ public void addValue(int docID, float[] vectorValue) throws IOException {
+ flatFieldVectorsWriter.addValue(docID, vectorValue);
+ accumulate(vectorValue);
+ }
+
+ @Override
+ float[] floatVectorValue(int ord) {
+ float[] vector = flatFieldVectorsWriter.getVectors().get(ord);
Review Comment:
I wonder if we should be normalizing "on the way in" -- in `addValue`?
##########
lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java:
##########
@@ -214,6 +264,256 @@ public void testQuantizedVectorsWriteAndRead() throws
IOException {
}
}
+ /**
+ * fp16 counterpart of {@link #testQuantizedVectorsWriteAndRead()}: indexes
float16 vectors and
+ * verifies the persisted quantized bytes + corrective terms match a
reference re-quantization.
+ * The reference mirrors the writer's fp16 path exactly — inflate
fp16->fp32 (normalizing
Review Comment:
wait! I see the dreaded mdash! Did AI write this??!
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java:
##########
@@ -0,0 +1,398 @@
+/*
+ * 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 static
org.apache.lucene.util.quantization.OptimizedScalarQuantizer.deQuantize;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.packed.DirectMonotonicReader;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Reads quantized vector values from the index input and returns float16
vector values after
+ * dequantizing them.
+ *
+ * <p>Used for read-only indexes whose raw float16 vectors have been dropped
to save storage: only
+ * the scalar-quantized bytes remain, so {@link #vectorValue(int)}
reconstructs float16 values by
+ * dequantizing them, with some precision loss.
Review Comment:
the precision loss is relative to the *original* fp16 vectors I guess, not
relative to the quantized vectors. Maybe just add "relative to the original
fp16 vectors" to be explicit
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -613,6 +668,75 @@ public long ramBytesUsed() {
}
}
+ private static class Float32FieldWriter extends FieldWriter<float[]> {
+ private final float[] normalized;
+
+ Float32FieldWriter(
+ FieldInfo fieldInfo, FlatFieldVectorsWriter<float[]>
flatFieldVectorsWriter) {
+ super(fieldInfo, flatFieldVectorsWriter);
+ this.normalized = new float[dim];
+ }
+
+ @Override
+ public void addValue(int docID, float[] vectorValue) throws IOException {
+ flatFieldVectorsWriter.addValue(docID, vectorValue);
+ accumulate(vectorValue);
+ }
+
+ @Override
+ float[] floatVectorValue(int ord) {
+ float[] vector = flatFieldVectorsWriter.getVectors().get(ord);
Review Comment:
I guess I don't really care about stupid `COSINE`
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java:
##########
@@ -111,6 +111,43 @@ public RandomVectorScorer getRandomVectorScorer(
public RandomVectorScorer getRandomVectorScorer(
VectorSimilarityFunction similarityFunction, KnnVectorValues
vectorValues, short[] target)
throws IOException {
+ if (vectorValues instanceof QuantizedByteVectorValues qv) {
+ FlatVectorsScorer.checkDimensions(target.length, qv.dimension());
+ OptimizedScalarQuantizer quantizer = qv.getQuantizer();
+ ScalarEncoding scalarEncoding = qv.getScalarEncoding();
+ byte[] scratch = new
byte[scalarEncoding.getDiscreteDimensions(qv.dimension())];
+ final byte[] targetQuantized;
+ if (scalarEncoding.isAsymmetric() == false) {
+ targetQuantized = scratch;
+ } else {
+ // This is asymmetric quantization, we will pack the vector
+ targetQuantized = new
byte[scalarEncoding.getQueryPackedLength(scratch.length)];
+ }
+ // Inflate the fp16 query to fp32 and normalize there; quantization
operates on fp32.
+ float[] copy = new float[target.length];
Review Comment:
this should all be a copy of the logic above in `float[]` case right? can we
factor out into a utility method?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -131,25 +128,11 @@ public FlatFieldVectorsWriter<?> addField(FieldInfo
fieldInfo) throws IOExceptio
@Override
public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
rawVectorDelegate.flush(maxDoc, sortMap);
- for (FieldWriter field : fields) {
Review Comment:
did this logic get moved? Oh I see .. into `centroid()`.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedVectorValues.java:
##########
@@ -447,6 +472,31 @@ public VectorScorer.Bulk bulk(DocIdSetIterator
matchingDocs) {
}
};
}
+
+ @Override
+ public VectorScorer scorer(short[] target) throws IOException {
+ assert isQuerySide == false;
+ SparseOffHeapVectorValues copy = copy();
+ DocIndexIterator iterator = copy.iterator();
+ RandomVectorScorer scorer =
+ vectorsScorer.getRandomVectorScorer(similarityFunction, copy,
target);
+ return new VectorScorer() {
Review Comment:
we have this pattern in so many places, it might be nice to refactor. But
for another issue
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -523,33 +538,44 @@ public long ramBytesUsed() {
return total;
}
- static class FieldWriter extends FlatFieldVectorsWriter<float[]> {
+ abstract static class FieldWriter<T> extends FlatFieldVectorsWriter<T> {
private static final long SHALLOW_SIZE =
shallowSizeOfInstance(FieldWriter.class);
- private final FieldInfo fieldInfo;
+ protected final FieldInfo fieldInfo;
private boolean finished;
- private final FlatFieldVectorsWriter<float[]> flatFieldVectorsWriter;
+ protected final FlatFieldVectorsWriter<T> flatFieldVectorsWriter;
private final float[] dimensionSums;
private final FloatArrayList magnitudes = new FloatArrayList();
+ protected final int dim;
- FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter<float[]>
flatFieldVectorsWriter) {
+ FieldWriter(FieldInfo fieldInfo, FlatFieldVectorsWriter<T>
flatFieldVectorsWriter) {
this.fieldInfo = fieldInfo;
this.flatFieldVectorsWriter = flatFieldVectorsWriter;
- this.dimensionSums = new float[fieldInfo.getVectorDimension()];
+ this.dim = fieldInfo.getVectorDimension();
+ this.dimensionSums = new float[dim];
+ }
+
+ @SuppressWarnings("unchecked")
+ static FieldWriter<?> create(
+ FieldInfo fieldInfo, FlatFieldVectorsWriter<?> flatFieldVectorsWriter)
{
+ return switch (fieldInfo.getVectorEncoding()) {
+ case BYTE -> throw new UnsupportedOperationException("Byte Vectors
aren't supported");
+ case FLOAT32 ->
+ new Float32FieldWriter(
+ fieldInfo, (FlatFieldVectorsWriter<float[]>)
flatFieldVectorsWriter);
+ case FLOAT16 ->
+ new Float16FieldWriter(
+ fieldInfo, (FlatFieldVectorsWriter<short[]>)
flatFieldVectorsWriter);
+ };
}
@Override
- public List<float[]> getVectors() {
+ public List<T> getVectors() {
return flatFieldVectorsWriter.getVectors();
}
- public void normalizeVectors() {
- for (int i = 0; i < flatFieldVectorsWriter.getVectors().size(); i++) {
- float[] vector = flatFieldVectorsWriter.getVectors().get(i);
- float magnitude = magnitudes.get(i);
- for (int j = 0; j < vector.length; j++) {
- vector[j] /= magnitude;
- }
- }
+ @Override
+ public T copyValue(T vectorValue) {
Review Comment:
did this get added?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -571,26 +597,55 @@ public boolean isFinished() {
return finished && flatFieldVectorsWriter.isFinished();
}
- @Override
- public void addValue(int docID, float[] vectorValue) throws IOException {
- flatFieldVectorsWriter.addValue(docID, vectorValue);
+ /** The ordinal's stored vector as fp32, ready for quantization
(unit-length for COSINE). */
Review Comment:
I think it will also be expected to be unit-length for `DOT_PRODUCT`? It's
just not guaranteed
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsWriter.java:
##########
@@ -428,21 +420,40 @@ static float[] getCentroid(KnnVectorsReader
vectorsReader, String fieldName) {
return null;
}
+ /**
+ * Returns the reader's floating-point vectors viewed as fp32, inflating
fp16 on read, or null
+ * when the field is absent from this reader or is byte-encoded.
+ */
+ private static FloatVectorValues floatingPointVectorValues(
+ KnnVectorsReader reader, FieldInfo fieldInfo) throws IOException {
+ return switch (fieldInfo.getVectorEncoding()) {
+ case FLOAT32 -> reader.getFloatVectorValues(fieldInfo.name);
+ case FLOAT16 -> {
+ Float16VectorValues f16 =
reader.getFloat16VectorValues(fieldInfo.name);
+ yield f16 == null ? null : new Float16AsFloatVectorValues(f16);
+ }
+ case BYTE -> null;
+ };
+ }
+
static int mergeAndRecalculateCentroids(
MergeState mergeState, FieldInfo fieldInfo, float[] mergedCentroid)
throws IOException {
boolean recalculate = false;
int totalVectorCount = 0;
for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i];
- if (knnVectorsReader == null
- || knnVectorsReader.getFloatVectorValues(fieldInfo.name) == null) {
+ if (knnVectorsReader == null) {
continue;
}
- float[] centroid = getCentroid(knnVectorsReader, fieldInfo.name);
- int vectorCount =
knnVectorsReader.getFloatVectorValues(fieldInfo.name).size();
+ KnnVectorValues values = floatingPointVectorValues(knnVectorsReader,
fieldInfo);
+ if (values == null) {
+ continue;
+ }
+ int vectorCount = values.size();
if (vectorCount == 0) {
continue;
}
+ float[] centroid = getCentroid(knnVectorsReader, fieldInfo.name);
Review Comment:
it feels a little confusing to have both `getCentroid()` and `centroid()`.
Maybe we can call the latter `computeCentroid()` to call out that it does some
significant work?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java:
##########
@@ -0,0 +1,398 @@
+/*
+ * 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 static
org.apache.lucene.util.quantization.OptimizedScalarQuantizer.deQuantize;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.packed.DirectMonotonicReader;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Reads quantized vector values from the index input and returns float16
vector values after
+ * dequantizing them.
+ *
+ * <p>Used for read-only indexes whose raw float16 vectors have been dropped
to save storage: only
+ * the scalar-quantized bytes remain, so {@link #vectorValue(int)}
reconstructs float16 values by
+ * dequantizing them, with some precision loss.
+ *
+ * @lucene.internal
+ */
+abstract class OffHeapScalarQuantizedFloat16VectorValues extends
Float16VectorValues
+ implements HasIndexSlice {
+
+ final int dimension;
+ final int size;
+ final VectorSimilarityFunction similarityFunction;
+ final FlatVectorsScorer vectorsScorer;
+
+ final IndexInput slice;
+ final short[] vectorValue;
+ final byte[] byteValue;
+ final ByteBuffer byteBuffer;
+ final byte[] unpackedByteVectorValue;
+ final int byteSize;
+ private int lastOrd = -1;
+ final float[] correctiveValues;
+ int quantizedComponentSum;
+ final ScalarEncoding encoding;
+ final float[] centroid;
+
+ OffHeapScalarQuantizedFloat16VectorValues(
+ int dimension,
+ int size,
+ float[] centroid,
+ ScalarEncoding encoding,
+ VectorSimilarityFunction similarityFunction,
+ FlatVectorsScorer vectorsScorer,
+ IndexInput slice) {
+ this.dimension = dimension;
+ this.size = size;
+ this.similarityFunction = similarityFunction;
+ this.vectorsScorer = vectorsScorer;
+ this.slice = slice;
+ this.centroid = centroid;
+ this.correctiveValues = new float[3];
+ this.encoding = encoding;
+ int docPackedLength = encoding.getDocPackedLength(dimension);
+ this.byteSize = docPackedLength + (Float.BYTES * 3) + Integer.BYTES;
+ this.byteBuffer = ByteBuffer.allocate(docPackedLength);
+ this.vectorValue = new short[dimension];
+ this.byteValue = byteBuffer.array();
+ this.unpackedByteVectorValue = new byte[dimension];
+ }
+
+ @Override
+ public int dimension() {
+ return dimension;
+ }
+
+ @Override
+ public int size() {
+ return size;
+ }
+
+ @Override
+ public short[] vectorValue(int targetOrd) throws IOException {
+ if (lastOrd == targetOrd) {
+ return vectorValue;
+ }
+
+ // read quantized byte vector, correctiveValues and quantizedComponentSum
+ slice.seek((long) targetOrd * byteSize);
+ slice.readBytes(byteBuffer.array(), byteBuffer.arrayOffset(),
byteValue.length);
+ slice.readFloats(correctiveValues, 0, 3);
+ quantizedComponentSum = slice.readInt();
+
+ // unpack bytes
+ switch (encoding) {
+ case PACKED_NIBBLE ->
+ OffHeapScalarQuantizedVectorValues.unpackNibbles(byteValue,
unpackedByteVectorValue);
+ case SINGLE_BIT_QUERY_NIBBLE ->
+ OptimizedScalarQuantizer.unpackBinary(byteValue,
unpackedByteVectorValue);
+ case DIBIT_QUERY_NIBBLE ->
+ OptimizedScalarQuantizer.untransposeDibit(byteValue,
unpackedByteVectorValue);
+ case UNSIGNED_BYTE, SEVEN_BIT -> {
+ deQuantize(
Review Comment:
could we assign `unpackedByteVectorValue = byteValue` and fall through to
make the flow more uniform?
##########
lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestLucene104ScalarQuantizedVectorsFormat.java:
##########
@@ -105,6 +119,42 @@ public void testSearch() throws Exception {
}
}
+ public void testFloat16Search() throws Exception {
+ String fieldName = "field";
+ int numVectors = random().nextInt(99, 500);
+ int dims = random().nextInt(4, 65);
+ if (dims % 2 == 1) {
Review Comment:
or `2 * random().nextInt(2, 33)`?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorScorer.java:
##########
@@ -111,6 +111,43 @@ public RandomVectorScorer getRandomVectorScorer(
public RandomVectorScorer getRandomVectorScorer(
VectorSimilarityFunction similarityFunction, KnnVectorValues
vectorValues, short[] target)
throws IOException {
+ if (vectorValues instanceof QuantizedByteVectorValues qv) {
+ FlatVectorsScorer.checkDimensions(target.length, qv.dimension());
+ OptimizedScalarQuantizer quantizer = qv.getQuantizer();
+ ScalarEncoding scalarEncoding = qv.getScalarEncoding();
+ byte[] scratch = new
byte[scalarEncoding.getDiscreteDimensions(qv.dimension())];
+ final byte[] targetQuantized;
+ if (scalarEncoding.isAsymmetric() == false) {
+ targetQuantized = scratch;
+ } else {
+ // This is asymmetric quantization, we will pack the vector
+ targetQuantized = new
byte[scalarEncoding.getQueryPackedLength(scratch.length)];
+ }
+ // Inflate the fp16 query to fp32 and normalize there; quantization
operates on fp32.
Review Comment:
BTW I just saw
https://opensearch.org/blog/accelerating-fp16-vector-search-performance-using-bulk-simd-in-opensearch-3-5/
maybe there is some goodness there we can incorporate?
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene104/OffHeapScalarQuantizedFloat16VectorValues.java:
##########
@@ -0,0 +1,398 @@
+/*
+ * 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 static
org.apache.lucene.util.quantization.OptimizedScalarQuantizer.deQuantize;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.codecs.lucene90.IndexedDISI;
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.packed.DirectMonotonicReader;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Reads quantized vector values from the index input and returns float16
vector values after
+ * dequantizing them.
+ *
+ * <p>Used for read-only indexes whose raw float16 vectors have been dropped
to save storage: only
+ * the scalar-quantized bytes remain, so {@link #vectorValue(int)}
reconstructs float16 values by
+ * dequantizing them, with some precision loss.
+ *
+ * @lucene.internal
+ */
+abstract class OffHeapScalarQuantizedFloat16VectorValues extends
Float16VectorValues
+ implements HasIndexSlice {
+
+ final int dimension;
+ final int size;
+ final VectorSimilarityFunction similarityFunction;
+ final FlatVectorsScorer vectorsScorer;
+
+ final IndexInput slice;
+ final short[] vectorValue;
+ final byte[] byteValue;
+ final ByteBuffer byteBuffer;
+ final byte[] unpackedByteVectorValue;
+ final int byteSize;
+ private int lastOrd = -1;
+ final float[] correctiveValues;
+ int quantizedComponentSum;
+ final ScalarEncoding encoding;
+ final float[] centroid;
+
+ OffHeapScalarQuantizedFloat16VectorValues(
+ int dimension,
+ int size,
+ float[] centroid,
+ ScalarEncoding encoding,
+ VectorSimilarityFunction similarityFunction,
+ FlatVectorsScorer vectorsScorer,
+ IndexInput slice) {
+ this.dimension = dimension;
+ this.size = size;
+ this.similarityFunction = similarityFunction;
+ this.vectorsScorer = vectorsScorer;
+ this.slice = slice;
+ this.centroid = centroid;
+ this.correctiveValues = new float[3];
+ this.encoding = encoding;
+ int docPackedLength = encoding.getDocPackedLength(dimension);
+ this.byteSize = docPackedLength + (Float.BYTES * 3) + Integer.BYTES;
Review Comment:
hmm I guess we must read corrective values, and ... an int
--
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]
