diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 979f742b3082..20b8c4fdf83f 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -321,6 +321,9 @@ New Features and document length instead of corpus statistics such as document frequency. It also supports k3 query-term frequency saturation. (Tianxiao Wei) +* GITHUB#15979: Add a de-duplicating HNSW vector format (Lucene106DedupHnswVectorsFormat) that stores + each distinct 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/lucene106/dedup/DedupFlatFieldVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java new file mode 100644 index 000000000000..b31545251137 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/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.codecs.lucene106.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 ordToVecOrd; + private int lastDocID; + private boolean finished; + + DedupFlatFieldVectorsWriter(DedupGroup group) { + this.group = group; + this.docsWithFieldSet = new DocsWithFieldSet(); + this.vectors = new ArrayList<>(); + this.ordToVecOrd = new IntArrayList(); + this.lastDocID = -1; + this.finished = false; + } + + @Override + public List getVectors() { + return vectors; + } + + @Override + public DocsWithFieldSet getDocsWithFieldSet() { + return docsWithFieldSet; + } + + IntArrayList getOrdToVecOrd() { + return ordToVecOrd; + } + + @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 + ordToVecOrd.add(cursor.index); // index in group + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + + docsWithFieldSet.ramBytesUsed() + + (long) vectors.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF + + ordToVecOrd.ramBytesUsed(); + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java new file mode 100644 index 000000000000..41d41be9905f --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.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 a per-document {@code ordToVecOrd} 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). + *
  • 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 ordToVecOrd} 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 ordToVecOrd} map in the .vdd file + *
  • [int64] length of this field's {@code ordToVecOrd} 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 ordToVecOrd} 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 ordToVecOrd} 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 ordToVecOrd} map stay off-heap; a read + * resolves a document ordinal to its vector via one extra {@code ordToVecOrd} lookup. + *
+ * + * @lucene.experimental + */ +final class DedupFlatVectorsFormat extends FlatVectorsFormat { + static final String NAME = "Lucene106DedupFlatVectorsFormat"; + + static final String META_CODEC_NAME = "Lucene106DedupFlatVectorsFormatMeta"; + static final String META_EXTENSION = "vdm"; + + static final String VECTOR_DATA_CODEC_NAME = "Lucene106DedupFlatVectorsFormatVectorData"; + 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/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java new file mode 100644 index 000000000000..bf016dad7647 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java @@ -0,0 +1,330 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VERSION_START; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.loadDedupBytes; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.loadDedupFloats; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.readFieldInfo; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.readGroupInfo; +import static org.apache.lucene.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; + +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.codecs.lucene106.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.ReadFieldInfo; +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.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.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 an {@code ordToVecOrd} 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 = readGroupInfo(meta); + if (groupInfo == null) { + break; + } + groupInfos.add(groupInfo); + } + + while (true) { + ReadFieldInfo fieldInfo = readFieldInfo(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( + "Invalid vector function: indexed=" + + fieldInfo.function() + + ", actual=" + + info.getVectorSimilarityFunction(), + meta); + } else if (fieldInfo.dimension() != info.getVectorDimension()) { + throw new CorruptIndexException( + "Invalid vector dimension: indexed=" + + fieldInfo.dimension() + + ", actual=" + + info.getVectorDimension(), + meta); + } else if (fieldInfo.encoding() != info.getVectorEncoding()) { + throw new CorruptIndexException( + "Invalid 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; + } + + private 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 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.groupSize(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.ordToVecOffset(), + entry.fieldInfo.ordToVecSize()); + } + + @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.groupSize(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.ordToVecOffset(), + entry.fieldInfo.ordToVecSize()); + } + + @Override + public ByteVectorValues getByteVectorValues(String field) throws IOException { + return getByteVectorValues(getEntry(field, BYTE)); + } + + @Override + public FlatVectorsReader getMergeInstance() { + // TODO: Can we improve performance using strictly sequential IO? + 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.ordToVecSize() + entry.groupInfo.vectorDataSize()); + } + + private 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/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java new file mode 100644 index 000000000000..4449808e7bc5 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java @@ -0,0 +1,170 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.SCRATCH_SIZE; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrd; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +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()}, mapped to + * group ordinals via {@link DedupVectorValues#getOrdToVecOrd()}. + * + * @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 delegate = + SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); + RandomVectorScorerSupplier groupView = + SCORER.getRandomVectorScorerSupplier(similarityFunction, dedupValues.getGroupView()); + return new RandomVectorScorerSupplierImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + } + return SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, float[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer delegate = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, byte[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer delegate = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + + private record RandomVectorScorerSupplierImpl( + RandomVectorScorerSupplier delegate, + RandomVectorScorerSupplier groupView, + OrdToVecOrd ordToVecOrd) + implements RandomVectorScorerSupplier { + + @Override + public UpdateableRandomVectorScorer scorer() throws IOException { + return new UpdateableRandomVectorScorerImpl( + delegate.scorer(), groupView.scorer(), ordToVecOrd); + } + + @Override + public RandomVectorScorerSupplier copy() throws IOException { + return new RandomVectorScorerSupplierImpl(delegate.copy(), groupView.copy(), ordToVecOrd); + } + } + + private static class RandomVectorScorerImpl implements RandomVectorScorer { + private final RandomVectorScorer delegate; + private final RandomVectorScorer groupView; + private final OrdToVecOrd ordToVecOrd; + private int[] scratch; + + RandomVectorScorerImpl( + RandomVectorScorer delegate, RandomVectorScorer groupView, OrdToVecOrd ordToVecOrd) { + this.delegate = delegate; + this.groupView = groupView; + this.ordToVecOrd = ordToVecOrd; + this.scratch = new int[SCRATCH_SIZE]; + } + + @Override + public int ordToDoc(int ord) { + return delegate.ordToDoc(ord); + } + + @Override + public Bits getAcceptOrds(Bits acceptDocs) { + return delegate.getAcceptOrds(acceptDocs); + } + + @Override + public float score(int node) throws IOException { + return groupView.score(ordToVecOrd.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] = ordToVecOrd.get(nodes[i]); + } + return groupView.bulkScore(scratch, scores, numNodes); + } + + @Override + public int maxOrd() { + return delegate.maxOrd(); + } + } + + private static final class UpdateableRandomVectorScorerImpl extends RandomVectorScorerImpl + implements UpdateableRandomVectorScorer { + private final UpdateableRandomVectorScorer groupView; + private final OrdToVecOrd ordToVecOrd; + + UpdateableRandomVectorScorerImpl( + UpdateableRandomVectorScorer delegate, + UpdateableRandomVectorScorer groupView, + OrdToVecOrd ordToVecOrd) { + super(delegate, groupView, ordToVecOrd); + this.groupView = groupView; + this.ordToVecOrd = ordToVecOrd; + } + + @Override + public void setScoringOrdinal(int node) throws IOException { + groupView.setScoringOrdinal(ordToVecOrd.get(node)); + } + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java new file mode 100644 index 000000000000..81099818dac3 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java @@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.codecs.lucene106.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; + IndexOutput m = null, v = null; + try { + m = state.directory.createOutput(metaFileName, state.context); + v = state.directory.createOutput(vectorDataFileName, state.context); + CodecUtil.writeIndexHeader( + m, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + CodecUtil.writeIndexHeader( + v, + VECTOR_DATA_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + this.meta = m; + this.vectorData = v; + success = true; + } finally { + if (success == false) { + IOUtils.closeWhileHandlingException(m, v); + } + } + } + + @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(); + } + + if (meta != null) { + CodecUtil.writeFooter(meta); + } + + if (vectorData != null) { + 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/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java new file mode 100644 index 000000000000..9e53f0636ffd --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java @@ -0,0 +1,268 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static java.nio.ByteOrder.LITTLE_ENDIAN; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.ORD_UNKNOWN; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeEndOfFields; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeFieldInfo; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeGroupInfo; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.FloatBuffer; +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.codecs.lucene106.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.GroupKey; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrd; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrdArrayList; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrdMappedArrayList; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.WriteFieldInfo; +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.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) { + return switch (groupKey.encoding()) { + case BYTE -> new ByteGroup(groupKey.dimension()); + case FLOAT32 -> new FloatGroup(groupKey.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(); + DedupGroup group = entry.getValue(); + + int groupSize = group.size(); + long vectorDataOffset = alignBytes(vectorData, groupKey.encoding()); + + // TODO: Write in sorted order for faster merge? (with sequential IO) + for (int ord = 0; ord < groupSize; ord++) { + byte[] bytes = group.serialize(ord); + vectorData.writeBytes(bytes, bytes.length); + } + long vectorDataSize = vectorData.getFilePointer() - vectorDataOffset; + + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); + + GroupInfo groupInfo = + new GroupInfo(groupOrd, dimension, encoding, groupSize, vectorDataOffset, vectorDataSize); + writeGroupInfo(meta, groupInfo); + + groupOrds.put(groupKey, groupOrd); + groupOrd++; + } + + writeEndOfGroups(meta); + + for (FieldData fieldData : fieldDataList) { + fieldData.fieldWriter.finish(); + + IntArrayList ordToVecOrd = fieldData.fieldWriter.getOrdToVecOrd(); + int vectorCount = ordToVecOrd.elementsCount; + + DocsWithFieldSet docs; + OrdToVecOrd ordToVecFinal; + if (sortMap == null) { + docs = fieldData.fieldWriter.getDocsWithFieldSet(); + ordToVecFinal = new OrdToVecOrdArrayList(ordToVecOrd); + } else { + DocsWithFieldSet oldDocs = fieldData.fieldWriter.getDocsWithFieldSet(); + docs = new DocsWithFieldSet(); + int[] new2OldOrd = new int[vectorCount]; + KnnVectorsWriter.mapOldOrdToNewOrd(oldDocs, sortMap, null, new2OldOrd, docs); + ordToVecFinal = new OrdToVecOrdMappedArrayList(new2OldOrd, ordToVecOrd); + } + + WriteFieldInfo fieldInfo = + new WriteFieldInfo( + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + vectorCount, + maxDoc, + docs, + ordToVecFinal); + writeFieldInfo(meta, vectorData, fieldInfo); + } + + writeEndOfFields(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() + size() * 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; + private int lastOrd; + + 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(); + this.lastOrd = ORD_UNKNOWN; + } + + @Override + public long hash(float[] vector) { + // the vector needs to be converted to bytes to use a utility hash function. + // the existing buffer is used for this conversion, so lastOrd is reset too. + buffer.put(0, vector); + lastOrd = ORD_UNKNOWN; + 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) { + if (ord != lastOrd) { + buffer.put(0, get(ord)); + lastOrd = ord; + } + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + size() * ramBytesPerVector; + } + } + + private record FieldData( + FieldInfo fieldInfo, GroupKey groupKey, DedupFlatFieldVectorsWriter fieldWriter) + implements Accountable { + + @Override + public long ramBytesUsed() { + return fieldWriter.ramBytesUsed(); + } + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java new file mode 100644 index 000000000000..4d46d0121e0e --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.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, + DedupMergeContext.DedupMergeGroup { + + private static final int ORD_NOT_FOUND = -1; + + 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 size() { + 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 ordHint = hashToOrdHint.getOrDefault(hash, ORD_NOT_FOUND); + if (ordHint == ORD_NOT_FOUND) { + groupOrd = vectors.size(); + ownedVector = copy(vectorValue); // only for unique vectors + hashToOrdHint.put(hash, groupOrd); + vectors.add(ownedVector); + break; + } else { + T other = vectors.get(ordHint); + if (equals(vectorValue, other)) { + groupOrd = ordHint; + 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/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java new file mode 100644 index 000000000000..db10db4ca5f3 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java @@ -0,0 +1,357 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static java.nio.ByteOrder.LITTLE_ENDIAN; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.ORD_UNKNOWN; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeEndOfFields; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeFieldInfo; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.writeGroupInfo; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.FloatBuffer; +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.codecs.lucene106.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.GroupKey; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrd; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrdArrayList; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.WriteFieldInfo; +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.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.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(); + long vectorDataOffset = alignBytes(vectorData, groupKey.encoding()); + + DedupMergeGroup mergeGroup = + switch (groupKey.encoding()) { + case BYTE -> new ByteGroup(); + case FLOAT32 -> new FloatGroup(groupKey.dimension()); + }; + + for (FieldData fieldData : entry.getValue()) { + mergeGroup.processField(fieldData, vectorData); + } + + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); + int groupSize = mergeGroup.size(); + long vectorDataSize = vectorData.getFilePointer() - vectorDataOffset; + + GroupInfo groupInfo = + new GroupInfo(groupOrd, dimension, encoding, groupSize, vectorDataOffset, vectorDataSize); + writeGroupInfo(meta, groupInfo); + + groupOrds.put(groupKey, groupOrd); + groupOrd++; + } + + writeEndOfGroups(meta); + + for (FieldData fieldData : fieldDataList) { + WriteFieldInfo fieldInfo = + new WriteFieldInfo( + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + fieldData.ordToVecOrd.elementsCount, + fieldData.maxDoc, + fieldData.docsWithFieldSet, + new OrdToVecOrdArrayList(fieldData.ordToVecOrd)); + writeFieldInfo(meta, vectorData, fieldInfo); + } + + writeEndOfFields(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 groupSize = size(); + + // add vector to group + ObjectCursor cursor = super.addUnique(vector); + if (cursor.index == groupSize) { // new addition + // already on-heap, write immediately to avoid another IO read + byte[] bytes = serialize(groupSize); + vectorData.writeBytes(bytes, bytes.length); + } + + // record hit and ord in group + fieldData.docsWithFieldSet.add(next.mappedDocID); + fieldData.ordToVecOrd.add(cursor.index); + } + } + } + + 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) { + OrdToVecOrd ordToVecOrd = dedup.getOrdToVecOrd(); + return ordToVecOrd.get(vector.ord) == ordToVecOrd.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() + size() * ByteVector.SHALLOW_SIZE; + } + } + + 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; + private int lastOrd; + + FloatGroup(int dimension) { + int length = dimension * Float.BYTES; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asFloatBuffer(); + this.lastOrd = ORD_UNKNOWN; + } + + @Override + FloatVector vectorFrom(Sub sub) { + return new FloatVector(sub.values, sub.iterator.index()); + } + + @Override + public long hash(FloatVector vector) throws IOException { + // the vector needs to be converted to bytes to use a utility hash function. + // the existing buffer is used for this conversion, so lastOrd is reset too. + buffer.put(0, vector.get()); + lastOrd = ORD_UNKNOWN; + 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) { + OrdToVecOrd ordToVecOrd = dedup.getOrdToVecOrd(); + return ordToVecOrd.get(vector.ord) == ordToVecOrd.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 { + if (ord != lastOrd) { + buffer.put(0, get(ord).get()); + lastOrd = ord; + } + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + size() * FloatVector.SHALLOW_SIZE; + } + } + + private record FieldData( + FieldInfo fieldInfo, + GroupKey groupKey, + DocsWithFieldSet docsWithFieldSet, + IntArrayList ordToVecOrd, + 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); + }; + + if (vectorValues == null) { + continue; + } + + subs.add(new Sub<>(mergeState.docMaps[i], vectorValues)); + } + + return DocIDMerger.of(subs, mergeState.needsIndexSort); + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java new file mode 100644 index 000000000000..a76c25747017 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java @@ -0,0 +1,562 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static org.apache.lucene.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +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.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene95.OffHeapByteVectorValues; +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.DocsWithFieldSet; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorEncoding; +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.IndexOutput; +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; +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 DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + + private static final int END_MARKER = -1; + + private static final int ORD_TO_VEC_ALIGN_BYTES = 4; + + // TODO: Evaluate using fewer bits. + private static final int ORD_TO_VEC_BITS_PER_VALUE = 32; + + static final int ORD_UNKNOWN = -1; + + static final int SCRATCH_SIZE = 16; + + /** Key used to group vectors (dimension + encoding). */ + record GroupKey(int dimension, VectorEncoding encoding) { + GroupKey(FieldInfo fieldInfo) { + this(fieldInfo.getVectorDimension(), fieldInfo.getVectorEncoding()); + } + } + + /** + * 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 delegate} 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 ordToVecOrd} translates a document ordinal 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()}. */ + OrdToVecOrd getOrdToVecOrd(); + } + + /** + * 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 OrdToVecOrd { + int get(int ord); + + OrdToVecOrd copy() throws IOException; + } + + record GroupInfo( + int groupOrd, + int dimension, + VectorEncoding encoding, + int groupSize, + long vectorDataOffset, + long vectorDataSize) {} + + static void writeGroupInfo(IndexOutput meta, GroupInfo groupInfo) throws IOException { + meta.writeInt(groupInfo.groupOrd); + meta.writeInt(groupInfo.dimension); + meta.writeInt(groupInfo.encoding.ordinal()); + meta.writeInt(groupInfo.groupSize); + meta.writeLong(groupInfo.vectorDataOffset); + meta.writeLong(groupInfo.vectorDataSize); + } + + static void writeEndOfGroups(IndexOutput meta) throws IOException { + meta.writeInt(END_MARKER); + } + + static GroupInfo readGroupInfo(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 groupSize = meta.readInt(); + long vectorDataOffset = meta.readLong(); + long vectorDataSize = meta.readLong(); + + return new GroupInfo( + groupOrd, dimension, encoding, groupSize, vectorDataOffset, vectorDataSize); + } + + record WriteFieldInfo( + int fieldNumber, + VectorSimilarityFunction function, + int dimension, + VectorEncoding encoding, + int groupOrd, + int vectorCount, + int maxDoc, + DocsWithFieldSet docs, + OrdToVecOrd ordToVecOrd) {} + + static void writeFieldInfo(IndexOutput meta, IndexOutput vectorData, WriteFieldInfo fieldInfo) + throws IOException { + + meta.writeInt(fieldInfo.fieldNumber); + meta.writeInt(fieldInfo.function.ordinal()); + meta.writeInt(fieldInfo.dimension); + meta.writeInt(fieldInfo.encoding.ordinal()); + meta.writeInt(fieldInfo.groupOrd); + meta.writeInt(fieldInfo.vectorCount); + + // write ordToDoc + OrdToDocDISIReaderConfiguration.writeStoredMeta( + DIRECT_MONOTONIC_BLOCK_SHIFT, + meta, + vectorData, + fieldInfo.vectorCount, + fieldInfo.maxDoc, + fieldInfo.docs); + + // write ordToVec + long ordToVecOffset = vectorData.alignFilePointer(ORD_TO_VEC_ALIGN_BYTES); + DirectWriter writer = + DirectWriter.getInstance(vectorData, fieldInfo.vectorCount, ORD_TO_VEC_BITS_PER_VALUE); + for (int i = 0; i < fieldInfo.vectorCount; i++) { + writer.add(fieldInfo.ordToVecOrd.get(i)); + } + writer.finish(); + long ordToVecSize = vectorData.getFilePointer() - ordToVecOffset; + + meta.writeLong(ordToVecOffset); + meta.writeLong(ordToVecSize); + } + + static void writeEndOfFields(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 ordToVecOffset, + long ordToVecSize) {} + + static ReadFieldInfo readFieldInfo(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 ordToVecOffset = meta.readLong(); + long ordToVecSize = meta.readLong(); + + return new ReadFieldInfo( + fieldNumber, + function, + dimension, + encoding, + groupOrd, + vectorCount, + ordToDoc, + ordToVecOffset, + ordToVecSize); + } + + 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 -> 64; + }; + return output.alignFilePointer(alignBytes); + } + + /** On-heap map used during a flush, backed directly by the buffered ordinals. */ + record OrdToVecOrdArrayList(IntArrayList ordToVecOrd) implements OrdToVecOrd { + @Override + public int get(int ord) { + return ordToVecOrd.get(ord); + } + + @Override + public OrdToVecOrd copy() { + return new OrdToVecOrdArrayList(ordToVecOrd); + } + } + + /** On-heap map used during a sorted flush, indirecting through a new-to-old ordinal map. */ + record OrdToVecOrdMappedArrayList(int[] map, IntArrayList ordToVecOrd) implements OrdToVecOrd { + @Override + public int get(int ord) { + return ordToVecOrd.get(map[ord]); + } + + @Override + public OrdToVecOrd copy() { + return new OrdToVecOrdMappedArrayList(map, ordToVecOrd); + } + } + + /** Off-heap map used while reading, backed by a {@link DirectReader}. */ + static final class OrdToVecOrdOffHeap implements OrdToVecOrd { + private final IndexInput vectorData; + private final long ordToVecOffset; + private final long ordToVecSize; + private final LongValues values; + + OrdToVecOrdOffHeap(IndexInput vectorData, long ordToVecOffset, long ordToVecSize) + throws IOException { + this.vectorData = vectorData; + this.ordToVecOffset = ordToVecOffset; + this.ordToVecSize = ordToVecSize; + + RandomAccessInput slice = vectorData.randomAccessSlice(ordToVecOffset, ordToVecSize); + this.values = DirectReader.getInstance(slice, ORD_TO_VEC_BITS_PER_VALUE); + } + + @Override + public int get(int v) { + return (int) values.get(v); + } + + @Override + public OrdToVecOrd copy() throws IOException { + return new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + } + } + + static ByteVectorValues loadDedupBytes( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupSize, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long ordToVecOffset, + long ordToVecSize) + throws IOException { + + final OffHeapByteVectorValues delegate = + OffHeapByteVectorValues.load( + function, vectorsScorer, configuration, BYTE, dimension, 0, 0, vectorData); + + final OffHeapByteVectorValues groupView = + new OffHeapByteVectorValues.DenseOffHeapVectorValues( + dimension, + groupSize, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + delegate.getVectorByteLength(), + vectorsScorer, + function); + + final OrdToVecOrd ordToVecOrd = + new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + + return new ByteImpl(vectorsScorer, function, delegate, groupView, ordToVecOrd); + } + + /** {@link DedupVectorValues} over byte vectors. */ + private static final class ByteImpl extends ByteVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final ByteVectorValues delegate; + private final ByteVectorValues groupView; + private final OrdToVecOrd ordToVecOrd; + private int[] scratch; + + ByteImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + ByteVectorValues delegate, + ByteVectorValues groupView, + OrdToVecOrd ordToVecOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.delegate = delegate; + this.groupView = groupView; + this.ordToVecOrd = ordToVecOrd; + this.scratch = new int[SCRATCH_SIZE]; + } + + @Override + public ByteVectorValues getGroupView() { + return groupView; + } + + @Override + public OrdToVecOrd getOrdToVecOrd() { + return ordToVecOrd; + } + + @Override + public int ordToDoc(int ord) { + return delegate.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] = ordToVecOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(ordToVecOrd.get(ord)); + } + + @Override + public int dimension() { + return delegate.dimension(); + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public ByteImpl copy() throws IOException { + return new ByteImpl( + vectorsScorer, function, delegate.copy(), groupView.copy(), ordToVecOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return delegate.iterator(); + } + + @Override + public VectorScorer scorer(byte[] target) throws IOException { + if (size() == 0) { + return null; + } + ByteImpl copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return vectorScorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerDense(vectorScorer, iterator, matchingDocs); + } + }; + } + } + + static FloatVectorValues loadDedupFloats( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupSize, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long ordToVecOffset, + long ordToVecSize) + throws IOException { + + final OffHeapFloatVectorValues delegate = + OffHeapFloatVectorValues.load( + function, vectorsScorer, configuration, FLOAT32, dimension, 0, 0, vectorData); + + final OffHeapFloatVectorValues groupView = + new OffHeapFloatVectorValues.DenseOffHeapVectorValues( + dimension, + groupSize, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + delegate.getVectorByteLength(), + vectorsScorer, + function); + + final OrdToVecOrd ordToVecOrd = + new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + + return new FloatImpl(vectorsScorer, function, delegate, groupView, ordToVecOrd); + } + + /** {@link DedupVectorValues} over float vectors. */ + private static final class FloatImpl extends FloatVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final FloatVectorValues delegate; + private final FloatVectorValues groupView; + private final OrdToVecOrd ordToVecOrd; + private int[] scratch; + + FloatImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + FloatVectorValues delegate, + FloatVectorValues groupView, + OrdToVecOrd ordToVecOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.delegate = delegate; + this.groupView = groupView; + this.ordToVecOrd = ordToVecOrd; + this.scratch = new int[SCRATCH_SIZE]; + } + + @Override + public FloatVectorValues getGroupView() { + return groupView; + } + + @Override + public OrdToVecOrd getOrdToVecOrd() { + return ordToVecOrd; + } + + @Override + public int ordToDoc(int ord) { + return delegate.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] = ordToVecOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public float[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(ordToVecOrd.get(ord)); + } + + @Override + public int dimension() { + return delegate.dimension(); + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public FloatImpl copy() throws IOException { + return new FloatImpl( + vectorsScorer, function, delegate.copy(), groupView.copy(), ordToVecOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return delegate.iterator(); + } + + @Override + public VectorScorer scorer(float[] target) throws IOException { + if (size() == 0) { + return null; + } + FloatImpl copy = copy(); + DocIndexIterator iterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + return new VectorScorer() { + @Override + public float score() throws IOException { + return vectorScorer.score(iterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return iterator; + } + + @Override + public Bulk bulk(DocIdSetIterator matchingDocs) { + return Bulk.fromRandomScorerDense(vectorScorer, iterator, matchingDocs); + } + }; + } + } +} diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java new file mode 100644 index 000000000000..ce4c6223bc68 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java @@ -0,0 +1,221 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.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 trades a small amount of indexing work for reduced storage + * when vectors repeat, e.g. multiple fields derived from the same embedding or heavily duplicated + * content. + * + * @lucene.experimental + */ +public final class Lucene106DedupHnswVectorsFormat extends KnnVectorsFormat { + private static final String NAME = "Lucene106DedupHnswVectorsFormat"; + + /** + * 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 Lucene106DedupHnswVectorsFormat() { + 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 Lucene106DedupHnswVectorsFormat(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 Lucene106DedupHnswVectorsFormat(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 Lucene106DedupHnswVectorsFormat( + 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 Lucene106DedupHnswVectorsFormat( + int maxConn, + int beamWidth, + int numMergeWorkers, + ExecutorService mergeExec, + int tinySegmentsThreshold) { + super("Lucene106DedupHnswVectorsFormat"); + 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 { + 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/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java new file mode 100644 index 000000000000..1b732db6a79b --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/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.codecs.lucene106.dedup.Lucene106DedupHnswVectorsFormat} for the entry point and + * {@link org.apache.lucene.codecs.lucene106.dedup.DedupFlatVectorsFormat} for the on-disk layout. + */ +package org.apache.lucene.codecs.lucene106.dedup; diff --git a/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 3ac106d11c84..9d562567d0da 100644 --- a/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/lucene/core/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -16,3 +16,4 @@ org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat org.apache.lucene.codecs.lucene104.Lucene104ScalarQuantizedVectorsFormat org.apache.lucene.codecs.lucene104.Lucene104HnswScalarQuantizedVectorsFormat +org.apache.lucene.codecs.lucene106.dedup.Lucene106DedupHnswVectorsFormat diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java new file mode 100644 index 000000000000..96f3aa295fda --- /dev/null +++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java @@ -0,0 +1,186 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnByteVectorField; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; + +/** + * Tests that {@link Lucene106DedupHnswVectorsFormat} 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 IndexWriterConfig config() { + return newIndexWriterConfig() + .setCodec(TestUtil.alwaysKnnVectorsFormat(new Lucene106DedupHnswVectorsFormat())); + } + + /** 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 (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)) { + FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); + assertEquals(docVectors.length, values.size()); // one entry per document + assertEquals(2, groupSize(values)); // only two distinct vectors stored + for (int ord = 0; ord < values.size(); ord++) { + assertArrayEquals(docVectors[ord], values.vectorValue(ord), 0f); + } + } + } + } + + /** 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, a, b, a, b}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (byte[] vector : docVectors) { + Document doc = new Document(); + doc.add(new KnnByteVectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + ByteVectorValues values = getOnlyLeafReader(reader).getByteVectorValues("f"); + assertEquals(docVectors.length, values.size()); + assertEquals(2, groupSize(values)); + for (int ord = 0; ord < values.size(); ord++) { + assertArrayEquals(docVectors[ord], values.vectorValue(ord)); + } + } + } + } + + /** Distinct vectors are all kept, i.e. nothing is collapsed by mistake. */ + public void testDistinctVectorsAllStored() throws Exception { + float[][] docVectors = {{1, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}}; + 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)) { + FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); + assertEquals(3, values.size()); + assertEquals(3, groupSize(values)); + } + } + } + + /** 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, EUCLIDEAN)); + w.addDocument(doc); + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + FloatVectorValues v1 = leaf.getFloatVectorValues("f1"); + FloatVectorValues v2 = leaf.getFloatVectorValues("f2"); + assertEquals(1, groupSize(v1)); // both fields resolve to the same one-vector group + assertEquals(1, groupSize(v2)); + assertArrayEquals(shared, v1.vectorValue(0), 0f); + assertArrayEquals(shared, v2.vectorValue(0), 0f); + } + } + } + + /** Fields differing in dimension use separate groups, even for otherwise similar vectors. */ + public void testDifferentDimensionsUseSeparateGroups() throws Exception { + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f2d", new float[] {1, 1}, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f3d", new float[] {1, 1, 1}, EUCLIDEAN)); + w.addDocument(doc); + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + assertEquals(1, groupSize(leaf.getFloatVectorValues("f2d"))); + assertEquals(1, groupSize(leaf.getFloatVectorValues("f3d"))); + assertArrayEquals(new float[] {1, 1}, leaf.getFloatVectorValues("f2d").vectorValue(0), 0f); + assertArrayEquals( + new float[] {1, 1, 1}, leaf.getFloatVectorValues("f3d").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 (float[] vector : docVectors) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + w.commit(); // one segment per document + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); + assertEquals(3, values.size()); + assertEquals(2, groupSize(values)); // a's duplicate collapsed across segments + for (int ord = 0; ord < values.size(); ord++) { + assertArrayEquals(docVectors[ord], values.vectorValue(ord), 0f); + } + } + } + } + + /** Number of distinct vectors physically stored for a field's group. */ + private static int groupSize(KnnVectorValues values) { + return ((DedupVectorValues) values).getGroupView().size(); + } +} diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java new file mode 100644 index 000000000000..de12c3d5d261 --- /dev/null +++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java @@ -0,0 +1,91 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +import java.io.IOException; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.simpletext.SimpleTextKnnVectorsReader; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.junit.Ignore; + +/** + * Runs the standard KNN vectors format suite against the de-duplicating HNSW format. De-duplication + * behavior itself is covered by {@link TestDedupFlatVectorsFormat}. + */ +public class TestLucene106DedupHnswVectorsFormat extends BaseKnnVectorsFormatTestCase { + + private final KnnVectorsFormat format = new Lucene106DedupHnswVectorsFormat(); + + @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 SimpleTextKnnVectorsReader) { + assertEquals(0L, offHeap.size()); // all vectors are in memory + assertEquals(0L, totalByteSize); + } else { + 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()); + } + } + + /** + * This test indexes random vectors of small dimensions with high duplicates, checking that RAM + * usage is above a threshold. The RAM usage assumption breaks with the de-duplicating format. + */ + @Override + @Ignore + public void testWriterRamEstimate() {} + + /** The de-duplicating vector format does not attribute vectors to per-field writers. */ + @Override + @Ignore + public void testWriterByteVectorRamEstimate() {} +} 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 7556c0b020a5..a0c2693db895 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 @@ -2346,7 +2346,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(); @@ -2371,7 +2372,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) {