diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 0e07088673d5..6ab1babee81e 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -337,6 +337,10 @@ New Features requiring the ICU analysis plugin. Uses a generated switch for fold exceptions with Character.toLowerCase() as the default fallback. (Robert Muir, Costin Leau) +* GITHUB#15979: Add a sandboxed de-duplicating HNSW vector format (DedupHnswVectorsFormat) that stores + each distinct full-precision vector once, shared across all documents and fields that reference it. + (Kaival Parikh) + Improvements --------------------- diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java index 59b766948740..f196df3fe3e6 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java @@ -151,6 +151,10 @@ public void finishMerge() throws IOException { flatVectorsReader.finishMerge(); } + public FlatVectorsReader getFlatVectorsReader() { + return flatVectorsReader; + } + private static IndexInput openDataInput( SegmentReadState state, int versionMeta, diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java index ef3c6f700b6e..e643596cd3ba 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java @@ -236,10 +236,10 @@ public void finish() throws IOException { @Override public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; + long total = SHALLOW_RAM_BYTES_USED + flatVectorWriter.ramBytesUsed(); for (FieldWriter field : fields) { // the field tracks the delegate field usage - total += field.ramBytesUsed(); + total += field.ownRamBytesUsed(); } return total; } @@ -813,13 +813,17 @@ OnHeapHnswGraph getGraph() throws IOException { } } - @Override - public long ramBytesUsed() { - long total = SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + private long ownRamBytesUsed() { + long total = SHALLOW_SIZE; if (hnswGraphBuilder != null) { total += hnswGraphBuilder.getGraph().ramBytesUsed(); } return total; } + + @Override + public long ramBytesUsed() { + return ownRamBytesUsed() + flatFieldVectorsWriter.ramBytesUsed(); + } } } diff --git a/lucene/sandbox/src/java/module-info.java b/lucene/sandbox/src/java/module-info.java index ee9be3227de2..8a3a60d7c544 100644 --- a/lucene/sandbox/src/java/module-info.java +++ b/lucene/sandbox/src/java/module-info.java @@ -22,6 +22,7 @@ requires org.apache.lucene.facet; exports org.apache.lucene.payloads; + exports org.apache.lucene.sandbox.codecs.dedup; exports org.apache.lucene.sandbox.codecs.faiss; exports org.apache.lucene.sandbox.codecs.idversion; exports org.apache.lucene.sandbox.codecs.quantization; @@ -41,5 +42,6 @@ provides org.apache.lucene.codecs.PostingsFormat with org.apache.lucene.sandbox.codecs.idversion.IDVersionPostingsFormat; provides org.apache.lucene.codecs.KnnVectorsFormat with - org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat; + org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat, + org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat; } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java new file mode 100644 index 000000000000..b4bf395538ea --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java @@ -0,0 +1,110 @@ +/* + * 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.ArrayList; +import java.util.List; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.internal.hppc.IntArrayList; +import org.apache.lucene.internal.hppc.ObjectCursor; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * Buffers one field's vectors during flush, de-duplicating them through a shared {@link + * DedupGroup}. + * + * @lucene.experimental + */ +final class DedupFlatFieldVectorsWriter extends FlatFieldVectorsWriter { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(DedupFlatFieldVectorsWriter.class); + + private final DedupGroup group; + private final DocsWithFieldSet docsWithFieldSet; + private final List vectors; + private final IntArrayList fieldOrdToGroupOrd; + private int lastDocID; + private boolean finished; + + DedupFlatFieldVectorsWriter(DedupGroup group) { + this.group = group; + this.docsWithFieldSet = new DocsWithFieldSet(); + this.vectors = new ArrayList<>(); + this.fieldOrdToGroupOrd = new IntArrayList(); + this.lastDocID = -1; + this.finished = false; + } + + @Override + public List getVectors() { + return vectors; + } + + @Override + public DocsWithFieldSet getDocsWithFieldSet() { + return docsWithFieldSet; + } + + IntArrayList getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public void finish() { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + } + + @Override + public boolean isFinished() { + return finished; + } + + @Override + public T copyValue(T vectorValue) { + throw new UnsupportedOperationException(); // handled inside group + } + + @Override + public void addValue(int docID, T vectorValue) throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } else if (docID <= lastDocID) { + throw new IllegalArgumentException( + "docID=" + docID + " not going forwards, indexed lastDocID=" + lastDocID); + } + + lastDocID = docID; + docsWithFieldSet.add(docID); + + ObjectCursor cursor = group.addUnique(vectorValue); + vectors.add(cursor.value); // owned vector value + fieldOrdToGroupOrd.add(cursor.index); // index in group + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + docsWithFieldSet.ramBytesUsed() + + (long) vectors.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF + + fieldOrdToGroupOrd.ramBytesUsed(); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java new file mode 100644 index 000000000000..7d640b0cc83f --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java @@ -0,0 +1,137 @@ +/* + * 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; + +/** + * Flat vector format that stores each distinct vector once. + * + *

Vectors that share the same dimension and encoding form a group. Within a group, an + * identical vector is stored a single time regardless of how many documents (across all fields that + * map to that group) reference it; each field then keeps an {@code fieldOrdToGroupOrd} map from its + * document ordinal to the group ordinal of the shared vector. This is well suited to indexes with + * repeated vectors, e.g. several fields derived from the same embedding, or heavily duplicated + * content. + * + *

.vdd (vector de-dup data) file

+ * + *
    + *
  • For each group, its distinct vectors, aligned to 4 bytes (BYTE) or 64 bytes (FLOAT32 and + * FLOAT16). This a best-effort alignment, because Arm Neoverse machines incur a performance + * penalty in reading data not aligned to 64 bytes. This penalty may be incurred for float + * vectors that do not have a dimension of a multiple of 16 (equivalent to 64 bytes), because + * the alignment will not hold for all vectors in the file. + *
  • For each field: + *
      + *
    • The sparse-encoding data (only when some documents lack the field): DocIds encoded by + * {@link + * org.apache.lucene.codecs.lucene90.IndexedDISI#writeBitSet(org.apache.lucene.search.DocIdSetIterator, + * org.apache.lucene.store.IndexOutput, byte)}, followed by the ordinal-to-doc mapping + * encoded by {@link org.apache.lucene.util.packed.DirectMonotonicWriter}. + *
    • The {@code fieldOrdToGroupOrd} map (aligned to 4 bytes): one entry per document + * ordinal giving the group ordinal of the shared vector, packed by {@link + * org.apache.lucene.util.packed.DirectWriter}. + *
    + *
+ * + *

.vdm (vector de-dup metadata) file

+ * + *

A list of groups, each: + * + *

    + *
  • [int32] group ordinal + *
  • [int32] vector dimension + *
  • [int32] vector encoding ordinal + *
  • [int32] group size (number of distinct vectors) + *
  • [int64] offset to this group's vectors in the .vdd file + *
  • [int64] length of this group's vectors, in bytes + *
+ * + *

terminated by [int32] {@code -1}, then a list of fields, each: + * + *

    + *
  • [int32] field number + *
  • [int32] vector similarity function ordinal + *
  • [int32] vector dimension + *
  • [int32] vector encoding ordinal + *
  • [int32] ordinal of the group holding this field's vectors + *
  • [int32] the number of documents having values for this field + *
  • the sparse-encoding metadata (docs-with-field offset/length and ordToDoc configuration), as + * written by {@link + * org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration#writeStoredMeta} + *
  • [int64] offset to this field's {@code fieldOrdToGroupOrd} map in the .vdd file + *
  • [int64] length of this field's {@code fieldOrdToGroupOrd} map, in bytes + *
+ * + *

also terminated by [int32] {@code -1}. + * + *

Complexity

+ * + *

Let {@code N} be the number of indexed vectors (one per document per field), {@code U} the + * number of distinct vectors, and {@code d} the dimension. De-duplication interns each vector via a + * hash lookup with linear probing, resolving hash collisions with a full equality check. + * + *

    + *
  • Indexing (flush): expected {@code O(N * d)} time (a hash plus occasional equality + * check per vector). Heap is {@code O(U * d)} for the distinct vectors held in the group, + * plus {@code O(N)} for the per-document references and {@code fieldOrdToGroupOrd} entries. + *
  • Merge: expected {@code O(N * d)} time; distinct vectors are written to disk as soon + * as they are first seen rather than buffered, so heap stays {@code O(N)} (the per-field + * {@code fieldOrdToGroupOrd} maps and light per-vector handles) with no {@code O(U * d)} + * term. When a source segment is itself in this format, equality is decided by comparing + * group ordinals in {@code O(1)} without reading the vectors back. + *
  • Reading: both the vectors and the {@code fieldOrdToGroupOrd} map stay off-heap; a + * read resolves a document ordinal to its vector via one extra {@code fieldOrdToGroupOrd} + * lookup. + *
+ * + * @lucene.experimental + */ +final class DedupFlatVectorsFormat extends FlatVectorsFormat { + static final String NAME = "DedupFlatVectorsFormat"; + + static final String META_CODEC_NAME = "DedupFlatVectorsFormatMeta"; + static final String META_EXTENSION = "vdm"; + + static final String VECTOR_DATA_CODEC_NAME = "DedupFlatVectorsFormatVectorData"; + static final String VECTOR_DATA_EXTENSION = "vdd"; + + static final int VERSION_START = 0; + static final int VERSION_CURRENT = VERSION_START; + + private static final DedupFlatVectorsScorer FLAT_VECTORS_SCORER = new DedupFlatVectorsScorer(); + + DedupFlatVectorsFormat() { + super(NAME); + } + + @Override + public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + return new DedupFlatVectorsWriter(state, FLAT_VECTORS_SCORER); + } + + @Override + public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new DedupFlatVectorsReader(state, FLAT_VECTORS_SCORER); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java new file mode 100644 index 000000000000..dfe5454cd0da --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java @@ -0,0 +1,361 @@ +/* + * 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.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT16; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_START; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupFloat16s; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupFloats; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.ReadFieldInfo; +import org.apache.lucene.store.ChecksumIndexInput; +import org.apache.lucene.store.DataAccessHint; +import org.apache.lucene.store.FileDataHint; +import org.apache.lucene.store.FileTypeHint; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.RandomVectorScorer; + +/** + * Reads de-duplicated flat vectors written by {@link DedupFlatVectorsWriter}. Each field exposes a + * view backed by its group's shared vectors and a {@code fieldOrdToGroupOrd} translation map. + * + * @lucene.experimental + */ +final class DedupFlatVectorsReader extends FlatVectorsReader { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(DedupFlatVectorsReader.class); + + private final FlatVectorsScorer vectorsScorer; + private final Map fields; + private final IndexInput vectorData; + + DedupFlatVectorsReader(SegmentReadState state, FlatVectorsScorer vectorsScorer) + throws IOException { + + this.vectorsScorer = vectorsScorer; + this.fields = new HashMap<>(); + + String metaFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + + int versionMeta; + try (ChecksumIndexInput meta = state.directory.openChecksumInput(metaFileName)) { + Throwable priorE = null; + try { + versionMeta = + CodecUtil.checkIndexHeader( + meta, + META_CODEC_NAME, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + readMetaBody(meta, state.fieldInfos); + } catch (Throwable e) { + priorE = e; + throw e; + } finally { + CodecUtil.checkFooter(meta, priorE); + } + } + + this.vectorData = openDataInput(state, versionMeta); + } + + private void readMetaBody(ChecksumIndexInput meta, FieldInfos fieldInfos) throws IOException { + List groupInfos = new ArrayList<>(); + while (true) { + GroupInfo groupInfo = GroupInfo.readFromMeta(meta); + if (groupInfo == null) { + break; + } + groupInfos.add(groupInfo); + } + + while (true) { + ReadFieldInfo fieldInfo = ReadFieldInfo.read(meta); + if (fieldInfo == null) { + break; + } + + FieldInfo info = fieldInfos.fieldInfo(fieldInfo.fieldNumber()); + if (info == null) { + throw new CorruptIndexException("Invalid field number: " + fieldInfo.fieldNumber(), meta); + } else if (fieldInfo.function() != info.getVectorSimilarityFunction()) { + throw new CorruptIndexException( + "Inconsistent vector function: indexed=" + + fieldInfo.function() + + ", actual=" + + info.getVectorSimilarityFunction(), + meta); + } else if (fieldInfo.dimension() != info.getVectorDimension()) { + throw new CorruptIndexException( + "Inconsistent vector dimension: indexed=" + + fieldInfo.dimension() + + ", actual=" + + info.getVectorDimension(), + meta); + } else if (fieldInfo.encoding() != info.getVectorEncoding()) { + throw new CorruptIndexException( + "Inconsistent vector encoding: indexed=" + + fieldInfo.encoding() + + ", actual=" + + info.getVectorEncoding(), + meta); + } + + if (fieldInfo.groupOrd() < 0 || fieldInfo.groupOrd() >= groupInfos.size()) { + throw new CorruptIndexException( + "Invalid groupId=" + fieldInfo.groupOrd() + ", numGroups=" + groupInfos.size(), meta); + } + + GroupInfo groupInfo = groupInfos.get(fieldInfo.groupOrd()); + if (fieldInfo.dimension() != groupInfo.dimension()) { + throw new CorruptIndexException( + "Vector dimension mismatch: field=" + + fieldInfo.dimension() + + ", group=" + + groupInfo.dimension(), + meta); + } else if (fieldInfo.encoding() != groupInfo.encoding()) { + throw new CorruptIndexException( + "Vector encoding mismatch: field=" + + fieldInfo.encoding() + + ", group=" + + groupInfo.encoding(), + meta); + } + + fields.put(info.name, new FieldEntry(fieldInfo, groupInfo)); + } + } + + private static IndexInput openDataInput(SegmentReadState state, int versionMeta) + throws IOException { + + String fileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); + + IOContext.FileOpenHint[] hints = { + FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM + }; + IOContext context = state.context.withHints(hints); + + IndexInput in = null; + boolean success = false; + try { + in = state.directory.openInput(fileName, context); + int versionVectorData = + CodecUtil.checkIndexHeader( + in, + VECTOR_DATA_CODEC_NAME, + VERSION_START, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + if (versionMeta != versionVectorData) { + throw new CorruptIndexException( + "Format versions mismatch: meta=" + + versionMeta + + ", " + + VECTOR_DATA_CODEC_NAME + + "=" + + versionVectorData, + in); + } + CodecUtil.retrieveChecksum(in); + success = true; + return in; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(in); + } + } + } + + @Override + public FlatVectorsScorer getFlatVectorScorer(String field) { + return vectorsScorer; + } + + // package-private for testing + FieldEntry getEntry(String field, VectorEncoding expected) { + FieldEntry entry = fields.get(field); + if (entry == null) { + throw new IllegalArgumentException("field=" + field + " not found"); + } else if (entry.fieldInfo.encoding() != expected) { + throw new IllegalArgumentException("field=" + field + " not indexed as " + expected); + } + return entry; + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException { + FieldEntry entry = getEntry(field, FLOAT32); + FloatVectorValues vectorValues = getFloatVectorValues(entry); + return vectorsScorer.getRandomVectorScorer(entry.fieldInfo.function(), vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException { + FieldEntry entry = getEntry(field, BYTE); + ByteVectorValues vectorValues = getByteVectorValues(entry); + return vectorsScorer.getRandomVectorScorer(entry.fieldInfo.function(), vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer(String field, short[] target) throws IOException { + FieldEntry entry = getEntry(field, FLOAT16); + Float16VectorValues vectorValues = getFloat16VectorValues(entry); + return vectorsScorer.getRandomVectorScorer(entry.fieldInfo.function(), vectorValues, target); + } + + @Override + public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException { + CodecUtil.checksumEntireFile(vectorData, merge); + } + + private FloatVectorValues getFloatVectorValues(FieldEntry entry) throws IOException { + return loadDedupFloats( + vectorsScorer, + entry.fieldInfo.function(), + entry.fieldInfo.ordToDoc(), + entry.fieldInfo.dimension(), + entry.groupInfo.groupNumVectors(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); + } + + @Override + public FloatVectorValues getFloatVectorValues(String field) throws IOException { + return getFloatVectorValues(getEntry(field, FLOAT32)); + } + + private ByteVectorValues getByteVectorValues(FieldEntry entry) throws IOException { + return loadDedupBytes( + vectorsScorer, + entry.fieldInfo.function(), + entry.fieldInfo.ordToDoc(), + entry.fieldInfo.dimension(), + entry.groupInfo.groupNumVectors(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return getByteVectorValues(getEntry(field, BYTE)); + } + + private Float16VectorValues getFloat16VectorValues(FieldEntry entry) throws IOException { + return loadDedupFloat16s( + vectorsScorer, + entry.fieldInfo.function(), + entry.fieldInfo.ordToDoc(), + entry.fieldInfo.dimension(), + entry.groupInfo.groupNumVectors(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); + } + + @Override + public Float16VectorValues getFloat16VectorValues(String field) throws IOException { + return getFloat16VectorValues(getEntry(field, FLOAT16)); + } + + @Override + public FlatVectorsReader getMergeInstance() { + // TODO: Can we improve performance using sequential IO + avoiding read-backs? One way is to + // de-duplicate only within a document, allowing for cleaner sort of vectors during flush + + // avoid read-backs for full equality checks during merge. + return this; + } + + @Override + public void finishMerge() { + // TODO: Converse of getMergeInstance() + } + + @Override + public void close() throws IOException { + IOUtils.close(vectorData); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + fields.size() * FieldEntry.SHALLOW_SIZE; + } + + @Override + public Map getOffHeapByteSize(FieldInfo fieldInfo) { + FieldEntry entry = fields.get(fieldInfo.name); + if (entry == null) { + return Map.of(); + } + // TODO: This is an over-estimation. + return Map.of( + VECTOR_DATA_EXTENSION, + entry.fieldInfo.fieldOrdToGroupOrdSize() + entry.groupInfo.vectorDataSize()); + } + + record FieldEntry(ReadFieldInfo fieldInfo, GroupInfo groupInfo) { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(FieldEntry.class) + + RamUsageEstimator.shallowSizeOfInstance(ReadFieldInfo.class) + + RamUsageEstimator.shallowSizeOfInstance(GroupInfo.class); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java new file mode 100644 index 000000000000..63da8b2d950b --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java @@ -0,0 +1,187 @@ +/* + * 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.sandbox.codecs.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; +import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer; + +/** + * Scorer for de-duplicated vectors. Performs doc operations on the original vector values, but + * delegates vector operations to the underlying {@link DedupVectorValues#getGroupView()}, mapping + * document ordinals to group ordinals via {@link DedupVectorValues#getFieldOrdToGroupOrd()}. + * + * @lucene.experimental + */ +final class DedupFlatVectorsScorer implements FlatVectorsScorer { + private static final FlatVectorsScorer SCORER = + FlatVectorScorerUtil.getLucene99FlatVectorsScorer(); + + @Override + public RandomVectorScorerSupplier getRandomVectorScorerSupplier( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorerSupplier fieldView = + SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); + RandomVectorScorerSupplier groupView = + SCORER.getRandomVectorScorerSupplier(similarityFunction, dedupValues.getGroupView()); + return new RandomVectorScorerSupplierImpl( + fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); + } + return SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, float[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer fieldView = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, byte[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer fieldView = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, short[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer fieldView = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + private record RandomVectorScorerSupplierImpl( + RandomVectorScorerSupplier fieldView, + RandomVectorScorerSupplier groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) + implements RandomVectorScorerSupplier { + + @Override + public UpdateableRandomVectorScorer scorer() throws IOException { + return new UpdateableRandomVectorScorerImpl( + fieldView.scorer(), groupView.scorer(), fieldOrdToGroupOrd); + } + + @Override + public RandomVectorScorerSupplier copy() throws IOException { + return new RandomVectorScorerSupplierImpl( + fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + } + + private static class RandomVectorScorerImpl implements RandomVectorScorer { + private final RandomVectorScorer fieldView; + private final RandomVectorScorer groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + RandomVectorScorerImpl( + RandomVectorScorer fieldView, + RandomVectorScorer groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return fieldView.getAcceptOrds(acceptDocs); + } + + @Override + public float score(int node) throws IOException { + return groupView.score(fieldOrdToGroupOrd.get(node)); + } + + @Override + public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException { + if (scratch.length < nodes.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, nodes.length); + } + for (int i = 0; i < numNodes; i++) { + scratch[i] = fieldOrdToGroupOrd.get(nodes[i]); + } + return groupView.bulkScore(scratch, scores, numNodes); + } + + @Override + public int maxOrd() { + return fieldView.maxOrd(); + } + } + + private static final class UpdateableRandomVectorScorerImpl extends RandomVectorScorerImpl + implements UpdateableRandomVectorScorer { + private final UpdateableRandomVectorScorer groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + + UpdateableRandomVectorScorerImpl( + UpdateableRandomVectorScorer fieldView, + UpdateableRandomVectorScorer groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + super(fieldView, groupView, fieldOrdToGroupOrd); + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + } + + @Override + public void setScoringOrdinal(int node) throws IOException { + groupView.setScoringOrdinal(fieldOrdToGroupOrd.get(node)); + } + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java new file mode 100644 index 000000000000..d8020fb2e855 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java @@ -0,0 +1,155 @@ +/* + * 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.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; + +import java.io.IOException; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * Writes de-duplicated flat vectors. A single instance is used for either flushing buffered vectors + * or merging existing segments (never both), delegating to {@link DedupFlushContext} or {@link + * DedupMergeContext} accordingly. + * + * @lucene.experimental + */ +final class DedupFlatVectorsWriter extends FlatVectorsWriter { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(DedupFlatVectorsWriter.class); + + private final IndexOutput meta; + private final IndexOutput vectorData; + private boolean finished; + + private final DedupFlushContext flushContext; + private boolean usedForFlush; + + private final DedupMergeContext mergeContext; + private boolean usedForMerge; + + DedupFlatVectorsWriter(SegmentWriteState state, FlatVectorsScorer vectorsScorer) + throws IOException { + super(vectorsScorer); + + this.finished = false; + this.flushContext = new DedupFlushContext(); + this.usedForFlush = false; + this.mergeContext = new DedupMergeContext(); + this.usedForMerge = false; + + String metaFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); + String vectorDataFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); + + boolean success = false; + try { + this.meta = state.directory.createOutput(metaFileName, state.context); + this.vectorData = state.directory.createOutput(vectorDataFileName, state.context); + CodecUtil.writeIndexHeader( + meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + CodecUtil.writeIndexHeader( + vectorData, + VECTOR_DATA_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(this); + } + } + } + + @Override + public FlatFieldVectorsWriter addField(FieldInfo fieldInfo) { + if (usedForMerge) { + throw new IllegalStateException("already used for merge"); + } + usedForFlush = true; + + return flushContext.addField(fieldInfo); + } + + @Override + public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { + if (usedForMerge) { + throw new IllegalStateException("already used for merge"); + } + usedForFlush = true; + + flushContext.flush(meta, vectorData, maxDoc, sortMap); + } + + @Override + public void finish() throws IOException { + if (finished) { + throw new IllegalStateException("already finished"); + } + finished = true; + + if (usedForMerge) { + finishMerge(); + } + + CodecUtil.writeFooter(meta); + CodecUtil.writeFooter(vectorData); + } + + @Override + public void mergeOneFlatVectorField(FieldInfo fieldInfo, MergeState mergeState) + throws IOException { + if (usedForFlush) { + throw new IllegalStateException("already used for flush"); + } + usedForMerge = true; + + mergeContext.addField(fieldInfo, mergeState); + } + + private void finishMerge() throws IOException { + mergeContext.finish(meta, vectorData); + } + + @Override + public void close() throws IOException { + IOUtils.close(meta, vectorData); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + flushContext.ramBytesUsed() + mergeContext.ramBytesUsed(); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java new file mode 100644 index 000000000000..59ad5465b83e --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java @@ -0,0 +1,306 @@ +/* + * 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 java.nio.ByteOrder.LITTLE_ENDIAN; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndMarker; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.FloatBuffer; +import java.nio.ShortBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.internal.hppc.IntArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdMappedArrayList; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * Buffers vectors added during a flush and de-duplicates them in memory. Fields sharing a {@link + * DedupUtil.GroupKey} intern into the same {@link DedupGroup}; on {@link #flush} each group's + * distinct vectors are written once, followed by per-field metadata mapping document ordinals to + * group ordinals. + * + * @lucene.experimental + */ +final class DedupFlushContext implements Accountable { + private final Map> groups; + private final List fieldDataList; + + DedupFlushContext() { + this.groups = new HashMap<>(); + this.fieldDataList = new ArrayList<>(); + } + + private static DedupGroup getGroup(GroupKey groupKey) { + int dimension = groupKey.dimension(); + return switch (groupKey.encoding()) { + case BYTE -> new ByteGroup(dimension); + case FLOAT32 -> new FloatGroup(dimension); + case FLOAT16 -> new Float16Group(dimension); + }; + } + + @Override + public long ramBytesUsed() { + long total = 0; + for (DedupGroup group : groups.values()) { + total += group.ramBytesUsed(); + } + for (FieldData data : fieldDataList) { + total += data.ramBytesUsed(); + } + return total; + } + + FlatFieldVectorsWriter addField(FieldInfo fieldInfo) { + GroupKey groupKey = new GroupKey(fieldInfo); + DedupGroup group = groups.computeIfAbsent(groupKey, DedupFlushContext::getGroup); + DedupFlatFieldVectorsWriter fieldVectorsWriter = new DedupFlatFieldVectorsWriter<>(group); + + fieldDataList.add(new FieldData(fieldInfo, groupKey, fieldVectorsWriter)); + return fieldVectorsWriter; + } + + void flush(IndexOutput meta, IndexOutput vectorData, int maxDoc, Sorter.DocMap sortMap) + throws IOException { + + Map groupOrds = new HashMap<>(); + + int groupOrd = 0; + for (Map.Entry> entry : groups.entrySet()) { + GroupKey groupKey = entry.getKey(); + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); + + DedupGroup group = entry.getValue(); + int groupNumVectors = group.numVectors(); + long vectorDataOffset = alignBytes(vectorData, encoding); + + // TODO: Write in sorted order for faster merge? (with sequential IO) + for (int ord = 0; ord < groupNumVectors; ord++) { + byte[] bytes = group.serialize(ord); + vectorData.writeBytes(bytes, bytes.length); + } + long vectorDataSize = vectorData.getFilePointer() - vectorDataOffset; + + GroupInfo groupInfo = + new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + groupInfo.write(meta); + + groupOrds.put(groupKey, groupOrd); + groupOrd++; + } + + writeEndMarker(meta); + + for (FieldData fieldData : fieldDataList) { + fieldData.fieldWriter.finish(); + + IntArrayList fieldOrdToGroupOrd = fieldData.fieldWriter.getFieldOrdToGroupOrd(); + int vectorCount = fieldOrdToGroupOrd.elementsCount; + + DocsWithFieldSet docs; + FieldOrdToGroupOrd fieldOrdToGroupOrdFinal; + if (sortMap == null) { + docs = fieldData.fieldWriter.getDocsWithFieldSet(); + fieldOrdToGroupOrdFinal = new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); + } else { + DocsWithFieldSet oldDocs = fieldData.fieldWriter.getDocsWithFieldSet(); + docs = new DocsWithFieldSet(); + int[] new2OldOrd = new int[vectorCount]; + KnnVectorsWriter.mapOldOrdToNewOrd(oldDocs, sortMap, null, new2OldOrd, docs); + fieldOrdToGroupOrdFinal = + new FieldOrdToGroupOrdMappedArrayList(new2OldOrd, fieldOrdToGroupOrd); + } + + writeFieldInfo( + meta, + vectorData, + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + vectorCount, + maxDoc, + docs, + fieldOrdToGroupOrdFinal); + } + + writeEndMarker(meta); + } + + static final class ByteGroup extends DedupGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ByteGroup.class); + private final long ramBytesPerVector; + + ByteGroup(int dimension) { + ramBytesPerVector = + RamUsageEstimator.NUM_BYTES_OBJECT_REF + + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + + dimension; + } + + @Override + public long hash(byte[] vector) { + return hashBytes(vector); + } + + @Override + public boolean equals(byte[] vector, byte[] other) { + return Arrays.equals(vector, other); + } + + @Override + public byte[] copy(byte[] vectorValue) { + return vectorValue.clone(); + } + + @Override + byte[] serialize(int ord) { + return get(ord); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; + } + } + + static final class FloatGroup extends DedupGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(FloatGroup.class); + + private final long ramBytesPerVector; + private final byte[] bytes; + private final FloatBuffer buffer; + + FloatGroup(int dimension) { + int length = dimension * Float.BYTES; + this.ramBytesPerVector = + RamUsageEstimator.NUM_BYTES_OBJECT_REF + + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + + length; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asFloatBuffer(); + } + + @Override + public long hash(float[] vector) { + buffer.put(0, vector); + return hashBytes(bytes); + } + + @Override + public boolean equals(float[] vector, float[] other) { + return Arrays.equals(vector, other); + } + + @Override + public float[] copy(float[] vectorValue) { + return vectorValue.clone(); + } + + @Override + byte[] serialize(int ord) { + buffer.put(0, get(ord)); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; + } + } + + static final class Float16Group extends DedupGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Group.class); + + private final long ramBytesPerVector; + private final byte[] bytes; + private final ShortBuffer buffer; + + Float16Group(int dimension) { + int length = dimension * Short.BYTES; + this.ramBytesPerVector = + RamUsageEstimator.NUM_BYTES_OBJECT_REF + + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + + length; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asShortBuffer(); + } + + @Override + public long hash(short[] vector) { + buffer.put(0, vector); + return hashBytes(bytes); + } + + @Override + public boolean equals(short[] vector, short[] other) { + return Arrays.equals(vector, other); + } + + @Override + public short[] copy(short[] vectorValue) { + return vectorValue.clone(); + } + + @Override + byte[] serialize(int ord) { + buffer.put(0, get(ord)); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; + } + } + + private record FieldData( + FieldInfo fieldInfo, GroupKey groupKey, DedupFlatFieldVectorsWriter fieldWriter) + implements Accountable { + + @Override + public long ramBytesUsed() { + return fieldWriter.ramBytesUsed(); + } + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java new file mode 100644 index 000000000000..e2b83cf67767 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java @@ -0,0 +1,105 @@ +/* + * 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.ArrayList; +import java.util.List; +import org.apache.lucene.internal.hppc.LongIntHashMap; +import org.apache.lucene.internal.hppc.ObjectCursor; +import org.apache.lucene.util.Accountable; + +/** + * Interns vectors so that each distinct value is stored once. {@link #addUnique} returns the group + * ordinal for a vector, adding it (via {@link #copy}) only if not already present. Callers with the + * same {@code (dimension, encoding)} share a group, so an identical vector across fields is stored + * a single time. + * + *

Not thread-safe; a group is confined to the writer that created it. + * + * @lucene.experimental + */ +abstract sealed class DedupGroup implements Accountable + permits DedupFlushContext.ByteGroup, + DedupFlushContext.FloatGroup, + DedupFlushContext.Float16Group, + DedupMergeContext.DedupMergeGroup { + + /** + * Map for vector hash -> group ord. Only a hint and not the ground truth due to possibility of + * hash collisions, where a full equality check must be performed. + */ + private final LongIntHashMap hashToOrdHint; + + private final List vectors; + + private final ObjectCursor current; // reuse from addUnique + + DedupGroup() { + this.hashToOrdHint = new LongIntHashMap(); + this.vectors = new ArrayList<>(); + this.current = new ObjectCursor<>(); + } + + abstract long hash(T vector) throws IOException; + + abstract boolean equals(T vector, T other) throws IOException; + + abstract T copy(T vector); + + abstract byte[] serialize(int ord) throws IOException; + + int numVectors() { + return vectors.size(); + } + + T get(int ord) { + return vectors.get(ord); + } + + ObjectCursor addUnique(T vectorValue) throws IOException { + final int groupOrd; + final T ownedVector; + for (long hash = hash(vectorValue); ; hash++) { // linear probing + int ordIndex = hashToOrdHint.indexOf(hash); + if (ordIndex < 0) { // not found + groupOrd = vectors.size(); + ownedVector = copy(vectorValue); // only for unique vectors + hashToOrdHint.put(hash, groupOrd); + vectors.add(ownedVector); + break; + } else { + int ord = hashToOrdHint.indexGet(ordIndex); + T other = vectors.get(ord); + if (equals(vectorValue, other)) { + groupOrd = ord; + ownedVector = other; + break; + } + // else continue probing + } + } + current.index = groupOrd; + current.value = ownedVector; + return current; + } + + @Override + public long ramBytesUsed() { + return hashToOrdHint.ramBytesUsed(); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java new file mode 100644 index 000000000000..bff40f470558 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java @@ -0,0 +1,234 @@ +/* + * 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.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.HNSW_GRAPH_THRESHOLD; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_BEAM_WIDTH; +import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.MAXIMUM_MAX_CONN; + +import java.io.IOException; +import java.util.concurrent.ExecutorService; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsWriter; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.MergeScheduler; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.search.TaskExecutor; +import org.apache.lucene.util.hnsw.HnswGraph; + +/** + * An HNSW vector format that de-duplicates raw vectors. + * + *

Graph construction and search are identical to {@link + * org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat}. A {@link DedupFlatVectorsFormat} is + * used for the flat vector storage, which stores each distinct vector exactly once, shared across + * all documents that reference it. + * + *

This format is suitable for high-performance filtered vector search when filter information is + * available at indexing time. In addition to the primary vector field, the user creates separate + * fields for each filter value (e.g. product categories in an e-commerce search engine) to build + * dedicated HNSW graphs that share the same raw vector storage. The responsibility of searching + * the right field at query-time lies on the user. + * + *

This scheme allows for more efficient search than query-time pre-filtering (i.e. {@link + * org.apache.lucene.search.AcceptDocs} derived from a {@link org.apache.lucene.search.Query} {@code + * filter}) at the expense of slower indexing and larger indexes due to additional HNSW graphs. + * + *

If you customize this format, be sure to share the same instance of the underlying + * {@link DedupFlatVectorsFormat} to de-duplicate raw vectors correctly. + * + * @lucene.experimental + */ +public final class DedupHnswVectorsFormat extends KnnVectorsFormat { + private static final String NAME = "DedupHnswVectorsFormat"; + + /** + * Controls how many of the nearest neighbor candidates are connected to the new node. Defaults to + * {@link org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat#DEFAULT_MAX_CONN}. See + * {@link HnswGraph} for more details. + */ + private final int maxConn; + + /** + * The number of candidate neighbors to track while searching the graph for each newly inserted + * node. Defaults to {@link + * org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat#DEFAULT_BEAM_WIDTH}. See {@link + * HnswGraph} for details. + */ + private final int beamWidth; + + /** The format for storing, reading, and merging vectors on disk. */ + private static final FlatVectorsFormat FORMAT = new DedupFlatVectorsFormat(); + + private final int numMergeWorkers; + private final TaskExecutor mergeExec; + + /** + * The threshold to use to bypass HNSW graph building for tiny segments in terms of k for a graph + * i.e. number of docs to match the query (default is {@link + * org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD}). + * + *

    + *
  • 0 indicates that the graph is always built. + *
  • Positive values require that many estimated visited nodes before a graph is built. + *
  • Negative values aren't allowed. + *
+ */ + private final int tinySegmentsThreshold; + + /** Constructs a format using default graph construction parameters */ + public DedupHnswVectorsFormat() { + this( + DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD); + } + + /** + * Constructs a format using the given graph construction parameters. + * + * @param maxConn the maximum number of connections to a node in the HNSW graph + * @param beamWidth the size of the queue maintained during graph construction. + */ + public DedupHnswVectorsFormat(int maxConn, int beamWidth) { + this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD); + } + + /** + * Constructs a format using the given graph construction parameters. + * + * @param maxConn the maximum number of connections to a node in the HNSW graph + * @param beamWidth the size of the queue maintained during graph construction. + * @param tinySegmentsThreshold the expected number of vector operations to return k nearest + * neighbors of the current graph size + */ + public DedupHnswVectorsFormat(int maxConn, int beamWidth, int tinySegmentsThreshold) { + this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, tinySegmentsThreshold); + } + + /** + * Constructs a format using the given graph construction parameters. + * + * @param maxConn the maximum number of connections to a node in the HNSW graph + * @param beamWidth the size of the queue maintained during graph construction. + * @param numMergeWorkers number of workers (threads) that will be used when doing merge. If + * larger than 1, a non-null {@link ExecutorService} must be passed as mergeExec + * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are + * generated by this format to do the merge. If null, the configured {@link + * MergeScheduler#getIntraMergeExecutor(MergePolicy.OneMerge)} is used. + */ + public DedupHnswVectorsFormat( + int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) { + this(maxConn, beamWidth, numMergeWorkers, mergeExec, HNSW_GRAPH_THRESHOLD); + } + + /** + * Constructs a format using the given graph construction parameters. + * + * @param maxConn the maximum number of connections to a node in the HNSW graph + * @param beamWidth the size of the queue maintained during graph construction. + * @param numMergeWorkers number of workers (threads) that will be used when doing merge. If + * larger than 1, a non-null {@link ExecutorService} must be passed as mergeExec + * @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are + * generated by this format to do the merge. If null, the configured {@link + * MergeScheduler#getIntraMergeExecutor(MergePolicy.OneMerge)} is used. + * @param tinySegmentsThreshold the expected number of vector operations to return k nearest + * neighbors of the current graph size + */ + public DedupHnswVectorsFormat( + int maxConn, + int beamWidth, + int numMergeWorkers, + ExecutorService mergeExec, + int tinySegmentsThreshold) { + super(NAME); + if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) { + throw new IllegalArgumentException( + "maxConn must be positive and less than or equal to " + + MAXIMUM_MAX_CONN + + "; maxConn=" + + maxConn); + } + if (beamWidth <= 0 || beamWidth > MAXIMUM_BEAM_WIDTH) { + throw new IllegalArgumentException( + "beamWidth must be positive and less than or equal to " + + MAXIMUM_BEAM_WIDTH + + "; beamWidth=" + + beamWidth); + } + this.maxConn = maxConn; + this.beamWidth = beamWidth; + this.tinySegmentsThreshold = tinySegmentsThreshold; + if (numMergeWorkers == 1 && mergeExec != null) { + throw new IllegalArgumentException( + "No executor service is needed as we'll use single thread to merge"); + } + this.numMergeWorkers = numMergeWorkers; + if (mergeExec != null) { + this.mergeExec = new TaskExecutor(mergeExec); + } else { + this.mergeExec = null; + } + } + + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + // TODO: Can we have an HNSW writer that uses de-duplication information to speed up graph + // construction? + return new Lucene99HnswVectorsWriter( + state, + maxConn, + beamWidth, + FORMAT, + FORMAT.fieldsWriter(state), + numMergeWorkers, + mergeExec, + tinySegmentsThreshold); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return new Lucene99HnswVectorsReader(state, FORMAT.fieldsReader(state)); + } + + @Override + public int getMaxDimensions(String fieldName) { + return DEFAULT_MAX_DIMENSIONS; + } + + @Override + public String toString() { + return NAME + + "(name=" + + NAME + + ", maxConn=" + + maxConn + + ", beamWidth=" + + beamWidth + + ", tinySegmentsThreshold=" + + tinySegmentsThreshold + + ", flatVectorFormat=" + + FORMAT + + ")"; + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java new file mode 100644 index 000000000000..55ca0e2cb1ac --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java @@ -0,0 +1,417 @@ +/* + * 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 java.nio.ByteOrder.LITTLE_ENDIAN; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndMarker; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.FloatBuffer; +import java.nio.ShortBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.DocIDMerger; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.internal.hppc.IntArrayList; +import org.apache.lucene.internal.hppc.ObjectCursor; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.IOSupplier; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * Merges de-duplicated flat vectors from several segments. Fields sharing a {@link + * DedupUtil.GroupKey} are merged into one group: their vectors are streamed in merged doc order + * through a {@link DedupGroup}, so a vector is written the first time it is seen and later + * occurrences (within or across fields) reuse that group ordinal. Vectors originating from a dedup + * source are compared by ordinal to avoid reading them back. + * + * @lucene.experimental + */ +final class DedupMergeContext implements Accountable { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(DedupMergeContext.class); + private final List fieldDataList; + + DedupMergeContext() { + this.fieldDataList = new ArrayList<>(); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + fieldDataList.size() * FieldData.SHALLOW_SIZE; + } + + void addField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + fieldDataList.add( + new FieldData( + fieldInfo, + new GroupKey(fieldInfo), + new DocsWithFieldSet(), + new IntArrayList(), + getVectorMerger(fieldInfo, mergeState), + mergeState.segmentInfo.maxDoc())); + } + + void finish(IndexOutput meta, IndexOutput vectorData) throws IOException { + + // Evaluate compatible fields together for correct de-duplication + Map> fieldGroups = + fieldDataList.stream().collect(Collectors.groupingBy(FieldData::groupKey)); + + Map groupOrds = new HashMap<>(); + int groupOrd = 0; + for (Map.Entry> entry : fieldGroups.entrySet()) { + GroupKey groupKey = entry.getKey(); + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); + + long vectorDataOffset = alignBytes(vectorData, encoding); + + DedupMergeGroup mergeGroup = + switch (encoding) { + case BYTE -> new ByteGroup(); + case FLOAT32 -> new FloatGroup(dimension); + case FLOAT16 -> new Float16Group(dimension); + }; + + for (FieldData fieldData : entry.getValue()) { + mergeGroup.processField(fieldData, vectorData); + } + + int groupNumVectors = mergeGroup.numVectors(); + long vectorDataSize = vectorData.getFilePointer() - vectorDataOffset; + + GroupInfo groupInfo = + new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + groupInfo.write(meta); + + groupOrds.put(groupKey, groupOrd); + groupOrd++; + } + + writeEndMarker(meta); + + for (FieldData fieldData : fieldDataList) { + writeFieldInfo( + meta, + vectorData, + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + fieldData.fieldOrdToGroupOrd.elementsCount, + fieldData.maxDoc, + fieldData.docsWithFieldSet, + new FieldOrdToGroupOrdArrayList(fieldData.fieldOrdToGroupOrd)); + } + + writeEndMarker(meta); + } + + abstract static sealed class DedupMergeGroup extends DedupGroup { + abstract T vectorFrom(Sub sub); + + void processField(FieldData fieldData, IndexOutput vectorData) throws IOException { + @SuppressWarnings("unchecked") + DocIDMerger> merger = (DocIDMerger>) fieldData.merger; + + // iterate merged docs one-by-one + for (Sub next = merger.next(); next != null; next = merger.next()) { + T vector = vectorFrom(next); + int groupNumVectors = numVectors(); + + // add vector to group + ObjectCursor cursor = addUnique(vector); + if (cursor.index == groupNumVectors) { // new addition + // already on-heap, write immediately to avoid another IO read + byte[] bytes = serialize(groupNumVectors); + vectorData.writeBytes(bytes, bytes.length); + } + + // record hit and ord in group + fieldData.docsWithFieldSet.add(next.mappedDocID); + fieldData.fieldOrdToGroupOrd.add(cursor.index); + } + } + } + + private record ByteVector(ByteVectorValues values, int ord) implements IOSupplier { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ByteVector.class); + + @Override + public byte[] get() throws IOException { + return values.vectorValue(ord); + } + } + + private static final class ByteGroup extends DedupMergeGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ByteGroup.class); + + @Override + ByteVector vectorFrom(Sub sub) { + return new ByteVector(sub.values, sub.iterator.index()); + } + + @Override + public long hash(ByteVector vector) throws IOException { + return hashBytes(vector.get()); + } + + @Override + public boolean equals(ByteVector vector, ByteVector other) throws IOException { + // Fast path: two docs from the same dedup source share a vector iff they map to the same + // group ordinal, so we can compare ordinals without reading the vectors back. + if (vector.values == other.values && vector.values instanceof DedupVectorValues dedup) { + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); + } + byte[] a = vector.get(); + if (vector.values == other.values) { + a = a.clone(); // same reader reuses one buffer; copy before reading the other vector + } + return Arrays.equals(a, other.get()); + } + + @Override + public ByteVector copy(ByteVector vectorValue) { + return vectorValue; + } + + @Override + byte[] serialize(int ord) throws IOException { + return get(ord).get(); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ByteVector.SHALLOW_SIZE; + } + } + + private record FloatVector(FloatVectorValues values, int ord) implements IOSupplier { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(FloatVector.class); + + @Override + public float[] get() throws IOException { + return values.vectorValue(ord); + } + } + + private static final class FloatGroup extends DedupMergeGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(FloatGroup.class); + + private final byte[] bytes; + private final FloatBuffer buffer; + + FloatGroup(int dimension) { + int length = dimension * Float.BYTES; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asFloatBuffer(); + } + + @Override + FloatVector vectorFrom(Sub sub) { + return new FloatVector(sub.values, sub.iterator.index()); + } + + @Override + public long hash(FloatVector vector) throws IOException { + buffer.put(0, vector.get()); + return hashBytes(bytes); + } + + @Override + public boolean equals(FloatVector vector, FloatVector other) throws IOException { + // Fast path: two docs from the same dedup source share a vector iff they map to the same + // group ordinal, so we can compare ordinals without reading the vectors back. + if (vector.values == other.values && vector.values instanceof DedupVectorValues dedup) { + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); + } + float[] a = vector.get(); + if (vector.values == other.values) { + a = a.clone(); // same reader reuses one buffer; copy before reading the other vector + } + return Arrays.equals(a, other.get()); + } + + @Override + public FloatVector copy(FloatVector vectorValue) { + return vectorValue; + } + + @Override + byte[] serialize(int ord) throws IOException { + buffer.put(0, get(ord).get()); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * FloatVector.SHALLOW_SIZE; + } + } + + private record Float16Vector(Float16VectorValues values, int ord) implements IOSupplier { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Vector.class); + + @Override + public short[] get() throws IOException { + return values.vectorValue(ord); + } + } + + private static final class Float16Group + extends DedupMergeGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Group.class); + + private final byte[] bytes; + private final ShortBuffer buffer; + + Float16Group(int dimension) { + int length = dimension * Short.BYTES; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asShortBuffer(); + } + + @Override + Float16Vector vectorFrom(Sub sub) { + return new Float16Vector(sub.values, sub.iterator.index()); + } + + @Override + public long hash(Float16Vector vector) throws IOException { + buffer.put(0, vector.get()); + return hashBytes(bytes); + } + + @Override + public boolean equals(Float16Vector vector, Float16Vector other) throws IOException { + // Fast path: two docs from the same dedup source share a vector iff they map to the same + // group ordinal, so we can compare ordinals without reading the vectors back. + if (vector.values == other.values && vector.values instanceof DedupVectorValues dedup) { + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); + } + short[] a = vector.get(); + if (vector.values == other.values) { + a = a.clone(); // same reader reuses one buffer; copy before reading the other vector + } + return Arrays.equals(a, other.get()); + } + + @Override + public Float16Vector copy(Float16Vector vectorValue) { + return vectorValue; + } + + @Override + byte[] serialize(int ord) throws IOException { + buffer.put(0, get(ord).get()); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * Float16Vector.SHALLOW_SIZE; + } + } + + private record FieldData( + FieldInfo fieldInfo, + GroupKey groupKey, + DocsWithFieldSet docsWithFieldSet, + IntArrayList fieldOrdToGroupOrd, + DocIDMerger merger, + int maxDoc) { + + static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(FieldData.class); + } + + private static class Sub extends DocIDMerger.Sub { + private final T values; + private final KnnVectorValues.DocIndexIterator iterator; + + Sub(MergeState.DocMap docMap, T values) { + super(docMap); + this.values = values; + iterator = values.iterator(); + } + + @Override + public int nextDoc() throws IOException { + return iterator.nextDoc(); + } + } + + private static DocIDMerger> getVectorMerger( + FieldInfo fieldInfo, MergeState mergeState) throws IOException { + + List> subs = new ArrayList<>(); + for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { + + if (mergeState.knnVectorsReaders[i] == null + || mergeState.fieldInfos[i].fieldInfo(fieldInfo.name) == null + || mergeState.fieldInfos[i].fieldInfo(fieldInfo.name).hasVectorValues() == false) { + continue; + } + + KnnVectorValues vectorValues = + switch (fieldInfo.getVectorEncoding()) { + case BYTE -> mergeState.knnVectorsReaders[i].getByteVectorValues(fieldInfo.name); + case FLOAT32 -> mergeState.knnVectorsReaders[i].getFloatVectorValues(fieldInfo.name); + case FLOAT16 -> mergeState.knnVectorsReaders[i].getFloat16VectorValues(fieldInfo.name); + }; + + if (vectorValues == null) { + continue; + } + + subs.add(new Sub<>(mergeState.docMaps[i], vectorValues)); + } + + return DocIDMerger.of(subs, mergeState.needsIndexSort); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java new file mode 100644 index 000000000000..594fb6abe730 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java @@ -0,0 +1,193 @@ +/* + * 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.util.StringHelper.GOOD_FAST_HASH_SEED; +import static org.apache.lucene.util.StringHelper.murmurhash3_x64_128; + +import java.io.IOException; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.packed.DirectWriter; + +/** + * Shared helpers for the de-duplicating flat format: reading / writing field and group metadata, + * vector hashing and alignment, and the {@link DedupVectorValues} views used on the read path. + * + * @lucene.experimental + */ +final class DedupUtil { + + private static final int ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + + private static final int END_MARKER = -1; + + /** Alignment bytes on disk for fieldOrdToGroupOrd. */ + private static final int FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES = 4; + + // TODO: This is the number of bits used to write each group ordinal in the index-backed per-field + // FieldOrdToGroupOrd mapping. Evaluate using fewer bits to reduce index size, at the expense of + // costlier lookups. + static final int FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE = 32; + + /** Initial allocation size for internal re-used int[] scratch buffers. */ + static final int SCRATCH_INITIAL_SIZE = 16; + + /** Key used to group vectors (dimension + encoding). */ + record GroupKey(int dimension, VectorEncoding encoding) { + GroupKey(FieldInfo fieldInfo) { + this(fieldInfo.getVectorDimension(), fieldInfo.getVectorEncoding()); + } + } + + record GroupInfo( + int groupOrd, + int dimension, + VectorEncoding encoding, + int groupNumVectors, + long vectorDataOffset, + long vectorDataSize) { + + void write(IndexOutput meta) throws IOException { + meta.writeInt(groupOrd); + meta.writeInt(dimension); + meta.writeInt(encoding.ordinal()); + meta.writeInt(groupNumVectors); + meta.writeLong(vectorDataOffset); + meta.writeLong(vectorDataSize); + } + + static GroupInfo readFromMeta(IndexInput meta) throws IOException { + int groupOrd = meta.readInt(); + if (groupOrd == END_MARKER) { + return null; + } + + int dimension = meta.readInt(); + VectorEncoding encoding = VectorEncoding.values()[meta.readInt()]; + int groupNumVectors = meta.readInt(); + long vectorDataOffset = meta.readLong(); + long vectorDataSize = meta.readLong(); + + return new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + } + } + + static void writeFieldInfo( + IndexOutput meta, + IndexOutput vectorData, + int fieldNumber, + VectorSimilarityFunction function, + int dimension, + VectorEncoding encoding, + int groupOrd, + int vectorCount, + int maxDoc, + DocsWithFieldSet docs, + FieldOrdToGroupOrd fieldOrdToGroupOrd) + throws IOException { + + meta.writeInt(fieldNumber); + meta.writeInt(function.ordinal()); + meta.writeInt(dimension); + meta.writeInt(encoding.ordinal()); + meta.writeInt(groupOrd); + meta.writeInt(vectorCount); + + // write ordToDoc + OrdToDocDISIReaderConfiguration.writeStoredMeta( + ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, vectorCount, maxDoc, docs); + + // write fieldOrdToGroupOrd + long fieldOrdToGroupOrdOffset = vectorData.alignFilePointer(FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES); + DirectWriter writer = + DirectWriter.getInstance(vectorData, vectorCount, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); + for (int i = 0; i < vectorCount; i++) { + writer.add(fieldOrdToGroupOrd.get(i)); + } + writer.finish(); + long fieldOrdToGroupOrdSize = vectorData.getFilePointer() - fieldOrdToGroupOrdOffset; + + meta.writeLong(fieldOrdToGroupOrdOffset); + meta.writeLong(fieldOrdToGroupOrdSize); + } + + static void writeEndMarker(IndexOutput meta) throws IOException { + meta.writeInt(END_MARKER); + } + + record ReadFieldInfo( + int fieldNumber, + VectorSimilarityFunction function, + int dimension, + VectorEncoding encoding, + int groupOrd, + int vectorCount, + OrdToDocDISIReaderConfiguration ordToDoc, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) { + + static ReadFieldInfo read(IndexInput meta) throws IOException { + + int fieldNumber = meta.readInt(); + if (fieldNumber == END_MARKER) { + return null; + } + + VectorSimilarityFunction function = VectorSimilarityFunction.values()[meta.readInt()]; + int dimension = meta.readInt(); + VectorEncoding encoding = VectorEncoding.values()[meta.readInt()]; + int groupOrd = meta.readInt(); + int vectorCount = meta.readInt(); + OrdToDocDISIReaderConfiguration ordToDoc = + OrdToDocDISIReaderConfiguration.fromStoredMeta(meta, vectorCount); + long fieldOrdToGroupOrdOffset = meta.readLong(); + long fieldOrdToGroupOrdSize = meta.readLong(); + + return new ReadFieldInfo( + fieldNumber, + function, + dimension, + encoding, + groupOrd, + vectorCount, + ordToDoc, + fieldOrdToGroupOrdOffset, + fieldOrdToGroupOrdSize); + } + } + + static long hashBytes(byte[] bytes) { + return murmurhash3_x64_128(bytes, 0, bytes.length, GOOD_FAST_HASH_SEED)[0]; + } + + static long alignBytes(IndexOutput output, VectorEncoding encoding) throws IOException { + int alignBytes = + switch (encoding) { + case BYTE -> 4; + case FLOAT32, FLOAT16 -> 64; + }; + return output.alignFilePointer(alignBytes); + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java new file mode 100644 index 000000000000..802f15d72404 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java @@ -0,0 +1,519 @@ +/* + * 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.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT16; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerDense; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerSparse; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene95.OffHeapByteVectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloat16VectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntArrayList; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.RandomAccessInput; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.LongValues; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.packed.DirectReader; + +/** + * Vector values that share a single copy of each distinct vector across the documents and fields + * that reference it. + * + *

Every instance is backed by two views: the {@code fieldView} maps ordinals to docs and drives + * iteration (one entry per document), while the {@code groupView} holds the de-duplicated vectors + * (one entry per distinct vector). {@code fieldOrdToGroupOrd} translates a document ordinal in the + * field into its group ordinal. + */ +sealed interface DedupVectorValues { + /** The dense view over distinct vectors, indexed by group ordinal. */ + KnnVectorValues getGroupView(); + + /** Maps a per-document ordinal to its group ordinal in {@link #getGroupView()}. */ + FieldOrdToGroupOrd getFieldOrdToGroupOrd(); + + /** + * Maps a field's per-document ordinal to the ordinal of its (shared) vector within the group. + * Backed on-heap while writing and off-heap while reading. + */ + sealed interface FieldOrdToGroupOrd { + int get(int ord); + + FieldOrdToGroupOrd copy() throws IOException; + } + + static ByteVectorValues loadDedupBytes( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapByteVectorValues fieldView = + OffHeapByteVectorValues.load( + function, vectorsScorer, configuration, BYTE, dimension, 0, 0, vectorData); + + final OffHeapByteVectorValues groupView = + new OffHeapByteVectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new ByteImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over byte vectors. */ + final class ByteImpl extends ByteVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final ByteVectorValues fieldView; + private final ByteVectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + ByteImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + ByteVectorValues fieldView, + ByteVectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public ByteVectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public ByteImpl copy() throws IOException { + return new ByteImpl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(byte[] target) throws IOException { + if (size() == 0) { + return null; + } + ByteImpl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = copy.fieldView instanceof OffHeapByteVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + static FloatVectorValues loadDedupFloats( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapFloatVectorValues fieldView = + OffHeapFloatVectorValues.load( + function, vectorsScorer, configuration, FLOAT32, dimension, 0, 0, vectorData); + + final OffHeapFloatVectorValues groupView = + new OffHeapFloatVectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new FloatImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over float vectors. */ + final class FloatImpl extends FloatVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final FloatVectorValues fieldView; + private final FloatVectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + FloatImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + FloatVectorValues fieldView, + FloatVectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public FloatVectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public float[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public FloatImpl copy() throws IOException { + return new FloatImpl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(float[] target) throws IOException { + if (size() == 0) { + return null; + } + FloatImpl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = copy.fieldView instanceof OffHeapFloatVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + static Float16VectorValues loadDedupFloat16s( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapFloat16VectorValues fieldView = + OffHeapFloat16VectorValues.load( + function, vectorsScorer, configuration, FLOAT16, dimension, 0, 0, vectorData); + + final OffHeapFloat16VectorValues groupView = + new OffHeapFloat16VectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new Float16Impl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over float16 vectors. */ + final class Float16Impl extends Float16VectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final Float16VectorValues fieldView; + private final Float16VectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + Float16Impl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + Float16VectorValues fieldView, + Float16VectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public Float16VectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public short[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public Float16Impl copy() throws IOException { + return new Float16Impl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + if (size() == 0) { + return null; + } + Float16Impl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = + copy.fieldView instanceof OffHeapFloat16VectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + record DedupVectorScorer( + KnnVectorValues.DocIndexIterator indexIterator, + RandomVectorScorer vectorScorer, + boolean isDense) + implements VectorScorer { + + @Override + public float score() throws IOException { + return vectorScorer.score(indexIterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return indexIterator; + } + + @Override + public Bulk bulk(DocIdSetIterator matchingDocs) { + if (isDense) { + return fromRandomScorerDense(vectorScorer, indexIterator, matchingDocs); + } else { + return fromRandomScorerSparse(vectorScorer, indexIterator, matchingDocs); + } + } + } + + /** On-heap map used during a flush, backed directly by the buffered ordinals. */ + record FieldOrdToGroupOrdArrayList(IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + + @Override + public int get(int ord) { + return fieldOrdToGroupOrd.get(ord); + } + + @Override + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); + } + } + + /** On-heap map used during a sorted flush, indirecting through a new-to-old ordinal map. */ + record FieldOrdToGroupOrdMappedArrayList(int[] map, IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + + @Override + public int get(int ord) { + return fieldOrdToGroupOrd.get(map[ord]); + } + + @Override + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdMappedArrayList(map, fieldOrdToGroupOrd); + } + } + + /** Off-heap map used while reading, backed by a {@link DirectReader}. */ + record FieldOrdToGroupOrdOffHeap( + IndexInput vectorData, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize, + LongValues values) + implements FieldOrdToGroupOrd { + + FieldOrdToGroupOrdOffHeap( + IndexInput vectorData, long fieldOrdToGroupOrdOffset, long fieldOrdToGroupOrdSize) + throws IOException { + RandomAccessInput slice = + vectorData.randomAccessSlice(fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + LongValues values = DirectReader.getInstance(slice, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); + this(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize, values); + } + + @Override + public int get(int v) { + return (int) values.get(v); + } + + @Override + public FieldOrdToGroupOrd copy() throws IOException { + return new FieldOrdToGroupOrdOffHeap( + vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + } + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java new file mode 100644 index 000000000000..31892094183d --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java @@ -0,0 +1,26 @@ +/* + * 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. + */ + +/** + * De-duplicating HNSW vector format. + * + *

Stores each distinct vector once and shares it across the documents and fields that reference + * it, while reusing the Lucene 9.9 HNSW graph. See {@link + * org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat} for the entry point and {@link + * org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat} for the on-disk layout. + */ +package org.apache.lucene.sandbox.codecs.dedup; diff --git a/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 29a44d2ecfa8..4e3e3fc1fca0 100644 --- a/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -14,3 +14,4 @@ # limitations under the License. org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat +org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java new file mode 100644 index 000000000000..bf8b1d4e9fc2 --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java @@ -0,0 +1,440 @@ +/* + * 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.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.hamcrest.Matchers.arrayContainingInAnyOrder; +import static org.hamcrest.Matchers.instanceOf; + +import java.util.ArrayList; +import java.util.Collections; +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.KnnFloat16VectorField; +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.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; +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.NumericDocValues; +import org.apache.lucene.internal.hppc.LongArrayList; +import org.apache.lucene.search.Query; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; + +/** + * Tests that {@link DedupHnswVectorsFormat} stores each distinct vector once. De-duplication is + * observed through the group view size: the number of distinct vectors physically stored, + * regardless of how many documents reference them. + */ +public class TestDedupFlatVectorsFormat extends LuceneTestCase { + + private static IndexWriterConfig config() { + return newIndexWriterConfig() + .setCodec(TestUtil.alwaysKnnVectorsFormat(new DedupHnswVectorsFormat())); + } + + /** Repeated float vectors within a field are stored once but still read back per document. */ + public void testFloatDuplicatesWithinField() throws Exception { + float[] a = {1, 2, 3, 4}; + float[] b = {5, 6, 7, 8}; + float[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int ord = 0; ord < docVectors.length; ord++) { + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloatVectorField("f", docVectors[ord], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(docVectors.length, values.size()); // one entry per document + assertEquals(2, groupNumVectors(values)); // only two distinct vectors stored + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord), 0f); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; + } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); + } + } + } + + /** Repeated float16 vectors within a field are stored once but still read back per document. */ + public void testFloat16DuplicatesWithinField() throws Exception { + short[] a = {Float.floatToFloat16(1f), Float.floatToFloat16(2f), Float.floatToFloat16(3f)}; + short[] b = {Float.floatToFloat16(4f), Float.floatToFloat16(5f), Float.floatToFloat16(6f)}; + short[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int ord = 0; ord < docVectors.length; ord++) { + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloat16VectorField("f", docVectors[ord], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + Float16VectorValues values = leafReader.getFloat16VectorValues("f"); + assertEquals(docVectors.length, values.size()); // one entry per document + assertEquals(2, groupNumVectors(values)); // only two distinct vectors stored + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord)); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; + } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); + } + } + } + + /** Repeated byte vectors within a field are stored once but still read back per document. */ + public void testByteDuplicatesWithinField() throws Exception { + byte[] a = {1, 2, 3, 4}; + byte[] b = {5, 6, 7, 8}; + byte[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int ord = 0; ord < docVectors.length; ord++) { + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnByteVectorField("f", docVectors[ord], EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + ByteVectorValues values = leafReader.getByteVectorValues("f"); + assertEquals(docVectors.length, values.size()); + assertEquals(2, groupNumVectors(values)); + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord)); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; + } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); + } + } + } + + /** Distinct vectors are all kept, i.e. nothing is collapsed by mistake. */ + public void testDistinctVectorsAllStored() throws Exception { + // Vectors that are close to each other in bit representations. + float[][] distinctDocVectors = { + {+0f}, {-0f}, {Math.nextUp(0f)}, {Math.nextDown(0f)}, {1f}, {Math.nextUp(1f)} + }; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (float[] vector : distinctDocVectors) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); + assertEquals(distinctDocVectors.length, values.size()); + assertEquals(distinctDocVectors.length, groupNumVectors(values)); + } + } + } + + /** Check off-heap size of de-duplicated vectors. */ + public void testOffHeapSize() throws Exception { + float[] a = {1, 2, 3, 4}; + float[] b = {5, 6, 7, 8}; + float[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (float[] vector : docVectors) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + DedupFlatVectorsReader dedupReader = getDedupReader(leafReader, "f"); + FieldInfo fieldInfo = leafReader.getFieldInfos().fieldInfo("f"); + + long expectedOffHeapSize = + (docVectors.length * Integer.BYTES) // fieldOrdToGroupOrd mapping + + (a.length + b.length) * Float.BYTES; // raw vector size + + assertEquals( + expectedOffHeapSize, + dedupReader + .getOffHeapByteSize(fieldInfo) + .get("vdd") // vector data extension + .longValue()); + } + } + } + + /** Fields with the same dimension and encoding share one copy of an identical vector. */ + public void testDuplicatesAcrossFieldsShareGroup() throws Exception { + float[] shared = {9, 8, 7, 6}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f1", shared, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f2", shared, DOT_PRODUCT)); // different function + w.addDocument(doc); + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + + DedupFlatVectorsReader dedupReader1 = getDedupReader(leaf, "f1"); + DedupFlatVectorsReader dedupReader2 = getDedupReader(leaf, "f2"); + assertEquals(dedupReader1, dedupReader2); // de-duplication happened correctly + + assertEquals( // both fields DO resolve to the same group + dedupReader1.getEntry("f1", FLOAT32).groupInfo(), + dedupReader2.getEntry("f2", FLOAT32).groupInfo()); + + FloatVectorValues v1 = leaf.getFloatVectorValues("f1"); + assertEquals(1, groupNumVectors(v1)); // the group has one vector + assertArrayEquals(shared, v1.vectorValue(0), 0f); + + FloatVectorValues v2 = leaf.getFloatVectorValues("f2"); + assertEquals(1, groupNumVectors(v2)); // the group has one vector + assertArrayEquals(shared, v2.vectorValue(0), 0f); + } + } + } + + /** Fields differing in dimension use separate groups, even for otherwise similar vectors. */ + public void testDifferentDimensionsUseSeparateGroups() throws Exception { + float[] vector1 = {1, 1}; + float[] vector2 = {1, 1, 0}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f2d", vector1, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f3d", vector2, EUCLIDEAN)); + w.addDocument(doc); + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + + DedupFlatVectorsReader dedupReader1 = getDedupReader(leaf, "f2d"); + DedupFlatVectorsReader dedupReader2 = getDedupReader(leaf, "f3d"); + assertEquals(dedupReader1, dedupReader2); // de-duplication happened correctly + + assertNotEquals( // both fields DO NOT resolve to the same group + dedupReader1.getEntry("f2d", FLOAT32).groupInfo(), + dedupReader2.getEntry("f3d", FLOAT32).groupInfo()); + + FloatVectorValues v1 = leaf.getFloatVectorValues("f2d"); + assertEquals(1, groupNumVectors(v1)); // the group has one vector + assertArrayEquals(vector1, v1.vectorValue(0), 0f); + + FloatVectorValues v2 = leaf.getFloatVectorValues("f3d"); + assertEquals(1, groupNumVectors(v2)); // the group has one vector + assertArrayEquals(vector2, v2.vectorValue(0), 0f); + } + } + } + + /** Duplicates spanning multiple segments collapse to a single copy when merged. */ + public void testDuplicatesAcrossSegmentsDedupOnMerge() throws Exception { + float[] a = {1, 1, 1, 1}; + float[] b = {2, 2, 2, 2}; + float[][] docVectors = {a, b, a}; // 3 docs across 3 segments, 2 distinct + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int ord = 0; ord < docVectors.length; ord++) { + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloatVectorField("f", docVectors[ord], EUCLIDEAN)); + w.addDocument(doc); + w.commit(); // one segment per document + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(docVectors.length, values.size()); + assertEquals(2, groupNumVectors(values)); // a's duplicate collapsed across segments + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord), 0f); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; + } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); + } + } + } + + /** Test that vectors not referenced are deleted from the group. */ + public void testDeletes() throws Exception { + float[] a = {1, 1, 1, 1}; + float[] b = {2, 2, 2, 2}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + + LongArrayList docsWithArrayB = new LongArrayList(); + boolean aIndexed = false, bIndexed = false; + for (int i = 0; i < 50; i++) { // many documents + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", i)); + if (random().nextBoolean()) { // index either a or b + doc.add(new KnnFloatVectorField("f", a)); + aIndexed = true; + } else { + doc.add(new KnnFloatVectorField("f", b)); + docsWithArrayB.add(i); + bIndexed = true; + } + w.addDocument(doc); + } + + assumeTrue("Both vectors a and b indexed", aIndexed && bIndexed); + + w.forceMerge(1); // de-duplicate everything + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(2, groupNumVectors(values)); // the group has both vectors + } + + Query matchDocsWithArrayB = + NumericDocValuesField.newSlowSetQuery("id", docsWithArrayB.toArray()); + w.deleteDocuments(matchDocsWithArrayB); // delete all docs with vector b + + w.forceMerge(1); // de-duplicate everything + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(1, groupNumVectors(values)); // the group now has one vector + } + } + } + + /** Test many duplicates spread across fields, documents, segments. */ + public void testManyDuplicate() throws Exception { + float[] shared = {1, 2, 3, 4}; + List fields = new ArrayList<>(List.of("a", "b", "c", "d", "e")); + boolean atLeastOne = false; + + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int i = 0; i < 50; i++) { // many documents + Document doc = new Document(); + + // randomly pick [0, N) fields to index the same vector + int numFields = random().nextInt(fields.size()); + Collections.shuffle(fields, random()); + for (int j = 0; j < numFields; j++) { + doc.add(new KnnFloatVectorField(fields.get(j), shared)); + atLeastOne = true; + } + + w.addDocument(doc); + + if (random().nextFloat() < 0.2f) { // randomly create segments + w.commit(); + } + } + + w.forceMerge(1); // de-duplicate everything + + assumeTrue("At least one vector indexed", atLeastOne); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + + DedupFlatVectorsReader dedupReader = getDedupReader(leaf, "a"); + for (String field : fields) { // all fields DO resolve to the same group + DedupFlatVectorsReader other = getDedupReader(leaf, field); + assertEquals(dedupReader, other); // de-duplication happened correctly + assertEquals( // both fields DO resolve to the same group + dedupReader.getEntry("a", FLOAT32).groupInfo(), + other.getEntry(field, FLOAT32).groupInfo()); + } + + FloatVectorValues values = leaf.getFloatVectorValues("a"); + assertEquals(1, groupNumVectors(values)); // the group has one vector + assertArrayEquals(shared, values.vectorValue(0), 0f); + } + } + } + + /** Number of distinct vectors physically stored for a field's group. */ + private static int groupNumVectors(KnnVectorValues values) { + assertThat(values, instanceOf(DedupVectorValues.class)); + return ((DedupVectorValues) values).getGroupView().size(); + } + + /** Get underlying dedup vector reader instance. */ + private static DedupFlatVectorsReader getDedupReader(LeafReader leafReader, String fieldName) { + assertThat(leafReader, instanceOf(CodecReader.class)); + KnnVectorsReader knnVectorsReader = ((CodecReader) leafReader).getVectorReader(); + knnVectorsReader = knnVectorsReader.unwrapReaderForField(fieldName); + + assertThat(knnVectorsReader, instanceOf(Lucene99HnswVectorsReader.class)); + FlatVectorsReader flatReader = + ((Lucene99HnswVectorsReader) knnVectorsReader).getFlatVectorsReader(); + + assertThat(flatReader, instanceOf(DedupFlatVectorsReader.class)); + return (DedupFlatVectorsReader) flatReader; + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java new file mode 100644 index 000000000000..6f4c67989ef7 --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java @@ -0,0 +1,255 @@ +/* + * 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.hamcrest.Matchers.greaterThan; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DocValuesSkipIndexType; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.StringHelper; +import org.apache.lucene.util.Version; + +/** + * Runs the standard KNN vectors format suite against the de-duplicating HNSW format. De-duplication + * behavior itself is covered by {@link TestDedupFlatVectorsFormat}. + */ +public class TestDedupHnswVectorsFormat extends BaseKnnVectorsFormatTestCase { + + private final KnnVectorsFormat format = new DedupHnswVectorsFormat(); + + @Override + protected Codec getCodec() { + return TestUtil.alwaysKnnVectorsFormat(format); + } + + @Override + protected boolean supportsFloatVectorFallback() { + return false; // stores raw vectors, no quantized fallback + } + + @Override + protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOException { + var fieldInfo = r.getFieldInfos().fieldInfo(fieldName); + + if (r instanceof CodecReader codecReader) { + KnnVectorsReader knnVectorsReader = codecReader.getVectorReader(); + knnVectorsReader = knnVectorsReader.unwrapReaderForField(fieldName); + var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); + long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum(); + if (knnVectorsReader instanceof Lucene99HnswVectorsReader) { + if (getNumVectors(knnVectorsReader, fieldInfo) == 0) { + assertEquals(0L, totalByteSize); + } else { + assertTrue(totalByteSize > 0); + assertTrue(offHeap.get("vdd") > 0L); // NOTE: different from vec + + if (hasHNSW(knnVectorsReader, fieldInfo)) { + assertTrue(offHeap.get("vex") > 0L); + } else { + assertTrue(offHeap.get("vex") == null || offHeap.get("vex") == 0); + } + } + } + } else { + throw new AssertionError("unexpected:" + r.getClass()); + } + } + + /** Near copy of the original test, this one checks for size of unique vector count. */ + @Override + @SuppressWarnings("unchecked") + public void testWriterRamEstimate() throws IOException { + final FieldInfos fieldInfos = new FieldInfos(new FieldInfo[0]); + final Directory dir = newDirectory(); + Codec codec = Codec.getDefault(); + final SegmentInfo si = + new SegmentInfo( + dir, + Version.LATEST, + Version.LATEST, + "0", + 10000, + false, + false, + codec, + Collections.emptyMap(), + StringHelper.randomId(), + new HashMap<>(), + null); + final SegmentWriteState state = + new SegmentWriteState( + InfoStream.getDefault(), dir, si, fieldInfos, null, newIOContext(random())); + final KnnVectorsFormat format = codec.knnVectorsFormat(); + try (KnnVectorsWriter writer = format.fieldsWriter(state)) { + final long ramBytesUsed = writer.ramBytesUsed(); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + ++dim; + } + int numDocs = atLeast(100); + Set unique = new HashSet<>(); + KnnFieldVectorsWriter fieldWriter = + (KnnFieldVectorsWriter) + writer.addField( + new FieldInfo( + "fieldA", + 0, + false, + false, + false, + IndexOptions.NONE, + DocValuesType.NONE, + DocValuesSkipIndexType.NONE, + -1, + Map.of(), + 0, + 0, + 0, + dim, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.DOT_PRODUCT, + false, + false)); + for (int i = 0; i < numDocs; i++) { + float[] vector = randomVector(dim); + unique.add(new FloatVector(vector)); + fieldWriter.addValue(i, vector); + } + final long ramBytesUsed2 = writer.ramBytesUsed(); + assertThat(ramBytesUsed2, greaterThan(ramBytesUsed)); + assertThat(ramBytesUsed2, greaterThan((long) dim * unique.size() * Float.BYTES)); + } + dir.close(); + } + + private record FloatVector(float[] vector) { + @Override + public boolean equals(Object obj) { + return obj instanceof FloatVector(float[] other) && Arrays.equals(vector, other); + } + + @Override + public int hashCode() { + return Arrays.hashCode(vector); + } + } + + /** Near copy of the original test, this one checks for size of unique vector count. */ + @Override + @SuppressWarnings("unchecked") + public void testWriterByteVectorRamEstimate() throws IOException { + final FieldInfos fieldInfos = new FieldInfos(new FieldInfo[0]); + final Directory dir = newDirectory(); + Codec codec = Codec.getDefault(); + final SegmentInfo si = + new SegmentInfo( + dir, + Version.LATEST, + Version.LATEST, + "0", + 10000, + false, + false, + codec, + Collections.emptyMap(), + StringHelper.randomId(), + new HashMap<>(), + null); + final SegmentWriteState state = + new SegmentWriteState( + InfoStream.getDefault(), dir, si, fieldInfos, null, newIOContext(random())); + final KnnVectorsFormat format = codec.knnVectorsFormat(); + try (KnnVectorsWriter writer = format.fieldsWriter(state)) { + final long ramBytesUsed = writer.ramBytesUsed(); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + ++dim; + } + int numDocs = atLeast(100); + Set unique = new HashSet<>(); + KnnFieldVectorsWriter fieldWriter = + (KnnFieldVectorsWriter) + writer.addField( + new FieldInfo( + "fieldA", + 0, + false, + false, + false, + IndexOptions.NONE, + DocValuesType.NONE, + DocValuesSkipIndexType.NONE, + -1, + Map.of(), + 0, + 0, + 0, + dim, + VectorEncoding.BYTE, + VectorSimilarityFunction.DOT_PRODUCT, + false, + false)); + for (int i = 0; i < numDocs; i++) { + byte[] vector = randomVector8(dim); + unique.add(new ByteVector(vector)); + fieldWriter.addValue(i, vector); + } + final long ramBytesUsed2 = writer.ramBytesUsed(); + assertThat(ramBytesUsed2, greaterThan(ramBytesUsed)); + assertThat(ramBytesUsed2, greaterThan((long) dim * unique.size() * Byte.BYTES)); + } + dir.close(); + } + + private record ByteVector(byte[] vector) { + @Override + public boolean equals(Object obj) { + return obj instanceof ByteVector(byte[] other) && Arrays.equals(vector, other); + } + + @Override + public int hashCode() { + return Arrays.hashCode(vector); + } + } +} diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java index bc3b93a77db8..18a96bea5678 100644 --- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java +++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/BaseKnnVectorsFormatTestCase.java @@ -2410,7 +2410,8 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx } } - static int getNumVectors(KnnVectorsReader reader, FieldInfo fieldInfo) throws IOException { + protected static int getNumVectors(KnnVectorsReader reader, FieldInfo fieldInfo) + throws IOException { return switch (fieldInfo.getVectorEncoding()) { case BYTE -> reader.getByteVectorValues(fieldInfo.getName()).size(); case FLOAT32 -> reader.getFloatVectorValues(fieldInfo.getName()).size(); @@ -2436,7 +2437,7 @@ static boolean hasQuantized(KnnVectorsReader knnVectorsReader, FieldInfo fieldIn return name.contains("quantized"); } - static boolean hasHNSW(KnnVectorsReader knnVectorsReader, FieldInfo fieldInfo) + protected static boolean hasHNSW(KnnVectorsReader knnVectorsReader, FieldInfo fieldInfo) throws IOException { if (knnVectorsReader instanceof AssertingKnnVectorsFormat.AssertingKnnVectorsReader assertingReader) {