msokolov commented on code in PR #16506:
URL: https://github.com/apache/lucene/pull/16506#discussion_r3830262854
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java:
##########
@@ -37,112 +36,123 @@
*
* @lucene.experimental
*/
-final class DedupFlatVectorsScorer implements FlatVectorsScorer {
- private static final FlatVectorsScorer SCORER =
+sealed class DedupFlatVectorsScorer implements FlatVectorsScorer
+ permits DedupScalarQuantizedVectorsScorer {
+
+ private static final FlatVectorsScorer FLAT_SCORER =
FlatVectorScorerUtil.getLucene99FlatVectorsScorer();
+ private final FlatVectorsScorer scorer;
+
+ DedupFlatVectorsScorer() {
+ this(FLAT_SCORER);
+ }
+
+ protected DedupFlatVectorsScorer(FlatVectorsScorer scorer) {
+ this.scorer = scorer;
+ }
+
+ /** Resolves the values to score; subclasses may unwrap composite values. */
Review Comment:
Can we call it `unwrap`? I was confused about what this was doing when I
saw it used below
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java:
##########
@@ -37,112 +36,123 @@
*
* @lucene.experimental
*/
-final class DedupFlatVectorsScorer implements FlatVectorsScorer {
- private static final FlatVectorsScorer SCORER =
+sealed class DedupFlatVectorsScorer implements FlatVectorsScorer
+ permits DedupScalarQuantizedVectorsScorer {
+
+ private static final FlatVectorsScorer FLAT_SCORER =
FlatVectorScorerUtil.getLucene99FlatVectorsScorer();
+ private final FlatVectorsScorer scorer;
Review Comment:
since this is a scorer that delegates can we name this the conventional way
as either `delegate` or `in`?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java:
##########
@@ -41,6 +35,10 @@
* or merging existing segments (never both), delegating to {@link
DedupFlushContext} or {@link
* DedupMergeContext} accordingly.
*
+ * <p>Also used by {@link DedupScalarQuantizedVectorsFormat} (with a non-null
{@link
+ * DedupQuantizer}) to additionally write a quantized copy of each FLOAT32
group, into a separate
Review Comment:
maybe in future they could be `FLOAT16`? Should we say each "full-precision"
or "unqunatized" group?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupQuantizer.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.codecs.lucene104.OffHeapScalarQuantizedVectorValues;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.VectorEncoding;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Write-side helper performing <i>data-blind</i> scalar quantization of
de-duplicated vectors.
+ *
+ * <p>Quantization assumes input vectors are evenly distributed around the
origin, i.e. the centroid
+ * is a zero vector. The quantized record of a vector (packed bytes, optimized
intervals, an
+ * additional correction and the component sum) is then a pure function of the
raw vector and the
+ * {@link Flavor} derived from the field's similarity function: independent of
the other vectors in
+ * the segment. Identical vectors thus quantize identically and the quantized
data de-duplicates
+ * along with the raw data, shared across all documents and fields of a group
with the same flavor.
+ *
+ * <p>For the same reason, a record already computed by a source segment does
not change on merge:
+ * when a distinct vector originates from a segment in this format (with the
same encoding and
+ * flavor), its quantized record is <b>copied</b> instead of re-reading the
raw vector and
+ * re-quantizing it.
+ *
+ * <p>Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions exactly
+ * (per flavor), so they are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
+ *
+ * @lucene.experimental
+ */
+record DedupQuantizer(ScalarEncoding encoding) {
+
+ /**
+ * How a vector is prepared and quantized, derived from the field's {@link
+ * VectorSimilarityFunction}. Fields of a group whose similarity functions
map to the same flavor
+ * share one quantized record per distinct vector; a group stores one block
per flavor in use.
+ */
+ enum Flavor {
+
+ /** Vectors quantized as-is; the additional correction holds the squared
norm. */
+ EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN, false),
+
+ /**
+ * Vectors quantized as-is; the additional correction holds the dot
product with the (zero)
+ * centroid, i.e. zero.
+ */
+ DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT, false),
+
+ /**
+ * Vectors are l2-normalized before quantization (for cosine similarity,
scored as a dot
+ * product); the additional correction is zero as for {@link #DOT_PRODUCT}.
+ */
+ NORMALIZED(VectorSimilarityFunction.DOT_PRODUCT, true);
+
+ private final VectorSimilarityFunction quantizerFunction;
+ private final boolean normalized;
+
+ Flavor(VectorSimilarityFunction quantizerFunction, boolean normalized) {
+ this.quantizerFunction = quantizerFunction;
+ this.normalized = normalized;
+ }
+
+ static Flavor of(VectorSimilarityFunction function) {
+ return switch (function) {
+ case EUCLIDEAN -> EUCLIDEAN;
+ case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> DOT_PRODUCT;
+ case COSINE -> NORMALIZED;
+ };
+ }
+
+ /**
+ * The quantizer producing this flavor's records against a zero centroid.
NOTE: {@link
+ * #NORMALIZED} quantizes with {@link
VectorSimilarityFunction#DOT_PRODUCT} (identical bytes and
+ * corrections, without the unit-centroid requirement of a
cosine-configured quantizer).
+ */
+ OptimizedScalarQuantizer quantizer() {
+ return new OptimizedScalarQuantizer(quantizerFunction);
+ }
+
+ boolean normalized() {
+ return normalized;
+ }
+ }
+
+ /** Supplies the distinct (raw) float vector at a group ordinal. */
+ interface FloatVectorSupplier {
Review Comment:
what's the visibility of this? If we later want to quantize Float16 (ie
`short[]`) or even `byte[]` will we need a breaking change?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupQuantizer.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.codecs.lucene104.OffHeapScalarQuantizedVectorValues;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.VectorEncoding;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Write-side helper performing <i>data-blind</i> scalar quantization of
de-duplicated vectors.
+ *
+ * <p>Quantization assumes input vectors are evenly distributed around the
origin, i.e. the centroid
+ * is a zero vector. The quantized record of a vector (packed bytes, optimized
intervals, an
+ * additional correction and the component sum) is then a pure function of the
raw vector and the
+ * {@link Flavor} derived from the field's similarity function: independent of
the other vectors in
+ * the segment. Identical vectors thus quantize identically and the quantized
data de-duplicates
+ * along with the raw data, shared across all documents and fields of a group
with the same flavor.
+ *
+ * <p>For the same reason, a record already computed by a source segment does
not change on merge:
+ * when a distinct vector originates from a segment in this format (with the
same encoding and
+ * flavor), its quantized record is <b>copied</b> instead of re-reading the
raw vector and
+ * re-quantizing it.
+ *
+ * <p>Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions exactly
+ * (per flavor), so they are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
+ *
+ * @lucene.experimental
+ */
+record DedupQuantizer(ScalarEncoding encoding) {
+
+ /**
+ * How a vector is prepared and quantized, derived from the field's {@link
+ * VectorSimilarityFunction}. Fields of a group whose similarity functions
map to the same flavor
+ * share one quantized record per distinct vector; a group stores one block
per flavor in use.
+ */
+ enum Flavor {
+
+ /** Vectors quantized as-is; the additional correction holds the squared
norm. */
+ EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN, false),
+
+ /**
+ * Vectors quantized as-is; the additional correction holds the dot
product with the (zero)
+ * centroid, i.e. zero.
+ */
+ DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT, false),
+
+ /**
+ * Vectors are l2-normalized before quantization (for cosine similarity,
scored as a dot
+ * product); the additional correction is zero as for {@link #DOT_PRODUCT}.
+ */
+ NORMALIZED(VectorSimilarityFunction.DOT_PRODUCT, true);
+
+ private final VectorSimilarityFunction quantizerFunction;
+ private final boolean normalized;
+
+ Flavor(VectorSimilarityFunction quantizerFunction, boolean normalized) {
+ this.quantizerFunction = quantizerFunction;
+ this.normalized = normalized;
+ }
+
+ static Flavor of(VectorSimilarityFunction function) {
+ return switch (function) {
+ case EUCLIDEAN -> EUCLIDEAN;
+ case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> DOT_PRODUCT;
+ case COSINE -> NORMALIZED;
+ };
+ }
+
+ /**
+ * The quantizer producing this flavor's records against a zero centroid.
NOTE: {@link
+ * #NORMALIZED} quantizes with {@link
VectorSimilarityFunction#DOT_PRODUCT} (identical bytes and
+ * corrections, without the unit-centroid requirement of a
cosine-configured quantizer).
+ */
+ OptimizedScalarQuantizer quantizer() {
+ return new OptimizedScalarQuantizer(quantizerFunction);
+ }
+
+ boolean normalized() {
+ return normalized;
+ }
+ }
+
+ /** Supplies the distinct (raw) float vector at a group ordinal. */
+ interface FloatVectorSupplier {
+ float[] get(int ord) throws IOException;
+ }
+
+ /**
+ * The already-quantized record of a distinct vector, held by a source
segment at {@code ord} with
+ * the given flavor.
+ */
+ record PreQuantized(QuantizedByteVectorValues values, Flavor flavor, int
ord) {}
+
+ /**
+ * Supplies the {@link PreQuantized} record of the distinct vector at a
group ordinal, or {@code
+ * null} when its source segment does not hold one.
+ */
+ interface PreQuantizedSupplier {
+ PreQuantized get(int ord) throws IOException;
+ }
+
+ /** The encoding, offset and size of one flavor's quantized data block. */
+ record QuantizedBlock(
+ ScalarEncoding encoding, long quantizedDataOffset, long
quantizedDataSize) {}
+
+ /** Writes an empty flavor-block list, for groups without quantized data. */
+ static void writeEmptyGroup(IndexOutput meta) throws IOException {
+ meta.writeVInt(0);
+ }
+
+ /** Reads the flavor-block list of a group, the counterpart of {@link
#writeGroup}. */
+ static Map<Flavor, QuantizedBlock> readGroup(IndexInput meta) throws
IOException {
+ int numFlavors = meta.readVInt();
+ Map<Flavor, QuantizedBlock> blocks = new EnumMap<>(Flavor.class);
+ for (int i = 0; i < numFlavors; i++) {
+ int flavorOrd = meta.readVInt();
+ if (flavorOrd < 0 || flavorOrd >= Flavor.values().length) {
+ throw new CorruptIndexException("Invalid flavor ordinal: " +
flavorOrd, meta);
+ }
+ Flavor flavor = Flavor.values()[flavorOrd];
Review Comment:
Lucene doesn't generally rely on Enum ordinals assigned by the JVM and
rather will explicitly define codes to store in the index, that can later be
mapped to enums as needed. This enables, for example, later removing support
for an ordinal. EG see the way VectorSimlarityFunction is handled in
`Lucene94FieldInfosFormat`, or indeed the way `ScalarEncoding` is being
deserialized here using a custom function (fromWireNumber)
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupQuantizer.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.codecs.lucene104.OffHeapScalarQuantizedVectorValues;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.VectorEncoding;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Write-side helper performing <i>data-blind</i> scalar quantization of
de-duplicated vectors.
+ *
+ * <p>Quantization assumes input vectors are evenly distributed around the
origin, i.e. the centroid
+ * is a zero vector. The quantized record of a vector (packed bytes, optimized
intervals, an
+ * additional correction and the component sum) is then a pure function of the
raw vector and the
+ * {@link Flavor} derived from the field's similarity function: independent of
the other vectors in
+ * the segment. Identical vectors thus quantize identically and the quantized
data de-duplicates
+ * along with the raw data, shared across all documents and fields of a group
with the same flavor.
+ *
+ * <p>For the same reason, a record already computed by a source segment does
not change on merge:
+ * when a distinct vector originates from a segment in this format (with the
same encoding and
+ * flavor), its quantized record is <b>copied</b> instead of re-reading the
raw vector and
+ * re-quantizing it.
+ *
+ * <p>Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions exactly
+ * (per flavor), so they are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
+ *
+ * @lucene.experimental
+ */
+record DedupQuantizer(ScalarEncoding encoding) {
+
+ /**
+ * How a vector is prepared and quantized, derived from the field's {@link
+ * VectorSimilarityFunction}. Fields of a group whose similarity functions
map to the same flavor
+ * share one quantized record per distinct vector; a group stores one block
per flavor in use.
+ */
+ enum Flavor {
+
+ /** Vectors quantized as-is; the additional correction holds the squared
norm. */
+ EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN, false),
+
+ /**
+ * Vectors quantized as-is; the additional correction holds the dot
product with the (zero)
+ * centroid, i.e. zero.
+ */
+ DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT, false),
+
+ /**
+ * Vectors are l2-normalized before quantization (for cosine similarity,
scored as a dot
+ * product); the additional correction is zero as for {@link #DOT_PRODUCT}.
+ */
+ NORMALIZED(VectorSimilarityFunction.DOT_PRODUCT, true);
+
+ private final VectorSimilarityFunction quantizerFunction;
+ private final boolean normalized;
+
+ Flavor(VectorSimilarityFunction quantizerFunction, boolean normalized) {
+ this.quantizerFunction = quantizerFunction;
+ this.normalized = normalized;
+ }
+
+ static Flavor of(VectorSimilarityFunction function) {
+ return switch (function) {
+ case EUCLIDEAN -> EUCLIDEAN;
+ case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> DOT_PRODUCT;
+ case COSINE -> NORMALIZED;
+ };
+ }
+
+ /**
+ * The quantizer producing this flavor's records against a zero centroid.
NOTE: {@link
+ * #NORMALIZED} quantizes with {@link
VectorSimilarityFunction#DOT_PRODUCT} (identical bytes and
+ * corrections, without the unit-centroid requirement of a
cosine-configured quantizer).
+ */
+ OptimizedScalarQuantizer quantizer() {
+ return new OptimizedScalarQuantizer(quantizerFunction);
+ }
+
+ boolean normalized() {
+ return normalized;
+ }
+ }
+
+ /** Supplies the distinct (raw) float vector at a group ordinal. */
+ interface FloatVectorSupplier {
+ float[] get(int ord) throws IOException;
+ }
+
+ /**
+ * The already-quantized record of a distinct vector, held by a source
segment at {@code ord} with
+ * the given flavor.
+ */
+ record PreQuantized(QuantizedByteVectorValues values, Flavor flavor, int
ord) {}
+
+ /**
+ * Supplies the {@link PreQuantized} record of the distinct vector at a
group ordinal, or {@code
+ * null} when its source segment does not hold one.
+ */
+ interface PreQuantizedSupplier {
+ PreQuantized get(int ord) throws IOException;
+ }
+
+ /** The encoding, offset and size of one flavor's quantized data block. */
+ record QuantizedBlock(
+ ScalarEncoding encoding, long quantizedDataOffset, long
quantizedDataSize) {}
+
+ /** Writes an empty flavor-block list, for groups without quantized data. */
+ static void writeEmptyGroup(IndexOutput meta) throws IOException {
+ meta.writeVInt(0);
+ }
+
+ /** Reads the flavor-block list of a group, the counterpart of {@link
#writeGroup}. */
+ static Map<Flavor, QuantizedBlock> readGroup(IndexInput meta) throws
IOException {
+ int numFlavors = meta.readVInt();
+ Map<Flavor, QuantizedBlock> blocks = new EnumMap<>(Flavor.class);
+ for (int i = 0; i < numFlavors; i++) {
+ int flavorOrd = meta.readVInt();
+ if (flavorOrd < 0 || flavorOrd >= Flavor.values().length) {
+ throw new CorruptIndexException("Invalid flavor ordinal: " +
flavorOrd, meta);
+ }
+ Flavor flavor = Flavor.values()[flavorOrd];
+ int wireNumber = meta.readVInt();
+ ScalarEncoding encoding =
+ ScalarEncoding.fromWireNumber(wireNumber)
+ .orElseThrow(
+ () ->
+ new CorruptIndexException(
+ "Invalid scalar encoding wire number: " +
wireNumber, meta));
+ long offset = meta.readLong();
+ long size = meta.readLong();
+ if (blocks.put(flavor, new QuantizedBlock(encoding, offset, size)) !=
null) {
+ throw new CorruptIndexException("Duplicate flavor: " + flavor, meta);
+ }
+ }
+ return blocks;
+ }
+
+ /**
+ * Quantizes and writes a group's distinct vectors: one block per flavor,
one record per group
+ * ordinal within each block. Wherever a {@code preQuantized} record with a
matching encoding and
+ * flavor is available (i.e. the vector originates from a segment in this
format), it is copied
+ * as-is; otherwise the raw vector is quantized. The block locations are
written to {@code meta};
+ * non-FLOAT32 groups (stored raw only) record an empty list.
+ */
+ void writeGroup(
+ IndexOutput meta,
+ IndexOutput quantizedVectorData,
+ VectorEncoding groupEncoding,
+ int dimension,
+ int numVectors,
+ Set<Flavor> flavors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ if (groupEncoding != VectorEncoding.FLOAT32) {
+ writeEmptyGroup(meta);
+ return;
+ }
+
+ meta.writeVInt(flavors.size());
+ for (Flavor flavor : Flavor.values()) {
+ if (flavors.contains(flavor) == false) {
+ continue;
+ }
+ QuantizedBlock block =
+ writeFlavorBlock(
+ quantizedVectorData, flavor, dimension, numVectors, vectors,
preQuantized);
+ meta.writeVInt(flavor.ordinal());
+ meta.writeVInt(block.encoding().getWireNumber());
+ meta.writeLong(block.quantizedDataOffset());
+ meta.writeLong(block.quantizedDataSize());
+ }
+ }
+
+ private QuantizedBlock writeFlavorBlock(
+ IndexOutput quantizedVectorData,
+ Flavor flavor,
+ int dimension,
+ int numVectors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ OptimizedScalarQuantizer quantizer = flavor.quantizer();
+ float[] zeroCentroid = new float[dimension];
+ float[] normalized = flavor.normalized() ? new float[dimension] : null;
+ byte[] scratch = new byte[encoding.getDiscreteDimensions(dimension)];
+ byte[] packed =
+ switch (encoding) {
+ case UNSIGNED_BYTE, SEVEN_BIT -> scratch;
+ case PACKED_NIBBLE, SINGLE_BIT_QUERY_NIBBLE, DIBIT_QUERY_NIBBLE ->
+ new byte[encoding.getDocPackedLength(scratch.length)];
+ };
+
+ long quantizedDataOffset =
quantizedVectorData.alignFilePointer(Float.BYTES);
+ for (int ord = 0; ord < numVectors; ord++) {
+ // The quantized record is a pure function of the raw vector and the
flavor: a record
+ // already computed by a source segment can be copied instead of
re-quantizing.
+ PreQuantized pre = preQuantized == null ? null : preQuantized.get(ord);
Review Comment:
I feel like this would read clearer with an explicit `if (preQuantized ==
null) {` wrapping the following `if`. We don't use `pre` anywhere below, so the
only value of assigning null to it is to skip the block, which we could do in a
more idiomatic way with an if statement.
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java:
##########
@@ -86,7 +88,18 @@ void addField(FieldInfo fieldInfo, MergeState mergeState)
throws IOException {
mergeState.segmentInfo.maxDoc()));
}
- void finish(IndexOutput meta, IndexOutput vectorData) throws IOException {
+ /**
+ * Merges each group's distinct vectors, followed by per-field metadata.
When {@code quantizer} is
+ * non-null, a quantized copy of each FLOAT32 group is also written and its
block location
Review Comment:
each "full precision" group?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupScalarQuantizedVectorsFormat.java:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * A scalar quantized version of {@link DedupFlatVectorsFormat} that stores
each distinct vector
+ * once, in both raw and quantized form.
+ *
+ * <p>Quantization is <i>data-blind</i>: input vectors are assumed to be
evenly distributed, i.e.
+ * centered on a zero vector (see {@link DedupQuantizer}). The quantized
record of a vector is then
+ * a pure function of the raw vector and the {@link DedupQuantizer.Flavor}
derived from the field's
+ * similarity function — independent of the other vectors in the segment — so
quantized records
+ * de-duplicate like raw vectors: identical vectors across documents and
fields of a group whose
+ * similarity functions map to the same flavor share one quantized record,
resolved through the same
+ * {@code fieldOrdToGroupOrd} translation map. Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions per flavor,
+ * and are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
Note the accuracy
+ * tradeoff relative to that format, which centers vectors on a per-field
centroid before
+ * quantizing.
+ *
+ * <p>Only {@link org.apache.lucene.index.VectorEncoding#FLOAT32} vectors are
quantized; BYTE and
+ * FLOAT16 vectors are stored raw only, identical to {@link
DedupFlatVectorsFormat}.
Review Comment:
ah, okay, it's by design. But I think we would eventually want to be able to
quantize FLOAT16 too?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupQuantizer.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.codecs.lucene104.OffHeapScalarQuantizedVectorValues;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.VectorEncoding;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Write-side helper performing <i>data-blind</i> scalar quantization of
de-duplicated vectors.
+ *
+ * <p>Quantization assumes input vectors are evenly distributed around the
origin, i.e. the centroid
+ * is a zero vector. The quantized record of a vector (packed bytes, optimized
intervals, an
+ * additional correction and the component sum) is then a pure function of the
raw vector and the
+ * {@link Flavor} derived from the field's similarity function: independent of
the other vectors in
+ * the segment. Identical vectors thus quantize identically and the quantized
data de-duplicates
+ * along with the raw data, shared across all documents and fields of a group
with the same flavor.
+ *
+ * <p>For the same reason, a record already computed by a source segment does
not change on merge:
+ * when a distinct vector originates from a segment in this format (with the
same encoding and
+ * flavor), its quantized record is <b>copied</b> instead of re-reading the
raw vector and
+ * re-quantizing it.
+ *
+ * <p>Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions exactly
+ * (per flavor), so they are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
+ *
+ * @lucene.experimental
+ */
+record DedupQuantizer(ScalarEncoding encoding) {
+
+ /**
+ * How a vector is prepared and quantized, derived from the field's {@link
+ * VectorSimilarityFunction}. Fields of a group whose similarity functions
map to the same flavor
+ * share one quantized record per distinct vector; a group stores one block
per flavor in use.
+ */
+ enum Flavor {
+
+ /** Vectors quantized as-is; the additional correction holds the squared
norm. */
+ EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN, false),
+
+ /**
+ * Vectors quantized as-is; the additional correction holds the dot
product with the (zero)
+ * centroid, i.e. zero.
+ */
+ DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT, false),
+
+ /**
+ * Vectors are l2-normalized before quantization (for cosine similarity,
scored as a dot
+ * product); the additional correction is zero as for {@link #DOT_PRODUCT}.
+ */
+ NORMALIZED(VectorSimilarityFunction.DOT_PRODUCT, true);
+
+ private final VectorSimilarityFunction quantizerFunction;
+ private final boolean normalized;
+
+ Flavor(VectorSimilarityFunction quantizerFunction, boolean normalized) {
+ this.quantizerFunction = quantizerFunction;
+ this.normalized = normalized;
+ }
+
+ static Flavor of(VectorSimilarityFunction function) {
+ return switch (function) {
+ case EUCLIDEAN -> EUCLIDEAN;
+ case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> DOT_PRODUCT;
+ case COSINE -> NORMALIZED;
+ };
+ }
+
+ /**
+ * The quantizer producing this flavor's records against a zero centroid.
NOTE: {@link
+ * #NORMALIZED} quantizes with {@link
VectorSimilarityFunction#DOT_PRODUCT} (identical bytes and
+ * corrections, without the unit-centroid requirement of a
cosine-configured quantizer).
+ */
+ OptimizedScalarQuantizer quantizer() {
+ return new OptimizedScalarQuantizer(quantizerFunction);
+ }
+
+ boolean normalized() {
+ return normalized;
+ }
+ }
+
+ /** Supplies the distinct (raw) float vector at a group ordinal. */
+ interface FloatVectorSupplier {
+ float[] get(int ord) throws IOException;
+ }
+
+ /**
+ * The already-quantized record of a distinct vector, held by a source
segment at {@code ord} with
+ * the given flavor.
+ */
+ record PreQuantized(QuantizedByteVectorValues values, Flavor flavor, int
ord) {}
+
+ /**
+ * Supplies the {@link PreQuantized} record of the distinct vector at a
group ordinal, or {@code
+ * null} when its source segment does not hold one.
+ */
+ interface PreQuantizedSupplier {
+ PreQuantized get(int ord) throws IOException;
+ }
+
+ /** The encoding, offset and size of one flavor's quantized data block. */
+ record QuantizedBlock(
+ ScalarEncoding encoding, long quantizedDataOffset, long
quantizedDataSize) {}
+
+ /** Writes an empty flavor-block list, for groups without quantized data. */
+ static void writeEmptyGroup(IndexOutput meta) throws IOException {
+ meta.writeVInt(0);
+ }
+
+ /** Reads the flavor-block list of a group, the counterpart of {@link
#writeGroup}. */
+ static Map<Flavor, QuantizedBlock> readGroup(IndexInput meta) throws
IOException {
+ int numFlavors = meta.readVInt();
+ Map<Flavor, QuantizedBlock> blocks = new EnumMap<>(Flavor.class);
+ for (int i = 0; i < numFlavors; i++) {
+ int flavorOrd = meta.readVInt();
+ if (flavorOrd < 0 || flavorOrd >= Flavor.values().length) {
+ throw new CorruptIndexException("Invalid flavor ordinal: " +
flavorOrd, meta);
+ }
+ Flavor flavor = Flavor.values()[flavorOrd];
+ int wireNumber = meta.readVInt();
+ ScalarEncoding encoding =
+ ScalarEncoding.fromWireNumber(wireNumber)
+ .orElseThrow(
+ () ->
+ new CorruptIndexException(
+ "Invalid scalar encoding wire number: " +
wireNumber, meta));
+ long offset = meta.readLong();
+ long size = meta.readLong();
+ if (blocks.put(flavor, new QuantizedBlock(encoding, offset, size)) !=
null) {
+ throw new CorruptIndexException("Duplicate flavor: " + flavor, meta);
+ }
+ }
+ return blocks;
+ }
+
+ /**
+ * Quantizes and writes a group's distinct vectors: one block per flavor,
one record per group
+ * ordinal within each block. Wherever a {@code preQuantized} record with a
matching encoding and
+ * flavor is available (i.e. the vector originates from a segment in this
format), it is copied
+ * as-is; otherwise the raw vector is quantized. The block locations are
written to {@code meta};
+ * non-FLOAT32 groups (stored raw only) record an empty list.
+ */
+ void writeGroup(
+ IndexOutput meta,
+ IndexOutput quantizedVectorData,
+ VectorEncoding groupEncoding,
+ int dimension,
+ int numVectors,
+ Set<Flavor> flavors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ if (groupEncoding != VectorEncoding.FLOAT32) {
+ writeEmptyGroup(meta);
+ return;
+ }
+
+ meta.writeVInt(flavors.size());
+ for (Flavor flavor : Flavor.values()) {
+ if (flavors.contains(flavor) == false) {
+ continue;
+ }
+ QuantizedBlock block =
+ writeFlavorBlock(
+ quantizedVectorData, flavor, dimension, numVectors, vectors,
preQuantized);
+ meta.writeVInt(flavor.ordinal());
+ meta.writeVInt(block.encoding().getWireNumber());
+ meta.writeLong(block.quantizedDataOffset());
+ meta.writeLong(block.quantizedDataSize());
+ }
+ }
+
+ private QuantizedBlock writeFlavorBlock(
+ IndexOutput quantizedVectorData,
+ Flavor flavor,
+ int dimension,
+ int numVectors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ OptimizedScalarQuantizer quantizer = flavor.quantizer();
+ float[] zeroCentroid = new float[dimension];
+ float[] normalized = flavor.normalized() ? new float[dimension] : null;
+ byte[] scratch = new byte[encoding.getDiscreteDimensions(dimension)];
+ byte[] packed =
+ switch (encoding) {
+ case UNSIGNED_BYTE, SEVEN_BIT -> scratch;
+ case PACKED_NIBBLE, SINGLE_BIT_QUERY_NIBBLE, DIBIT_QUERY_NIBBLE ->
+ new byte[encoding.getDocPackedLength(scratch.length)];
+ };
+
+ long quantizedDataOffset =
quantizedVectorData.alignFilePointer(Float.BYTES);
+ for (int ord = 0; ord < numVectors; ord++) {
+ // The quantized record is a pure function of the raw vector and the
flavor: a record
+ // already computed by a source segment can be copied instead of
re-quantizing.
+ PreQuantized pre = preQuantized == null ? null : preQuantized.get(ord);
+ if (pre != null && pre.flavor() == flavor &&
pre.values().getScalarEncoding() == encoding) {
+ // NOTE: read the packed bytes before the corrective terms, which are
then served from the
+ // record cached by the read of the packed bytes
+ writeRecord(
+ quantizedVectorData,
+ pre.values().vectorValue(pre.ord()),
+ pre.values().getCorrectiveTerms(pre.ord()));
+ continue;
+ }
+
+ float[] vector = vectors.get(ord);
+ if (flavor.normalized()) {
+ // normalize a copy: the source buffer is shared / owned by the group
+ System.arraycopy(vector, 0, normalized, 0, dimension);
+ VectorUtil.l2normalize(normalized);
+ vector = normalized;
+ }
+ // NOTE: scalarQuantize subtracts the centroid from the input in place,
but with a zero
+ // centroid the values are unchanged, so shared / owned buffers are safe
to pass directly.
+ OptimizedScalarQuantizer.QuantizationResult corrections =
+ quantizer.scalarQuantize(vector, scratch, encoding.getBits(),
zeroCentroid);
+ switch (encoding) {
+ case PACKED_NIBBLE ->
OffHeapScalarQuantizedVectorValues.packNibbles(scratch, packed);
+ case SINGLE_BIT_QUERY_NIBBLE ->
OptimizedScalarQuantizer.packAsBinary(scratch, packed);
+ case DIBIT_QUERY_NIBBLE ->
OptimizedScalarQuantizer.transposeDibit(scratch, packed);
+ case UNSIGNED_BYTE, SEVEN_BIT -> {}
+ }
+ writeRecord(quantizedVectorData, packed, corrections);
+ }
+ long quantizedDataSize = quantizedVectorData.getFilePointer() -
quantizedDataOffset;
+
+ return new QuantizedBlock(encoding, quantizedDataOffset,
quantizedDataSize);
+ }
+
+ private static void writeRecord(
+ IndexOutput quantizedVectorData,
+ byte[] packed,
+ OptimizedScalarQuantizer.QuantizationResult corrections)
+ throws IOException {
+ quantizedVectorData.writeBytes(packed, packed.length);
+
quantizedVectorData.writeInt(Float.floatToIntBits(corrections.lowerInterval()));
Review Comment:
for flavors where the corrections are uniformly zero, can we skip writing
them?
##########
lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupScalarQuantizedVectorsFormat.java:
##########
@@ -0,0 +1,430 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import static
org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
+import static
org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
+import static org.apache.lucene.index.VectorEncoding.BYTE;
+import static org.apache.lucene.index.VectorEncoding.FLOAT32;
+import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
+import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT;
+import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
+import static
org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT;
+import static org.hamcrest.Matchers.instanceOf;
+
+import java.util.Arrays;
+import java.util.List;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.KnnByteVectorField;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.document.NumericDocValuesField;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.AcceptDocs;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.VectorUtil;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Tests that {@link DedupHnswScalarQuantizedVectorsFormat} stores each
distinct vector once, in
+ * both raw and quantized form. General de-duplication behavior is covered by
{@link
+ * TestDedupFlatVectorsFormat}; this test focuses on the quantized side.
+ */
+public class TestDedupScalarQuantizedVectorsFormat extends LuceneTestCase {
Review Comment:
since we explicitly forbade copy() on some the vector values, let's add a
test ensuring that those methods throw Unsupported exception
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupQuantizer.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.sandbox.codecs.dedup;
+
+import java.io.IOException;
+import java.util.EnumMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.codecs.lucene104.OffHeapScalarQuantizedVectorValues;
+import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.VectorEncoding;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
+
+/**
+ * Write-side helper performing <i>data-blind</i> scalar quantization of
de-duplicated vectors.
+ *
+ * <p>Quantization assumes input vectors are evenly distributed around the
origin, i.e. the centroid
+ * is a zero vector. The quantized record of a vector (packed bytes, optimized
intervals, an
+ * additional correction and the component sum) is then a pure function of the
raw vector and the
+ * {@link Flavor} derived from the field's similarity function: independent of
the other vectors in
+ * the segment. Identical vectors thus quantize identically and the quantized
data de-duplicates
+ * along with the raw data, shared across all documents and fields of a group
with the same flavor.
+ *
+ * <p>For the same reason, a record already computed by a source segment does
not change on merge:
+ * when a distinct vector originates from a segment in this format (with the
same encoding and
+ * flavor), its quantized record is <b>copied</b> instead of re-reading the
raw vector and
+ * re-quantizing it.
+ *
+ * <p>Records follow the {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat}
conventions exactly
+ * (per flavor), so they are scored by the stock {@link
+ * org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorScorer}.
+ *
+ * @lucene.experimental
+ */
+record DedupQuantizer(ScalarEncoding encoding) {
+
+ /**
+ * How a vector is prepared and quantized, derived from the field's {@link
+ * VectorSimilarityFunction}. Fields of a group whose similarity functions
map to the same flavor
+ * share one quantized record per distinct vector; a group stores one block
per flavor in use.
+ */
+ enum Flavor {
+
+ /** Vectors quantized as-is; the additional correction holds the squared
norm. */
+ EUCLIDEAN(VectorSimilarityFunction.EUCLIDEAN, false),
+
+ /**
+ * Vectors quantized as-is; the additional correction holds the dot
product with the (zero)
+ * centroid, i.e. zero.
+ */
+ DOT_PRODUCT(VectorSimilarityFunction.DOT_PRODUCT, false),
+
+ /**
+ * Vectors are l2-normalized before quantization (for cosine similarity,
scored as a dot
+ * product); the additional correction is zero as for {@link #DOT_PRODUCT}.
+ */
+ NORMALIZED(VectorSimilarityFunction.DOT_PRODUCT, true);
+
+ private final VectorSimilarityFunction quantizerFunction;
+ private final boolean normalized;
+
+ Flavor(VectorSimilarityFunction quantizerFunction, boolean normalized) {
+ this.quantizerFunction = quantizerFunction;
+ this.normalized = normalized;
+ }
+
+ static Flavor of(VectorSimilarityFunction function) {
+ return switch (function) {
+ case EUCLIDEAN -> EUCLIDEAN;
+ case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> DOT_PRODUCT;
+ case COSINE -> NORMALIZED;
+ };
+ }
+
+ /**
+ * The quantizer producing this flavor's records against a zero centroid.
NOTE: {@link
+ * #NORMALIZED} quantizes with {@link
VectorSimilarityFunction#DOT_PRODUCT} (identical bytes and
+ * corrections, without the unit-centroid requirement of a
cosine-configured quantizer).
+ */
+ OptimizedScalarQuantizer quantizer() {
+ return new OptimizedScalarQuantizer(quantizerFunction);
+ }
+
+ boolean normalized() {
+ return normalized;
+ }
+ }
+
+ /** Supplies the distinct (raw) float vector at a group ordinal. */
+ interface FloatVectorSupplier {
+ float[] get(int ord) throws IOException;
+ }
+
+ /**
+ * The already-quantized record of a distinct vector, held by a source
segment at {@code ord} with
+ * the given flavor.
+ */
+ record PreQuantized(QuantizedByteVectorValues values, Flavor flavor, int
ord) {}
+
+ /**
+ * Supplies the {@link PreQuantized} record of the distinct vector at a
group ordinal, or {@code
+ * null} when its source segment does not hold one.
+ */
+ interface PreQuantizedSupplier {
+ PreQuantized get(int ord) throws IOException;
+ }
+
+ /** The encoding, offset and size of one flavor's quantized data block. */
+ record QuantizedBlock(
+ ScalarEncoding encoding, long quantizedDataOffset, long
quantizedDataSize) {}
+
+ /** Writes an empty flavor-block list, for groups without quantized data. */
+ static void writeEmptyGroup(IndexOutput meta) throws IOException {
+ meta.writeVInt(0);
+ }
+
+ /** Reads the flavor-block list of a group, the counterpart of {@link
#writeGroup}. */
+ static Map<Flavor, QuantizedBlock> readGroup(IndexInput meta) throws
IOException {
+ int numFlavors = meta.readVInt();
+ Map<Flavor, QuantizedBlock> blocks = new EnumMap<>(Flavor.class);
+ for (int i = 0; i < numFlavors; i++) {
+ int flavorOrd = meta.readVInt();
+ if (flavorOrd < 0 || flavorOrd >= Flavor.values().length) {
+ throw new CorruptIndexException("Invalid flavor ordinal: " +
flavorOrd, meta);
+ }
+ Flavor flavor = Flavor.values()[flavorOrd];
+ int wireNumber = meta.readVInt();
+ ScalarEncoding encoding =
+ ScalarEncoding.fromWireNumber(wireNumber)
+ .orElseThrow(
+ () ->
+ new CorruptIndexException(
+ "Invalid scalar encoding wire number: " +
wireNumber, meta));
+ long offset = meta.readLong();
+ long size = meta.readLong();
+ if (blocks.put(flavor, new QuantizedBlock(encoding, offset, size)) !=
null) {
+ throw new CorruptIndexException("Duplicate flavor: " + flavor, meta);
+ }
+ }
+ return blocks;
+ }
+
+ /**
+ * Quantizes and writes a group's distinct vectors: one block per flavor,
one record per group
+ * ordinal within each block. Wherever a {@code preQuantized} record with a
matching encoding and
+ * flavor is available (i.e. the vector originates from a segment in this
format), it is copied
+ * as-is; otherwise the raw vector is quantized. The block locations are
written to {@code meta};
+ * non-FLOAT32 groups (stored raw only) record an empty list.
+ */
+ void writeGroup(
+ IndexOutput meta,
+ IndexOutput quantizedVectorData,
+ VectorEncoding groupEncoding,
+ int dimension,
+ int numVectors,
+ Set<Flavor> flavors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ if (groupEncoding != VectorEncoding.FLOAT32) {
+ writeEmptyGroup(meta);
+ return;
+ }
+
+ meta.writeVInt(flavors.size());
+ for (Flavor flavor : Flavor.values()) {
+ if (flavors.contains(flavor) == false) {
+ continue;
+ }
+ QuantizedBlock block =
+ writeFlavorBlock(
+ quantizedVectorData, flavor, dimension, numVectors, vectors,
preQuantized);
+ meta.writeVInt(flavor.ordinal());
+ meta.writeVInt(block.encoding().getWireNumber());
+ meta.writeLong(block.quantizedDataOffset());
+ meta.writeLong(block.quantizedDataSize());
+ }
+ }
+
+ private QuantizedBlock writeFlavorBlock(
+ IndexOutput quantizedVectorData,
+ Flavor flavor,
+ int dimension,
+ int numVectors,
+ FloatVectorSupplier vectors,
+ PreQuantizedSupplier preQuantized)
+ throws IOException {
+
+ OptimizedScalarQuantizer quantizer = flavor.quantizer();
+ float[] zeroCentroid = new float[dimension];
+ float[] normalized = flavor.normalized() ? new float[dimension] : null;
+ byte[] scratch = new byte[encoding.getDiscreteDimensions(dimension)];
+ byte[] packed =
+ switch (encoding) {
+ case UNSIGNED_BYTE, SEVEN_BIT -> scratch;
+ case PACKED_NIBBLE, SINGLE_BIT_QUERY_NIBBLE, DIBIT_QUERY_NIBBLE ->
+ new byte[encoding.getDocPackedLength(scratch.length)];
+ };
+
+ long quantizedDataOffset =
quantizedVectorData.alignFilePointer(Float.BYTES);
+ for (int ord = 0; ord < numVectors; ord++) {
+ // The quantized record is a pure function of the raw vector and the
flavor: a record
+ // already computed by a source segment can be copied instead of
re-quantizing.
+ PreQuantized pre = preQuantized == null ? null : preQuantized.get(ord);
+ if (pre != null && pre.flavor() == flavor &&
pre.values().getScalarEncoding() == encoding) {
+ // NOTE: read the packed bytes before the corrective terms, which are
then served from the
+ // record cached by the read of the packed bytes
+ writeRecord(
+ quantizedVectorData,
+ pre.values().vectorValue(pre.ord()),
+ pre.values().getCorrectiveTerms(pre.ord()));
+ continue;
+ }
+
+ float[] vector = vectors.get(ord);
+ if (flavor.normalized()) {
+ // normalize a copy: the source buffer is shared / owned by the group
+ System.arraycopy(vector, 0, normalized, 0, dimension);
Review Comment:
do we need to make a copy? I think we are free to normalize `vector` in
place?
##########
lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java:
##########
@@ -469,7 +475,7 @@ public int get(int ord) {
@Override
public FieldOrdToGroupOrd copy() {
- return new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd);
+ throw new UnsupportedOperationException("not meant for copying");
Review Comment:
I guess we just never needed this?
--
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]
