From 23ed23a017ed7935ad74dfb5fc0d41bb0caf7495 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Thu, 16 Jul 2026 03:23:54 +0000 Subject: [PATCH 01/12] Add a de-duplicating vector format --- lucene/CHANGES.txt | 3 + .../dedup/DedupFlatFieldVectorsWriter.java | 110 ++++ .../dedup/DedupFlatVectorsFormat.java | 132 ++++ .../dedup/DedupFlatVectorsReader.java | 330 ++++++++++ .../dedup/DedupFlatVectorsScorer.java | 170 ++++++ .../dedup/DedupFlatVectorsWriter.java | 163 +++++ .../lucene106/dedup/DedupFlushContext.java | 268 +++++++++ .../codecs/lucene106/dedup/DedupGroup.java | 100 ++++ .../lucene106/dedup/DedupMergeContext.java | 357 +++++++++++ .../codecs/lucene106/dedup/DedupUtil.java | 562 ++++++++++++++++++ .../Lucene106DedupHnswVectorsFormat.java | 255 ++++++++ .../codecs/lucene106/dedup/package-info.java | 26 + .../org.apache.lucene.codecs.KnnVectorsFormat | 1 + .../dedup/TestDedupFlatVectorsFormat.java | 186 ++++++ .../TestLucene106DedupHnswVectorsFormat.java | 91 +++ .../index/BaseKnnVectorsFormatTestCase.java | 5 +- 16 files changed, 2757 insertions(+), 2 deletions(-) create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java create mode 100644 lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java create mode 100644 lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java create mode 100644 lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java 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..a8e989a50f2b --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java @@ -0,0 +1,255 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.codecs.lucene106.dedup; + +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; +import org.apache.lucene.util.hnsw.HnswGraphBuilder; + +/** + * 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 { + public static final int VERSION_START = 0; + public static final int VERSION_CURRENT = VERSION_START; + + /** + * A maximum configurable maximum max conn. + * + *

NOTE: We eagerly populate `float[MAX_CONN*2]` and `int[MAX_CONN*2]`, so exceptionally large + * numbers here will use an inordinate amount of heap + */ + public static final int MAXIMUM_MAX_CONN = 512; + + /** Default number of maximum connections per node */ + public static final int DEFAULT_MAX_CONN = HnswGraphBuilder.DEFAULT_MAX_CONN; + + /** + * The maximum size of the queue to maintain while searching during graph construction. This + * maximum value preserves the ratio of the `DEFAULT_BEAM_WIDTH`/`DEFAULT_MAX_CONN` (i.e. `6.25 * + * 16 = 3200`). + */ + public static final int MAXIMUM_BEAM_WIDTH = 3200; + + /** + * Default number of the size of the queue maintained while searching during a graph construction. + */ + public static final int DEFAULT_BEAM_WIDTH = HnswGraphBuilder.DEFAULT_BEAM_WIDTH; + + /** Default to use single-thread merge */ + public static final int DEFAULT_NUM_MERGE_WORKER = 1; + + /** + * Minimum estimated search effort (in terms of expected visited nodes) required before building + * an HNSW graph for a segment. + * + *

This threshold is compared against the value produced by {@link + * org.apache.lucene.util.hnsw.HnswGraphSearcher#expectedVisitedNodes(int, int)}, which estimates + * how many nodes would be visited during a vector search based on the current graph size and + * {@code k} (neighbours to find). + * + *

If the estimated number of visited nodes falls below this threshold, HNSW graph construction + * is skipped for that segment - typically for small flushes or low document count segments - + * since the overhead of building the graph would outweigh its search benefits. + * + *

Default: {@code 100} + */ + public static final int HNSW_GRAPH_THRESHOLD = 100; + + /** + * Controls how many of the nearest neighbor candidates are connected to the new node. Defaults to + * {@link Lucene106DedupHnswVectorsFormat#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 Lucene106DedupHnswVectorsFormat#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 + * Lucene106DedupHnswVectorsFormat#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 + */ + 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 1024; + } + + @Override + public String toString() { + return "Lucene106DedupHnswVectorsFormat(name=Lucene106DedupHnswVectorsFormat, 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) { From 514f8984dd9f55aaaadad0979c2ebdc5e63e1977 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Fri, 17 Jul 2026 15:50:40 +0000 Subject: [PATCH 02/12] iter Small refactor of Lucene106DedupHnswVectorsFormat. --- .../Lucene106DedupHnswVectorsFormat.java | 72 +++++-------------- 1 file changed, 19 insertions(+), 53 deletions(-) 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 index a8e989a50f2b..ce4c6223bc68 100644 --- 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 @@ -16,6 +16,13 @@ */ 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; @@ -30,7 +37,6 @@ import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraphBuilder; /** * An HNSW vector format that de-duplicates raw vectors. @@ -45,62 +51,19 @@ * @lucene.experimental */ public final class Lucene106DedupHnswVectorsFormat extends KnnVectorsFormat { - public static final int VERSION_START = 0; - public static final int VERSION_CURRENT = VERSION_START; - - /** - * A maximum configurable maximum max conn. - * - *

NOTE: We eagerly populate `float[MAX_CONN*2]` and `int[MAX_CONN*2]`, so exceptionally large - * numbers here will use an inordinate amount of heap - */ - public static final int MAXIMUM_MAX_CONN = 512; - - /** Default number of maximum connections per node */ - public static final int DEFAULT_MAX_CONN = HnswGraphBuilder.DEFAULT_MAX_CONN; - - /** - * The maximum size of the queue to maintain while searching during graph construction. This - * maximum value preserves the ratio of the `DEFAULT_BEAM_WIDTH`/`DEFAULT_MAX_CONN` (i.e. `6.25 * - * 16 = 3200`). - */ - public static final int MAXIMUM_BEAM_WIDTH = 3200; - - /** - * Default number of the size of the queue maintained while searching during a graph construction. - */ - public static final int DEFAULT_BEAM_WIDTH = HnswGraphBuilder.DEFAULT_BEAM_WIDTH; - - /** Default to use single-thread merge */ - public static final int DEFAULT_NUM_MERGE_WORKER = 1; - - /** - * Minimum estimated search effort (in terms of expected visited nodes) required before building - * an HNSW graph for a segment. - * - *

This threshold is compared against the value produced by {@link - * org.apache.lucene.util.hnsw.HnswGraphSearcher#expectedVisitedNodes(int, int)}, which estimates - * how many nodes would be visited during a vector search based on the current graph size and - * {@code k} (neighbours to find). - * - *

If the estimated number of visited nodes falls below this threshold, HNSW graph construction - * is skipped for that segment - typically for small flushes or low document count segments - - * since the overhead of building the graph would outweigh its search benefits. - * - *

Default: {@code 100} - */ - public static final int HNSW_GRAPH_THRESHOLD = 100; + private static final String NAME = "Lucene106DedupHnswVectorsFormat"; /** * Controls how many of the nearest neighbor candidates are connected to the new node. Defaults to - * {@link Lucene106DedupHnswVectorsFormat#DEFAULT_MAX_CONN}. See {@link HnswGraph} for more - * details. + * {@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 Lucene106DedupHnswVectorsFormat#DEFAULT_BEAM_WIDTH}. See {@link + * node. Defaults to {@link + * org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat#DEFAULT_BEAM_WIDTH}. See {@link * HnswGraph} for details. */ private final int beamWidth; @@ -114,7 +77,7 @@ public final class Lucene106DedupHnswVectorsFormat extends KnnVectorsFormat { /** * 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 - * Lucene106DedupHnswVectorsFormat#HNSW_GRAPH_THRESHOLD}). + * org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat#HNSW_GRAPH_THRESHOLD}). * *

    *
  • 0 indicates that the graph is always built. @@ -181,7 +144,7 @@ public Lucene106DedupHnswVectorsFormat( * @param tinySegmentsThreshold the expected number of vector operations to return k nearest * neighbors of the current graph size */ - Lucene106DedupHnswVectorsFormat( + public Lucene106DedupHnswVectorsFormat( int maxConn, int beamWidth, int numMergeWorkers, @@ -237,12 +200,15 @@ public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException @Override public int getMaxDimensions(String fieldName) { - return 1024; + return DEFAULT_MAX_DIMENSIONS; } @Override public String toString() { - return "Lucene106DedupHnswVectorsFormat(name=Lucene106DedupHnswVectorsFormat, maxConn=" + return NAME + + "(name=" + + NAME + + ", maxConn=" + maxConn + ", beamWidth=" + beamWidth From 91936b9b43d7244b2663de254ef1c0715c59d754 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Wed, 22 Jul 2026 16:07:19 +0000 Subject: [PATCH 03/12] iter Small refactor of DedupFlatVectorsFormat and Lucene106DedupHnswVectorsFormat. --- .../codecs/lucene106/dedup/DedupFlatVectorsFormat.java | 6 +++--- .../lucene106/dedup/Lucene106DedupHnswVectorsFormat.java | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) 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 index 41d41be9905f..e019b8399b35 100644 --- 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 @@ -103,12 +103,12 @@ * @lucene.experimental */ final class DedupFlatVectorsFormat extends FlatVectorsFormat { - static final String NAME = "Lucene106DedupFlatVectorsFormat"; + static final String NAME = "DedupFlatVectorsFormat"; - static final String META_CODEC_NAME = "Lucene106DedupFlatVectorsFormatMeta"; + static final String META_CODEC_NAME = "DedupFlatVectorsFormatMeta"; static final String META_EXTENSION = "vdm"; - static final String VECTOR_DATA_CODEC_NAME = "Lucene106DedupFlatVectorsFormatVectorData"; + static final String VECTOR_DATA_CODEC_NAME = "DedupFlatVectorsFormatVectorData"; static final String VECTOR_DATA_EXTENSION = "vdd"; static final int VERSION_START = 0; 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 index ce4c6223bc68..e0e28cd35a6e 100644 --- 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 @@ -150,7 +150,7 @@ public Lucene106DedupHnswVectorsFormat( int numMergeWorkers, ExecutorService mergeExec, int tinySegmentsThreshold) { - super("Lucene106DedupHnswVectorsFormat"); + super(NAME); if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) { throw new IllegalArgumentException( "maxConn must be positive and less than or equal to " From cbce9f0f7803db4d52785f5b59fa021b8aa8a50e Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Mon, 27 Jul 2026 18:05:20 +0000 Subject: [PATCH 04/12] address comments --- .../dedup/DedupFlatVectorsScorer.java | 34 +++--- .../dedup/DedupFlatVectorsWriter.java | 22 ++-- .../codecs/lucene106/dedup/DedupGroup.java | 16 ++- .../lucene106/dedup/DedupMergeContext.java | 13 ++- .../codecs/lucene106/dedup/DedupUtil.java | 56 ++++----- .../Lucene106DedupHnswVectorsFormat.java | 2 + .../lucene99/Lucene99HnswVectorsWriter.java | 14 ++- .../dedup/TestDedupFlatVectorsFormat.java | 61 ++++++++-- .../TestLucene106DedupHnswVectorsFormat.java | 106 +++++++++++++++++- 9 files changed, 233 insertions(+), 91 deletions(-) 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 index 4449808e7bc5..658f13198d52 100644 --- 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 @@ -47,11 +47,11 @@ public RandomVectorScorerSupplier getRandomVectorScorerSupplier( VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues) throws IOException { if (vectorValues instanceof DedupVectorValues dedupValues) { - RandomVectorScorerSupplier delegate = + RandomVectorScorerSupplier fieldView = SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); RandomVectorScorerSupplier groupView = SCORER.getRandomVectorScorerSupplier(similarityFunction, dedupValues.getGroupView()); - return new RandomVectorScorerSupplierImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerSupplierImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); } return SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); } @@ -61,11 +61,11 @@ public RandomVectorScorer getRandomVectorScorer( VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, float[] target) throws IOException { if (vectorValues instanceof DedupVectorValues dedupValues) { - RandomVectorScorer delegate = + RandomVectorScorer fieldView = SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); RandomVectorScorer groupView = SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); - return new RandomVectorScorerImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); } return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } @@ -75,17 +75,17 @@ public RandomVectorScorer getRandomVectorScorer( VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, byte[] target) throws IOException { if (vectorValues instanceof DedupVectorValues dedupValues) { - RandomVectorScorer delegate = + RandomVectorScorer fieldView = SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); RandomVectorScorer groupView = SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); - return new RandomVectorScorerImpl(delegate, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); } return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } private record RandomVectorScorerSupplierImpl( - RandomVectorScorerSupplier delegate, + RandomVectorScorerSupplier fieldView, RandomVectorScorerSupplier groupView, OrdToVecOrd ordToVecOrd) implements RandomVectorScorerSupplier { @@ -93,24 +93,24 @@ private record RandomVectorScorerSupplierImpl( @Override public UpdateableRandomVectorScorer scorer() throws IOException { return new UpdateableRandomVectorScorerImpl( - delegate.scorer(), groupView.scorer(), ordToVecOrd); + fieldView.scorer(), groupView.scorer(), ordToVecOrd); } @Override public RandomVectorScorerSupplier copy() throws IOException { - return new RandomVectorScorerSupplierImpl(delegate.copy(), groupView.copy(), ordToVecOrd); + return new RandomVectorScorerSupplierImpl(fieldView.copy(), groupView.copy(), ordToVecOrd); } } private static class RandomVectorScorerImpl implements RandomVectorScorer { - private final RandomVectorScorer delegate; + private final RandomVectorScorer fieldView; private final RandomVectorScorer groupView; private final OrdToVecOrd ordToVecOrd; private int[] scratch; RandomVectorScorerImpl( - RandomVectorScorer delegate, RandomVectorScorer groupView, OrdToVecOrd ordToVecOrd) { - this.delegate = delegate; + RandomVectorScorer fieldView, RandomVectorScorer groupView, OrdToVecOrd ordToVecOrd) { + this.fieldView = fieldView; this.groupView = groupView; this.ordToVecOrd = ordToVecOrd; this.scratch = new int[SCRATCH_SIZE]; @@ -118,12 +118,12 @@ private static class RandomVectorScorerImpl implements RandomVectorScorer { @Override public int ordToDoc(int ord) { - return delegate.ordToDoc(ord); + return fieldView.ordToDoc(ord); } @Override public Bits getAcceptOrds(Bits acceptDocs) { - return delegate.getAcceptOrds(acceptDocs); + return fieldView.getAcceptOrds(acceptDocs); } @Override @@ -144,7 +144,7 @@ public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOExcep @Override public int maxOrd() { - return delegate.maxOrd(); + return fieldView.maxOrd(); } } @@ -154,10 +154,10 @@ private static final class UpdateableRandomVectorScorerImpl extends RandomVector private final OrdToVecOrd ordToVecOrd; UpdateableRandomVectorScorerImpl( - UpdateableRandomVectorScorer delegate, + UpdateableRandomVectorScorer fieldView, UpdateableRandomVectorScorer groupView, OrdToVecOrd ordToVecOrd) { - super(delegate, groupView, ordToVecOrd); + super(fieldView, groupView, ordToVecOrd); this.groupView = groupView; this.ordToVecOrd = ordToVecOrd; } 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 index 81099818dac3..3d0ed67345ea 100644 --- 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 @@ -74,24 +74,21 @@ final class DedupFlatVectorsWriter extends FlatVectorsWriter { 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); + this.meta = state.directory.createOutput(metaFileName, state.context); + this.vectorData = state.directory.createOutput(vectorDataFileName, state.context); CodecUtil.writeIndexHeader( - m, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); + meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); CodecUtil.writeIndexHeader( - v, + vectorData, 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); + IOUtils.closeWhileHandlingException(this); } } } @@ -127,13 +124,8 @@ public void finish() throws IOException { finishMerge(); } - if (meta != null) { - CodecUtil.writeFooter(meta); - } - - if (vectorData != null) { - CodecUtil.writeFooter(vectorData); - } + CodecUtil.writeFooter(meta); + CodecUtil.writeFooter(vectorData); } @Override 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 index 4d46d0121e0e..9b869da5a340 100644 --- 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 @@ -38,9 +38,12 @@ abstract sealed class DedupGroup implements Accountable DedupFlushContext.FloatGroup, DedupMergeContext.DedupMergeGroup { - private static final int ORD_NOT_FOUND = -1; - + /** + * Map for vector hash -> group ord. Only a hint and not the ground truth due to possibility of + * hash collisions, where a full equality check must be performed. + */ private final LongIntHashMap hashToOrdHint; + private final List vectors; private final ObjectCursor current; // reuse from addUnique @@ -71,17 +74,18 @@ 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) { + int ordIndex = hashToOrdHint.indexOf(hash); + if (ordIndex < 0) { // not found groupOrd = vectors.size(); ownedVector = copy(vectorValue); // only for unique vectors hashToOrdHint.put(hash, groupOrd); vectors.add(ownedVector); break; } else { - T other = vectors.get(ordHint); + int ord = hashToOrdHint.indexGet(ordIndex); + T other = vectors.get(ord); if (equals(vectorValue, other)) { - groupOrd = ordHint; + groupOrd = ord; ownedVector = other; break; } 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 index db10db4ca5f3..41a80344147d 100644 --- 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 @@ -99,20 +99,21 @@ void finish(IndexOutput meta, IndexOutput vectorData) throws IOException { int groupOrd = 0; for (Map.Entry> entry : fieldGroups.entrySet()) { GroupKey groupKey = entry.getKey(); - long vectorDataOffset = alignBytes(vectorData, groupKey.encoding()); + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); + + long vectorDataOffset = alignBytes(vectorData, encoding); DedupMergeGroup mergeGroup = - switch (groupKey.encoding()) { + switch (encoding) { case BYTE -> new ByteGroup(); - case FLOAT32 -> new FloatGroup(groupKey.dimension()); + case FLOAT32 -> new FloatGroup(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; @@ -157,7 +158,7 @@ void processField(FieldData fieldData, IndexOutput vectorData) throws IOExceptio int groupSize = size(); // add vector to group - ObjectCursor cursor = super.addUnique(vector); + ObjectCursor cursor = addUnique(vector); if (cursor.index == groupSize) { // new addition // already on-heap, write immediately to avoid another IO read byte[] bytes = serialize(groupSize); 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 index a76c25747017..98e6af6483d1 100644 --- 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 @@ -59,7 +59,9 @@ final class DedupUtil { private static final int ORD_TO_VEC_ALIGN_BYTES = 4; - // TODO: Evaluate using fewer bits. + // TODO: This is the number of bits used to write each group ordinal in the index-backed per-field + // OrdToVecOrd mapping. Evaluate using fewer bits to reduce index size, at the expense of + // costlier lookups. private static final int ORD_TO_VEC_BITS_PER_VALUE = 32; static final int ORD_UNKNOWN = -1; @@ -77,10 +79,10 @@ record GroupKey(int dimension, VectorEncoding encoding) { * 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. + *

    Every instance is backed by two views: the {@code fieldView} maps ordinals to docs and + * drives iteration (one entry per document), while the {@code groupView} holds the de-duplicated + * vectors (one entry per distinct vector). {@code ordToVecOrd} translates a document ordinal into + * its group ordinal. */ sealed interface DedupVectorValues { /** The dense view over distinct vectors, indexed by group ordinal. */ @@ -305,7 +307,7 @@ static ByteVectorValues loadDedupBytes( long ordToVecSize) throws IOException { - final OffHeapByteVectorValues delegate = + final OffHeapByteVectorValues fieldView = OffHeapByteVectorValues.load( function, vectorsScorer, configuration, BYTE, dimension, 0, 0, vectorData); @@ -314,21 +316,21 @@ static ByteVectorValues loadDedupBytes( dimension, groupSize, vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), - delegate.getVectorByteLength(), + fieldView.getVectorByteLength(), vectorsScorer, function); final OrdToVecOrd ordToVecOrd = new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); - return new ByteImpl(vectorsScorer, function, delegate, groupView, ordToVecOrd); + return new ByteImpl(vectorsScorer, function, fieldView, 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 fieldView; private final ByteVectorValues groupView; private final OrdToVecOrd ordToVecOrd; private int[] scratch; @@ -336,12 +338,12 @@ private static final class ByteImpl extends ByteVectorValues implements DedupVec ByteImpl( FlatVectorsScorer vectorsScorer, VectorSimilarityFunction function, - ByteVectorValues delegate, + ByteVectorValues fieldView, ByteVectorValues groupView, OrdToVecOrd ordToVecOrd) { this.vectorsScorer = vectorsScorer; this.function = function; - this.delegate = delegate; + this.fieldView = fieldView; this.groupView = groupView; this.ordToVecOrd = ordToVecOrd; this.scratch = new int[SCRATCH_SIZE]; @@ -359,7 +361,7 @@ public OrdToVecOrd getOrdToVecOrd() { @Override public int ordToDoc(int ord) { - return delegate.ordToDoc(ord); + return fieldView.ordToDoc(ord); } @Override @@ -380,23 +382,23 @@ public byte[] vectorValue(int ord) throws IOException { @Override public int dimension() { - return delegate.dimension(); + return fieldView.dimension(); } @Override public int size() { - return delegate.size(); + return fieldView.size(); } @Override public ByteImpl copy() throws IOException { return new ByteImpl( - vectorsScorer, function, delegate.copy(), groupView.copy(), ordToVecOrd.copy()); + vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); } @Override public DocIndexIterator iterator() { - return delegate.iterator(); + return fieldView.iterator(); } @Override @@ -439,7 +441,7 @@ static FloatVectorValues loadDedupFloats( long ordToVecSize) throws IOException { - final OffHeapFloatVectorValues delegate = + final OffHeapFloatVectorValues fieldView = OffHeapFloatVectorValues.load( function, vectorsScorer, configuration, FLOAT32, dimension, 0, 0, vectorData); @@ -448,21 +450,21 @@ static FloatVectorValues loadDedupFloats( dimension, groupSize, vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), - delegate.getVectorByteLength(), + fieldView.getVectorByteLength(), vectorsScorer, function); final OrdToVecOrd ordToVecOrd = new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); - return new FloatImpl(vectorsScorer, function, delegate, groupView, ordToVecOrd); + return new FloatImpl(vectorsScorer, function, fieldView, 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 fieldView; private final FloatVectorValues groupView; private final OrdToVecOrd ordToVecOrd; private int[] scratch; @@ -470,12 +472,12 @@ private static final class FloatImpl extends FloatVectorValues implements DedupV FloatImpl( FlatVectorsScorer vectorsScorer, VectorSimilarityFunction function, - FloatVectorValues delegate, + FloatVectorValues fieldView, FloatVectorValues groupView, OrdToVecOrd ordToVecOrd) { this.vectorsScorer = vectorsScorer; this.function = function; - this.delegate = delegate; + this.fieldView = fieldView; this.groupView = groupView; this.ordToVecOrd = ordToVecOrd; this.scratch = new int[SCRATCH_SIZE]; @@ -493,7 +495,7 @@ public OrdToVecOrd getOrdToVecOrd() { @Override public int ordToDoc(int ord) { - return delegate.ordToDoc(ord); + return fieldView.ordToDoc(ord); } @Override @@ -514,23 +516,23 @@ public float[] vectorValue(int ord) throws IOException { @Override public int dimension() { - return delegate.dimension(); + return fieldView.dimension(); } @Override public int size() { - return delegate.size(); + return fieldView.size(); } @Override public FloatImpl copy() throws IOException { return new FloatImpl( - vectorsScorer, function, delegate.copy(), groupView.copy(), ordToVecOrd.copy()); + vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); } @Override public DocIndexIterator iterator() { - return delegate.iterator(); + return fieldView.iterator(); } @Override 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 index e0e28cd35a6e..d18d6bbf76fe 100644 --- 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 @@ -182,6 +182,8 @@ public Lucene106DedupHnswVectorsFormat( @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + // TODO: Can we have an HNSW writer that uses de-duplication information to speed up graph + // construction? return new Lucene99HnswVectorsWriter( state, maxConn, diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java index 0b0bf5bda506..667e9acf0a34 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java @@ -236,10 +236,10 @@ public void finish() throws IOException { @Override public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; + long total = SHALLOW_RAM_BYTES_USED + flatVectorWriter.ramBytesUsed(); for (FieldWriter field : fields) { // the field tracks the delegate field usage - total += field.ramBytesUsed(); + total += field.ownRamBytesUsed(); } return total; } @@ -803,13 +803,17 @@ OnHeapHnswGraph getGraph() throws IOException { } } - @Override - public long ramBytesUsed() { - long total = SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed(); + private long ownRamBytesUsed() { + long total = SHALLOW_SIZE; if (hnswGraphBuilder != null) { total += hnswGraphBuilder.getGraph().ramBytesUsed(); } return total; } + + @Override + public long ramBytesUsed() { + return ownRamBytesUsed() + flatFieldVectorsWriter.ramBytesUsed(); + } } } diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java index 96f3aa295fda..870c681c2eb5 100644 --- 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 @@ -17,13 +17,17 @@ package org.apache.lucene.codecs.lucene106.dedup; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.hamcrest.Matchers.instanceOf; +import org.apache.lucene.codecs.KnnVectorsReader; 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.CodecReader; import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; @@ -95,10 +99,13 @@ public void testByteDuplicatesWithinField() throws Exception { /** 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}}; + // Vectors that are close to each other in bit representations. + float[][] distinctDocVectors = { + {+0f}, {-0f}, {Math.nextUp(0f)}, {Math.nextDown(0f)}, {1f}, {Math.nextUp(1f)} + }; try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, config())) { - for (float[] vector : docVectors) { + for (float[] vector : distinctDocVectors) { Document doc = new Document(); doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); w.addDocument(doc); @@ -106,8 +113,43 @@ public void testDistinctVectorsAllStored() throws Exception { w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); - assertEquals(3, values.size()); - assertEquals(3, groupSize(values)); + assertEquals(distinctDocVectors.length, values.size()); + assertEquals(distinctDocVectors.length, groupSize(values)); + } + } + } + + /** Check off-heap size of de-duplicated vectors. */ + public void testOffHeapSize() throws Exception { + float[] a = {1, 2, 3, 4}; + float[] b = {5, 6, 7, 8}; + float[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (float[] vector : docVectors) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + assertThat(leafReader, instanceOf(CodecReader.class)); + + FieldInfo fieldInfo = leafReader.getFieldInfos().fieldInfo("f"); + KnnVectorsReader knnVectorsReader = ((CodecReader) leafReader).getVectorReader(); + knnVectorsReader = knnVectorsReader.unwrapReaderForField("f"); + + long expectedOffHeapSize = + (docVectors.length * Integer.BYTES) // ordToVec mapping + + (a.length + b.length) * Float.BYTES; // raw vector size + + assertEquals( + expectedOffHeapSize, + knnVectorsReader + .getOffHeapByteSize(fieldInfo) + .get("vdd") // vector data extension + .longValue()); } } } @@ -136,20 +178,21 @@ public void testDuplicatesAcrossFieldsShareGroup() throws Exception { /** Fields differing in dimension use separate groups, even for otherwise similar vectors. */ public void testDifferentDimensionsUseSeparateGroups() throws Exception { + float[] vector1 = {1, 1}; + float[] vector2 = {1, 1, 0}; try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, config())) { Document doc = new Document(); - doc.add(new KnnFloatVectorField("f2d", new float[] {1, 1}, EUCLIDEAN)); - doc.add(new KnnFloatVectorField("f3d", new float[] {1, 1, 1}, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f2d", vector1, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f3d", vector2, EUCLIDEAN)); w.addDocument(doc); w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader leaf = getOnlyLeafReader(reader); 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); + assertArrayEquals(vector1, leaf.getFloatVectorValues("f2d").vectorValue(0), 0f); + assertArrayEquals(vector2, leaf.getFloatVectorValues("f3d").vectorValue(0), 0f); } } } 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 index de12c3d5d261..565bc5f0558f 100644 --- 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 @@ -16,15 +16,38 @@ */ package org.apache.lucene.codecs.lucene106.dedup; +import static org.hamcrest.Matchers.greaterThan; + import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.simpletext.SimpleTextKnnVectorsReader; import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DocValuesSkipIndexType; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.IndexOptions; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.StringHelper; +import org.apache.lucene.util.Version; import org.junit.Ignore; /** @@ -76,13 +99,84 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx } } - /** - * 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. - */ + /** Near copy of the original test, this one checks for size of unique vector count. */ @Override - @Ignore - public void testWriterRamEstimate() {} + @SuppressWarnings("unchecked") + public void testWriterRamEstimate() throws IOException { + final FieldInfos fieldInfos = new FieldInfos(new FieldInfo[0]); + final Directory dir = newDirectory(); + Codec codec = Codec.getDefault(); + final SegmentInfo si = + new SegmentInfo( + dir, + Version.LATEST, + Version.LATEST, + "0", + 10000, + false, + false, + codec, + Collections.emptyMap(), + StringHelper.randomId(), + new HashMap<>(), + null); + final SegmentWriteState state = + new SegmentWriteState( + InfoStream.getDefault(), dir, si, fieldInfos, null, newIOContext(random())); + final KnnVectorsFormat format = codec.knnVectorsFormat(); + try (KnnVectorsWriter writer = format.fieldsWriter(state)) { + final long ramBytesUsed = writer.ramBytesUsed(); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + ++dim; + } + int numDocs = atLeast(100); + Set unique = new HashSet<>(); + KnnFieldVectorsWriter fieldWriter = + (KnnFieldVectorsWriter) + writer.addField( + new FieldInfo( + "fieldA", + 0, + false, + false, + false, + IndexOptions.NONE, + DocValuesType.NONE, + DocValuesSkipIndexType.NONE, + -1, + Map.of(), + 0, + 0, + 0, + dim, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.DOT_PRODUCT, + false, + false)); + for (int i = 0; i < numDocs; i++) { + float[] vector = randomVector(dim); + unique.add(new FloatVector(vector)); + fieldWriter.addValue(i, vector); + } + final long ramBytesUsed2 = writer.ramBytesUsed(); + assertThat(ramBytesUsed2, greaterThan(ramBytesUsed)); + assertThat(ramBytesUsed2, greaterThan((long) dim * unique.size() * Float.BYTES)); + } + dir.close(); + } + + private record FloatVector(float[] vector) { + @Override + public boolean equals(Object obj) { + return obj instanceof FloatVector(float[] other) && Arrays.equals(vector, other); + } + + @Override + public int hashCode() { + return Arrays.hashCode(vector); + } + } /** The de-duplicating vector format does not attribute vectors to per-field writers. */ @Override From caa1d32ddddffdce683ad4b3e54211186a46b506 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Mon, 27 Jul 2026 20:54:07 +0000 Subject: [PATCH 05/12] add fp16 support --- .../dedup/DedupFlatVectorsFormat.java | 3 +- .../dedup/DedupFlatVectorsReader.java | 29 ++++ .../dedup/DedupFlatVectorsScorer.java | 14 ++ .../lucene106/dedup/DedupFlushContext.java | 59 ++++++-- .../codecs/lucene106/dedup/DedupGroup.java | 1 + .../lucene106/dedup/DedupMergeContext.java | 86 +++++++++-- .../codecs/lucene106/dedup/DedupUtil.java | 141 +++++++++++++++++- .../dedup/TestDedupFlatVectorsFormat.java | 26 ++++ .../TestLucene106DedupHnswVectorsFormat.java | 81 +++++++++- 9 files changed, 410 insertions(+), 30 deletions(-) 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 index e019b8399b35..ffee331c0537 100644 --- 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 @@ -36,7 +36,8 @@ *

    .vdd (vector de-dup data) file

    * *
      - *
    • For each group, its distinct vectors, aligned to 4 bytes (BYTE) or 64 bytes (FLOAT32). + *
    • For each group, its distinct vectors, aligned to 4 bytes (BYTE) or 64 bytes (FLOAT32 and + * FLOAT16). *
    • For each field: *
        *
      • The sparse-encoding data (only when some documents lack the field): DocIds encoded by 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 index bf016dad7647..b522d5c3079b 100644 --- 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 @@ -23,10 +23,12 @@ 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.loadDedupFloat16s; 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.FLOAT16; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import java.io.IOException; @@ -43,6 +45,7 @@ import org.apache.lucene.index.CorruptIndexException; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.MergePolicy; @@ -246,6 +249,13 @@ public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) thr return vectorsScorer.getRandomVectorScorer(entry.fieldInfo.function(), vectorValues, target); } + @Override + public RandomVectorScorer getRandomVectorScorer(String field, short[] target) throws IOException { + FieldEntry entry = getEntry(field, FLOAT16); + Float16VectorValues vectorValues = getFloat16VectorValues(entry); + return vectorsScorer.getRandomVectorScorer(entry.fieldInfo.function(), vectorValues, target); + } + @Override public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException { CodecUtil.checksumEntireFile(vectorData, merge); @@ -289,6 +299,25 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException { return getByteVectorValues(getEntry(field, BYTE)); } + private Float16VectorValues getFloat16VectorValues(FieldEntry entry) throws IOException { + return loadDedupFloat16s( + vectorsScorer, + entry.fieldInfo.function(), + entry.fieldInfo.ordToDoc(), + entry.fieldInfo.dimension(), + entry.groupInfo.groupSize(), + vectorData, + entry.groupInfo.vectorDataOffset(), + entry.groupInfo.vectorDataSize(), + entry.fieldInfo.ordToVecOffset(), + entry.fieldInfo.ordToVecSize()); + } + + @Override + public Float16VectorValues getFloat16VectorValues(String field) throws IOException { + return getFloat16VectorValues(getEntry(field, FLOAT16)); + } + @Override public FlatVectorsReader getMergeInstance() { // TODO: Can we improve performance using strictly sequential IO? 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 index 658f13198d52..8650549c4acf 100644 --- 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 @@ -84,6 +84,20 @@ public RandomVectorScorer getRandomVectorScorer( return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } + @Override + public RandomVectorScorer getRandomVectorScorer( + VectorSimilarityFunction similarityFunction, KnnVectorValues vectorValues, short[] target) + throws IOException { + if (vectorValues instanceof DedupVectorValues dedupValues) { + RandomVectorScorer fieldView = + SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + RandomVectorScorer groupView = + SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); + } + return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); + } + private record RandomVectorScorerSupplierImpl( RandomVectorScorerSupplier fieldView, RandomVectorScorerSupplier groupView, 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 index 9e53f0636ffd..3abb6dfb4685 100644 --- 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 @@ -17,7 +17,6 @@ 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; @@ -28,6 +27,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.FloatBuffer; +import java.nio.ShortBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -71,6 +71,7 @@ private static DedupGroup getGroup(GroupKey groupKey) { return switch (groupKey.encoding()) { case BYTE -> new ByteGroup(groupKey.dimension()); case FLOAT32 -> new FloatGroup(groupKey.dimension()); + case FLOAT16 -> new Float16Group(groupKey.dimension()); }; } @@ -209,7 +210,6 @@ static final class FloatGroup extends DedupGroup { private final long ramBytesPerVector; private final byte[] bytes; private final FloatBuffer buffer; - private int lastOrd; FloatGroup(int dimension) { int length = dimension * Float.BYTES; @@ -219,15 +219,11 @@ static final class FloatGroup extends DedupGroup { + 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); } @@ -243,10 +239,53 @@ public float[] copy(float[] vectorValue) { @Override byte[] serialize(int ord) { - if (ord != lastOrd) { - buffer.put(0, get(ord)); - lastOrd = ord; - } + buffer.put(0, get(ord)); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + size() * ramBytesPerVector; + } + } + + static final class Float16Group extends DedupGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Group.class); + + private final long ramBytesPerVector; + private final byte[] bytes; + private final ShortBuffer buffer; + + Float16Group(int dimension) { + int length = dimension * Short.BYTES; + this.ramBytesPerVector = + RamUsageEstimator.NUM_BYTES_OBJECT_REF + + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + + length; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asShortBuffer(); + } + + @Override + public long hash(short[] vector) { + buffer.put(0, vector); + return hashBytes(bytes); + } + + @Override + public boolean equals(short[] vector, short[] other) { + return Arrays.equals(vector, other); + } + + @Override + public short[] copy(short[] vectorValue) { + return vectorValue.clone(); + } + + @Override + byte[] serialize(int ord) { + buffer.put(0, get(ord)); return bytes; } 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 index 9b869da5a340..4662be8ceb38 100644 --- 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 @@ -36,6 +36,7 @@ abstract sealed class DedupGroup implements Accountable permits DedupFlushContext.ByteGroup, DedupFlushContext.FloatGroup, + DedupFlushContext.Float16Group, DedupMergeContext.DedupMergeGroup { /** 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 index 41a80344147d..5ccaf3469816 100644 --- 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 @@ -17,7 +17,6 @@ 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; @@ -28,6 +27,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.FloatBuffer; +import java.nio.ShortBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -44,6 +44,7 @@ import org.apache.lucene.index.DocIDMerger; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.MergeState; @@ -108,6 +109,7 @@ void finish(IndexOutput meta, IndexOutput vectorData) throws IOException { switch (encoding) { case BYTE -> new ByteGroup(); case FLOAT32 -> new FloatGroup(dimension); + case FLOAT16 -> new Float16Group(dimension); }; for (FieldData fieldData : entry.getValue()) { @@ -172,7 +174,7 @@ void processField(FieldData fieldData, IndexOutput vectorData) throws IOExceptio } } - record ByteVector(ByteVectorValues values, int ord) implements IOSupplier { + private record ByteVector(ByteVectorValues values, int ord) implements IOSupplier { private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(ByteVector.class); @@ -227,7 +229,7 @@ public long ramBytesUsed() { } } - record FloatVector(FloatVectorValues values, int ord) implements IOSupplier { + private record FloatVector(FloatVectorValues values, int ord) implements IOSupplier { private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(FloatVector.class); @@ -243,13 +245,11 @@ private static final class FloatGroup extends DedupMergeGroup sub) { @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); } @@ -288,10 +285,7 @@ public FloatVector copy(FloatVector vectorValue) { @Override byte[] serialize(int ord) throws IOException { - if (ord != lastOrd) { - buffer.put(0, get(ord).get()); - lastOrd = ord; - } + buffer.put(0, get(ord).get()); return bytes; } @@ -301,6 +295,73 @@ public long ramBytesUsed() { } } + private record Float16Vector(Float16VectorValues values, int ord) implements IOSupplier { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Vector.class); + + @Override + public short[] get() throws IOException { + return values.vectorValue(ord); + } + } + + private static final class Float16Group + extends DedupMergeGroup { + private static final long SHALLOW_SIZE = + RamUsageEstimator.shallowSizeOfInstance(Float16Group.class); + + private final byte[] bytes; + private final ShortBuffer buffer; + + Float16Group(int dimension) { + int length = dimension * Short.BYTES; + this.bytes = new byte[length]; + this.buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN).asShortBuffer(); + } + + @Override + Float16Vector vectorFrom(Sub sub) { + return new Float16Vector(sub.values, sub.iterator.index()); + } + + @Override + public long hash(Float16Vector vector) throws IOException { + buffer.put(0, vector.get()); + return hashBytes(bytes); + } + + @Override + public boolean equals(Float16Vector vector, Float16Vector other) throws IOException { + // Fast path: two docs from the same dedup source share a vector iff they map to the same + // group ordinal, so we can compare ordinals without reading the vectors back. + if (vector.values == other.values && vector.values instanceof DedupVectorValues dedup) { + OrdToVecOrd ordToVecOrd = dedup.getOrdToVecOrd(); + return ordToVecOrd.get(vector.ord) == ordToVecOrd.get(other.ord); + } + short[] a = vector.get(); + if (vector.values == other.values) { + a = a.clone(); // same reader reuses one buffer; copy before reading the other vector + } + return Arrays.equals(a, other.get()); + } + + @Override + public Float16Vector copy(Float16Vector vectorValue) { + return vectorValue; + } + + @Override + byte[] serialize(int ord) throws IOException { + buffer.put(0, get(ord).get()); + return bytes; + } + + @Override + public long ramBytesUsed() { + return SHALLOW_SIZE + super.ramBytesUsed() + size() * Float16Vector.SHALLOW_SIZE; + } + } + private record FieldData( FieldInfo fieldInfo, GroupKey groupKey, @@ -344,6 +405,7 @@ private static DocIDMerger> getVectorMerger( switch (fieldInfo.getVectorEncoding()) { case BYTE -> mergeState.knnVectorsReaders[i].getByteVectorValues(fieldInfo.name); case FLOAT32 -> mergeState.knnVectorsReaders[i].getFloatVectorValues(fieldInfo.name); + case FLOAT16 -> mergeState.knnVectorsReaders[i].getFloat16VectorValues(fieldInfo.name); }; if (vectorValues == null) { 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 index 98e6af6483d1..3af1623a5f75 100644 --- 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 @@ -17,6 +17,7 @@ package org.apache.lucene.codecs.lucene106.dedup; import static org.apache.lucene.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT16; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.StringHelper.GOOD_FAST_HASH_SEED; import static org.apache.lucene.util.StringHelper.murmurhash3_x64_128; @@ -24,11 +25,13 @@ import java.io.IOException; import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; import org.apache.lucene.codecs.lucene95.OffHeapByteVectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloat16VectorValues; import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues; import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.VectorEncoding; @@ -64,8 +67,6 @@ final class DedupUtil { // costlier lookups. 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). */ @@ -235,7 +236,7 @@ static long alignBytes(IndexOutput output, VectorEncoding encoding) throws IOExc int alignBytes = switch (encoding) { case BYTE -> 4; - case FLOAT32 -> 64; + case FLOAT32, FLOAT16 -> 64; }; return output.alignFilePointer(alignBytes); } @@ -561,4 +562,138 @@ public Bulk bulk(DocIdSetIterator matchingDocs) { }; } } + + static Float16VectorValues loadDedupFloat16s( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupSize, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long ordToVecOffset, + long ordToVecSize) + throws IOException { + + final OffHeapFloat16VectorValues fieldView = + OffHeapFloat16VectorValues.load( + function, vectorsScorer, configuration, FLOAT16, dimension, 0, 0, vectorData); + + final OffHeapFloat16VectorValues groupView = + new OffHeapFloat16VectorValues.DenseOffHeapVectorValues( + dimension, + groupSize, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final OrdToVecOrd ordToVecOrd = + new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + + return new Float16Impl(vectorsScorer, function, fieldView, groupView, ordToVecOrd); + } + + /** {@link DedupVectorValues} over float16 vectors. */ + private static final class Float16Impl extends Float16VectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final Float16VectorValues fieldView; + private final Float16VectorValues groupView; + private final OrdToVecOrd ordToVecOrd; + private int[] scratch; + + Float16Impl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + Float16VectorValues fieldView, + Float16VectorValues groupView, + OrdToVecOrd ordToVecOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.ordToVecOrd = ordToVecOrd; + this.scratch = new int[SCRATCH_SIZE]; + } + + @Override + public Float16VectorValues getGroupView() { + return groupView; + } + + @Override + public OrdToVecOrd getOrdToVecOrd() { + return ordToVecOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = ordToVecOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public short[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(ordToVecOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public Float16Impl copy() throws IOException { + return new Float16Impl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + if (size() == 0) { + return null; + } + Float16Impl copy = copy(); + DocIndexIterator 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/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java index 870c681c2eb5..af94f8da0d0e 100644 --- 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 @@ -23,11 +23,13 @@ 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.KnnFloat16VectorField; import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.Float16VectorValues; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; @@ -73,6 +75,30 @@ public void testFloatDuplicatesWithinField() throws Exception { } } + /** Repeated float16 vectors within a field are stored once but still read back per document. */ + public void testFloat16DuplicatesWithinField() throws Exception { + short[] a = {Float.floatToFloat16(1f), Float.floatToFloat16(2f), Float.floatToFloat16(3f)}; + short[] b = {Float.floatToFloat16(4f), Float.floatToFloat16(5f), Float.floatToFloat16(6f)}; + short[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (short[] vector : docVectors) { + Document doc = new Document(); + doc.add(new KnnFloat16VectorField("f", vector, EUCLIDEAN)); + w.addDocument(doc); + } + w.forceMerge(1); + try (DirectoryReader reader = DirectoryReader.open(w)) { + Float16VectorValues values = getOnlyLeafReader(reader).getFloat16VectorValues("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)); + } + } + } + } + /** 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}; 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 index 565bc5f0558f..8715e9638593 100644 --- 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 @@ -48,7 +48,6 @@ import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.StringHelper; import org.apache.lucene.util.Version; -import org.junit.Ignore; /** * Runs the standard KNN vectors format suite against the de-duplicating HNSW format. De-duplication @@ -178,8 +177,82 @@ public int hashCode() { } } - /** The de-duplicating vector format does not attribute vectors to per-field writers. */ + /** Near copy of the original test, this one checks for size of unique vector count. */ @Override - @Ignore - public void testWriterByteVectorRamEstimate() {} + @SuppressWarnings("unchecked") + public void testWriterByteVectorRamEstimate() throws IOException { + final FieldInfos fieldInfos = new FieldInfos(new FieldInfo[0]); + final Directory dir = newDirectory(); + Codec codec = Codec.getDefault(); + final SegmentInfo si = + new SegmentInfo( + dir, + Version.LATEST, + Version.LATEST, + "0", + 10000, + false, + false, + codec, + Collections.emptyMap(), + StringHelper.randomId(), + new HashMap<>(), + null); + final SegmentWriteState state = + new SegmentWriteState( + InfoStream.getDefault(), dir, si, fieldInfos, null, newIOContext(random())); + final KnnVectorsFormat format = codec.knnVectorsFormat(); + try (KnnVectorsWriter writer = format.fieldsWriter(state)) { + final long ramBytesUsed = writer.ramBytesUsed(); + int dim = random().nextInt(64) + 1; + if (dim % 2 == 1) { + ++dim; + } + int numDocs = atLeast(100); + Set unique = new HashSet<>(); + KnnFieldVectorsWriter fieldWriter = + (KnnFieldVectorsWriter) + writer.addField( + new FieldInfo( + "fieldA", + 0, + false, + false, + false, + IndexOptions.NONE, + DocValuesType.NONE, + DocValuesSkipIndexType.NONE, + -1, + Map.of(), + 0, + 0, + 0, + dim, + VectorEncoding.BYTE, + VectorSimilarityFunction.DOT_PRODUCT, + false, + false)); + for (int i = 0; i < numDocs; i++) { + byte[] vector = randomVector8(dim); + unique.add(new ByteVector(vector)); + fieldWriter.addValue(i, vector); + } + final long ramBytesUsed2 = writer.ramBytesUsed(); + assertThat(ramBytesUsed2, greaterThan(ramBytesUsed)); + assertThat(ramBytesUsed2, greaterThan((long) dim * unique.size() * Byte.BYTES)); + } + dir.close(); + } + + private record ByteVector(byte[] vector) { + @Override + public boolean equals(Object obj) { + return obj instanceof ByteVector(byte[] other) && Arrays.equals(vector, other); + } + + @Override + public int hashCode() { + return Arrays.hashCode(vector); + } + } } From 4052a01e6b1bf0bcc78d1aec9a28ace464103352 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Tue, 28 Jul 2026 18:21:59 +0000 Subject: [PATCH 06/12] address comments --- lucene/CHANGES.txt | 3 +- .../dedup/DedupFlatFieldVectorsWriter.java | 12 +- .../dedup/DedupFlatVectorsFormat.java | 34 ++-- .../dedup/DedupFlatVectorsReader.java | 23 +-- .../dedup/DedupFlatVectorsScorer.java | 48 ++--- .../lucene106/dedup/DedupFlushContext.java | 19 +- .../lucene106/dedup/DedupMergeContext.java | 24 +-- .../codecs/lucene106/dedup/DedupUtil.java | 179 +++++++++--------- .../Lucene106DedupHnswVectorsFormat.java | 17 +- .../dedup/TestDedupFlatVectorsFormat.java | 2 +- 10 files changed, 196 insertions(+), 165 deletions(-) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 0ef338d6f1b3..f67fdb42e2e4 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -331,7 +331,8 @@ New Features * GITHUB#16383: Add fp16 vector encoding support. (Pulkit Gupta) * 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) + each distinct full-precision vector once, shared across all documents and fields that reference it. + (Kaival Parikh) Improvements --------------------- diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java index b31545251137..a4ffbbc57608 100644 --- 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 @@ -38,7 +38,7 @@ final class DedupFlatFieldVectorsWriter extends FlatFieldVectorsWriter { private final DedupGroup group; private final DocsWithFieldSet docsWithFieldSet; private final List vectors; - private final IntArrayList ordToVecOrd; + private final IntArrayList fieldOrdToGroupOrd; private int lastDocID; private boolean finished; @@ -46,7 +46,7 @@ final class DedupFlatFieldVectorsWriter extends FlatFieldVectorsWriter { this.group = group; this.docsWithFieldSet = new DocsWithFieldSet(); this.vectors = new ArrayList<>(); - this.ordToVecOrd = new IntArrayList(); + this.fieldOrdToGroupOrd = new IntArrayList(); this.lastDocID = -1; this.finished = false; } @@ -61,8 +61,8 @@ public DocsWithFieldSet getDocsWithFieldSet() { return docsWithFieldSet; } - IntArrayList getOrdToVecOrd() { - return ordToVecOrd; + IntArrayList getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; } @Override @@ -97,7 +97,7 @@ public void addValue(int docID, T vectorValue) throws IOException { ObjectCursor cursor = group.addUnique(vectorValue); vectors.add(cursor.value); // owned vector value - ordToVecOrd.add(cursor.index); // index in group + fieldOrdToGroupOrd.add(cursor.index); // index in group } @Override @@ -105,6 +105,6 @@ public long ramBytesUsed() { return SHALLOW_SIZE + docsWithFieldSet.ramBytesUsed() + (long) vectors.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF - + ordToVecOrd.ramBytesUsed(); + + fieldOrdToGroupOrd.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 index ffee331c0537..ba5e94db1481 100644 --- 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 @@ -28,16 +28,19 @@ * *

        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. + * map to that group) reference it; each field then keeps an {@code fieldOrdToGroupOrd} map from its + * document ordinal to the group ordinal of the shared vector. This is well suited to indexes with + * repeated vectors, e.g. several fields derived from the same embedding, or heavily duplicated + * content. * *

        .vdd (vector de-dup data) file

        * *
          *
        • For each group, its distinct vectors, aligned to 4 bytes (BYTE) or 64 bytes (FLOAT32 and - * FLOAT16). + * FLOAT16). This a best-effort alignment, because Arm Neoverse machines incur a performance + * penalty in reading data not aligned to 64 bytes. This penalty may be incurred for float + * vectors that do not have a dimension of a multiple of 16 (equivalent to 64 bytes), because + * the alignment will not hold for all vectors in the file. *
        • For each field: *
            *
          • The sparse-encoding data (only when some documents lack the field): DocIds encoded by @@ -45,8 +48,8 @@ * 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 + *
          • The {@code fieldOrdToGroupOrd} map (aligned to 4 bytes): one entry per document + * ordinal giving the group ordinal of the shared vector, packed by {@link * org.apache.lucene.util.packed.DirectWriter}. *
          *
        @@ -76,8 +79,8 @@ *
      • 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 + *
      • [int64] offset to this field's {@code fieldOrdToGroupOrd} map in the .vdd file + *
      • [int64] length of this field's {@code fieldOrdToGroupOrd} map, in bytes *
      * *

      also terminated by [int32] {@code -1}. @@ -91,14 +94,15 @@ *

        *
      • 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. + * plus {@code O(N)} for the per-document references and {@code fieldOrdToGroupOrd} entries. *
      • Merge: expected {@code O(N * d)} time; distinct vectors are written to disk as soon * as they are first seen rather than buffered, so heap stays {@code O(N)} (the per-field - * {@code 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. + * {@code fieldOrdToGroupOrd} maps and light per-vector handles) with no {@code O(U * d)} + * term. When a source segment is itself in this format, equality is decided by comparing + * group ordinals in {@code O(1)} without reading the vectors back. + *
      • Reading: both the vectors and the {@code fieldOrdToGroupOrd} map stay off-heap; a + * read resolves a document ordinal to its vector via one extra {@code fieldOrdToGroupOrd} + * lookup. *
      * * @lucene.experimental 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 index b522d5c3079b..14d12d513f6a 100644 --- 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 @@ -63,7 +63,7 @@ /** * 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. + * view backed by its group's shared vectors and a {@code fieldOrdToGroupOrd} translation map. * * @lucene.experimental */ @@ -129,21 +129,21 @@ private void readMetaBody(ChecksumIndexInput meta, FieldInfos fieldInfos) throws throw new CorruptIndexException("Invalid field number: " + fieldInfo.fieldNumber(), meta); } else if (fieldInfo.function() != info.getVectorSimilarityFunction()) { throw new CorruptIndexException( - "Invalid vector function: indexed=" + "Inconsistent vector function: indexed=" + fieldInfo.function() + ", actual=" + info.getVectorSimilarityFunction(), meta); } else if (fieldInfo.dimension() != info.getVectorDimension()) { throw new CorruptIndexException( - "Invalid vector dimension: indexed=" + "Inconsistent vector dimension: indexed=" + fieldInfo.dimension() + ", actual=" + info.getVectorDimension(), meta); } else if (fieldInfo.encoding() != info.getVectorEncoding()) { throw new CorruptIndexException( - "Invalid vector encoding: indexed=" + "Inconsistent vector encoding: indexed=" + fieldInfo.encoding() + ", actual=" + info.getVectorEncoding(), @@ -271,8 +271,8 @@ private FloatVectorValues getFloatVectorValues(FieldEntry entry) throws IOExcept vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), - entry.fieldInfo.ordToVecOffset(), - entry.fieldInfo.ordToVecSize()); + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); } @Override @@ -290,8 +290,8 @@ private ByteVectorValues getByteVectorValues(FieldEntry entry) throws IOExceptio vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), - entry.fieldInfo.ordToVecOffset(), - entry.fieldInfo.ordToVecSize()); + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); } @Override @@ -309,8 +309,8 @@ private Float16VectorValues getFloat16VectorValues(FieldEntry entry) throws IOEx vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), - entry.fieldInfo.ordToVecOffset(), - entry.fieldInfo.ordToVecSize()); + entry.fieldInfo.fieldOrdToGroupOrdOffset(), + entry.fieldInfo.fieldOrdToGroupOrdSize()); } @Override @@ -347,7 +347,8 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { } // TODO: This is an over-estimation. return Map.of( - VECTOR_DATA_EXTENSION, entry.fieldInfo.ordToVecSize() + entry.groupInfo.vectorDataSize()); + VECTOR_DATA_EXTENSION, + entry.fieldInfo.fieldOrdToGroupOrdSize() + entry.groupInfo.vectorDataSize()); } private record FieldEntry(ReadFieldInfo fieldInfo, GroupInfo groupInfo) { 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 index 8650549c4acf..6885567e37b3 100644 --- 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 @@ -16,13 +16,13 @@ */ package org.apache.lucene.codecs.lucene106.dedup; -import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.SCRATCH_SIZE; +import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; import java.io.IOException; import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.DedupVectorValues; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.OrdToVecOrd; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrd; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.util.ArrayUtil; @@ -33,8 +33,8 @@ /** * 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()}. + * delegates vector operations to the underlying {@link DedupVectorValues#getGroupView()}, mapping + * document ordinals to group ordinals via {@link DedupVectorValues#getFieldOrdToGroupOrd()}. * * @lucene.experimental */ @@ -51,7 +51,8 @@ public RandomVectorScorerSupplier getRandomVectorScorerSupplier( SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); RandomVectorScorerSupplier groupView = SCORER.getRandomVectorScorerSupplier(similarityFunction, dedupValues.getGroupView()); - return new RandomVectorScorerSupplierImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerSupplierImpl( + fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); } return SCORER.getRandomVectorScorerSupplier(similarityFunction, vectorValues); } @@ -65,7 +66,7 @@ public RandomVectorScorer getRandomVectorScorer( SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); RandomVectorScorer groupView = SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); - return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); } return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } @@ -79,7 +80,7 @@ public RandomVectorScorer getRandomVectorScorer( SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); RandomVectorScorer groupView = SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); - return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); } return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } @@ -93,7 +94,7 @@ public RandomVectorScorer getRandomVectorScorer( SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); RandomVectorScorer groupView = SCORER.getRandomVectorScorer(similarityFunction, dedupValues.getGroupView(), target); - return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getOrdToVecOrd()); + return new RandomVectorScorerImpl(fieldView, groupView, dedupValues.getFieldOrdToGroupOrd()); } return SCORER.getRandomVectorScorer(similarityFunction, vectorValues, target); } @@ -101,33 +102,36 @@ public RandomVectorScorer getRandomVectorScorer( private record RandomVectorScorerSupplierImpl( RandomVectorScorerSupplier fieldView, RandomVectorScorerSupplier groupView, - OrdToVecOrd ordToVecOrd) + FieldOrdToGroupOrd fieldOrdToGroupOrd) implements RandomVectorScorerSupplier { @Override public UpdateableRandomVectorScorer scorer() throws IOException { return new UpdateableRandomVectorScorerImpl( - fieldView.scorer(), groupView.scorer(), ordToVecOrd); + fieldView.scorer(), groupView.scorer(), fieldOrdToGroupOrd); } @Override public RandomVectorScorerSupplier copy() throws IOException { - return new RandomVectorScorerSupplierImpl(fieldView.copy(), groupView.copy(), ordToVecOrd); + return new RandomVectorScorerSupplierImpl( + fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd); } } private static class RandomVectorScorerImpl implements RandomVectorScorer { private final RandomVectorScorer fieldView; private final RandomVectorScorer groupView; - private final OrdToVecOrd ordToVecOrd; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; private int[] scratch; RandomVectorScorerImpl( - RandomVectorScorer fieldView, RandomVectorScorer groupView, OrdToVecOrd ordToVecOrd) { + RandomVectorScorer fieldView, + RandomVectorScorer groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { this.fieldView = fieldView; this.groupView = groupView; - this.ordToVecOrd = ordToVecOrd; - this.scratch = new int[SCRATCH_SIZE]; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; } @Override @@ -142,7 +146,7 @@ public Bits getAcceptOrds(Bits acceptDocs) { @Override public float score(int node) throws IOException { - return groupView.score(ordToVecOrd.get(node)); + return groupView.score(fieldOrdToGroupOrd.get(node)); } @Override @@ -151,7 +155,7 @@ public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOExcep scratch = ArrayUtil.grow(scratch, nodes.length); } for (int i = 0; i < numNodes; i++) { - scratch[i] = ordToVecOrd.get(nodes[i]); + scratch[i] = fieldOrdToGroupOrd.get(nodes[i]); } return groupView.bulkScore(scratch, scores, numNodes); } @@ -165,20 +169,20 @@ public int maxOrd() { private static final class UpdateableRandomVectorScorerImpl extends RandomVectorScorerImpl implements UpdateableRandomVectorScorer { private final UpdateableRandomVectorScorer groupView; - private final OrdToVecOrd ordToVecOrd; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; UpdateableRandomVectorScorerImpl( UpdateableRandomVectorScorer fieldView, UpdateableRandomVectorScorer groupView, - OrdToVecOrd ordToVecOrd) { - super(fieldView, groupView, ordToVecOrd); + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + super(fieldView, groupView, fieldOrdToGroupOrd); this.groupView = groupView; - this.ordToVecOrd = ordToVecOrd; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; } @Override public void setScoringOrdinal(int node) throws IOException { - groupView.setScoringOrdinal(ordToVecOrd.get(node)); + groupView.setScoringOrdinal(fieldOrdToGroupOrd.get(node)); } } } 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 index 3abb6dfb4685..68273b9f3c77 100644 --- 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 @@ -35,11 +35,11 @@ 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.FieldOrdToGroupOrd; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdMappedArrayList; 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; @@ -132,20 +132,21 @@ void flush(IndexOutput meta, IndexOutput vectorData, int maxDoc, Sorter.DocMap s for (FieldData fieldData : fieldDataList) { fieldData.fieldWriter.finish(); - IntArrayList ordToVecOrd = fieldData.fieldWriter.getOrdToVecOrd(); - int vectorCount = ordToVecOrd.elementsCount; + IntArrayList fieldOrdToGroupOrd = fieldData.fieldWriter.getFieldOrdToGroupOrd(); + int vectorCount = fieldOrdToGroupOrd.elementsCount; DocsWithFieldSet docs; - OrdToVecOrd ordToVecFinal; + FieldOrdToGroupOrd fieldOrdToGroupOrdFinal; if (sortMap == null) { docs = fieldData.fieldWriter.getDocsWithFieldSet(); - ordToVecFinal = new OrdToVecOrdArrayList(ordToVecOrd); + fieldOrdToGroupOrdFinal = new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); } else { DocsWithFieldSet oldDocs = fieldData.fieldWriter.getDocsWithFieldSet(); docs = new DocsWithFieldSet(); int[] new2OldOrd = new int[vectorCount]; KnnVectorsWriter.mapOldOrdToNewOrd(oldDocs, sortMap, null, new2OldOrd, docs); - ordToVecFinal = new OrdToVecOrdMappedArrayList(new2OldOrd, ordToVecOrd); + fieldOrdToGroupOrdFinal = + new FieldOrdToGroupOrdMappedArrayList(new2OldOrd, fieldOrdToGroupOrd); } WriteFieldInfo fieldInfo = @@ -158,7 +159,7 @@ void flush(IndexOutput meta, IndexOutput vectorData, int maxDoc, Sorter.DocMap s vectorCount, maxDoc, docs, - ordToVecFinal); + fieldOrdToGroupOrdFinal); writeFieldInfo(meta, vectorData, fieldInfo); } 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 index 5ccaf3469816..4e25900740e9 100644 --- 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 @@ -35,10 +35,10 @@ 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.FieldOrdToGroupOrd; +import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; 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; @@ -137,10 +137,10 @@ void finish(IndexOutput meta, IndexOutput vectorData) throws IOException { fieldData.fieldInfo.getVectorDimension(), fieldData.fieldInfo.getVectorEncoding(), groupOrds.get(fieldData.groupKey), - fieldData.ordToVecOrd.elementsCount, + fieldData.fieldOrdToGroupOrd.elementsCount, fieldData.maxDoc, fieldData.docsWithFieldSet, - new OrdToVecOrdArrayList(fieldData.ordToVecOrd)); + new FieldOrdToGroupOrdArrayList(fieldData.fieldOrdToGroupOrd)); writeFieldInfo(meta, vectorData, fieldInfo); } @@ -169,7 +169,7 @@ void processField(FieldData fieldData, IndexOutput vectorData) throws IOExceptio // record hit and ord in group fieldData.docsWithFieldSet.add(next.mappedDocID); - fieldData.ordToVecOrd.add(cursor.index); + fieldData.fieldOrdToGroupOrd.add(cursor.index); } } } @@ -203,8 +203,8 @@ 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); + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); } byte[] a = vector.get(); if (vector.values == other.values) { @@ -268,8 +268,8 @@ 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); + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); } float[] a = vector.get(); if (vector.values == other.values) { @@ -335,8 +335,8 @@ public boolean equals(Float16Vector vector, Float16Vector other) throws IOExcept // 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); + FieldOrdToGroupOrd fieldOrdToGroupOrd = dedup.getFieldOrdToGroupOrd(); + return fieldOrdToGroupOrd.get(vector.ord) == fieldOrdToGroupOrd.get(other.ord); } short[] a = vector.get(); if (vector.values == other.values) { @@ -366,7 +366,7 @@ private record FieldData( FieldInfo fieldInfo, GroupKey groupKey, DocsWithFieldSet docsWithFieldSet, - IntArrayList ordToVecOrd, + IntArrayList fieldOrdToGroupOrd, DocIDMerger merger, int maxDoc) { 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 index 3af1623a5f75..e10bc93e9837 100644 --- 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 @@ -56,18 +56,19 @@ */ final class DedupUtil { - private static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; + private static final int ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT = 16; private static final int END_MARKER = -1; - private static final int ORD_TO_VEC_ALIGN_BYTES = 4; + private static final int FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES = 4; // TODO: This is the number of bits used to write each group ordinal in the index-backed per-field - // OrdToVecOrd mapping. Evaluate using fewer bits to reduce index size, at the expense of + // FieldOrdToGroupOrd mapping. Evaluate using fewer bits to reduce index size, at the expense of // costlier lookups. - private static final int ORD_TO_VEC_BITS_PER_VALUE = 32; + private static final int FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE = 32; - static final int SCRATCH_SIZE = 16; + /** Initial allocation size for internal re-used int[] scratch buffers. */ + static final int SCRATCH_INITIAL_SIZE = 16; /** Key used to group vectors (dimension + encoding). */ record GroupKey(int dimension, VectorEncoding encoding) { @@ -82,25 +83,25 @@ record GroupKey(int dimension, VectorEncoding encoding) { * *

      Every instance is backed by two views: the {@code fieldView} maps ordinals to docs and * drives iteration (one entry per document), while the {@code groupView} holds the de-duplicated - * vectors (one entry per distinct vector). {@code ordToVecOrd} translates a document ordinal into - * its group ordinal. + * vectors (one entry per distinct vector). {@code fieldOrdToGroupOrd} translates a document + * ordinal in the field into its group ordinal. */ sealed interface DedupVectorValues { /** The dense view over distinct vectors, indexed by group ordinal. */ KnnVectorValues getGroupView(); /** Maps a per-document ordinal to its group ordinal in {@link #getGroupView()}. */ - OrdToVecOrd getOrdToVecOrd(); + FieldOrdToGroupOrd getFieldOrdToGroupOrd(); } /** * Maps a field's per-document ordinal to the ordinal of its (shared) vector within the group. * Backed on-heap while writing and off-heap while reading. */ - sealed interface OrdToVecOrd { + sealed interface FieldOrdToGroupOrd { int get(int ord); - OrdToVecOrd copy() throws IOException; + FieldOrdToGroupOrd copy() throws IOException; } record GroupInfo( @@ -149,7 +150,7 @@ record WriteFieldInfo( int vectorCount, int maxDoc, DocsWithFieldSet docs, - OrdToVecOrd ordToVecOrd) {} + FieldOrdToGroupOrd fieldOrdToGroupOrd) {} static void writeFieldInfo(IndexOutput meta, IndexOutput vectorData, WriteFieldInfo fieldInfo) throws IOException { @@ -163,25 +164,26 @@ static void writeFieldInfo(IndexOutput meta, IndexOutput vectorData, WriteFieldI // write ordToDoc OrdToDocDISIReaderConfiguration.writeStoredMeta( - DIRECT_MONOTONIC_BLOCK_SHIFT, + ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, fieldInfo.vectorCount, fieldInfo.maxDoc, fieldInfo.docs); - // write ordToVec - long ordToVecOffset = vectorData.alignFilePointer(ORD_TO_VEC_ALIGN_BYTES); + // write fieldOrdToGroupOrd + long fieldOrdToGroupOrdOffset = vectorData.alignFilePointer(FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES); DirectWriter writer = - DirectWriter.getInstance(vectorData, fieldInfo.vectorCount, ORD_TO_VEC_BITS_PER_VALUE); + DirectWriter.getInstance( + vectorData, fieldInfo.vectorCount, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); for (int i = 0; i < fieldInfo.vectorCount; i++) { - writer.add(fieldInfo.ordToVecOrd.get(i)); + writer.add(fieldInfo.fieldOrdToGroupOrd.get(i)); } writer.finish(); - long ordToVecSize = vectorData.getFilePointer() - ordToVecOffset; + long fieldOrdToGroupOrdSize = vectorData.getFilePointer() - fieldOrdToGroupOrdOffset; - meta.writeLong(ordToVecOffset); - meta.writeLong(ordToVecSize); + meta.writeLong(fieldOrdToGroupOrdOffset); + meta.writeLong(fieldOrdToGroupOrdSize); } static void writeEndOfFields(IndexOutput meta) throws IOException { @@ -196,8 +198,8 @@ record ReadFieldInfo( int groupOrd, int vectorCount, OrdToDocDISIReaderConfiguration ordToDoc, - long ordToVecOffset, - long ordToVecSize) {} + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) {} static ReadFieldInfo readFieldInfo(IndexInput meta) throws IOException { @@ -213,8 +215,8 @@ static ReadFieldInfo readFieldInfo(IndexInput meta) throws IOException { int vectorCount = meta.readInt(); OrdToDocDISIReaderConfiguration ordToDoc = OrdToDocDISIReaderConfiguration.fromStoredMeta(meta, vectorCount); - long ordToVecOffset = meta.readLong(); - long ordToVecSize = meta.readLong(); + long fieldOrdToGroupOrdOffset = meta.readLong(); + long fieldOrdToGroupOrdSize = meta.readLong(); return new ReadFieldInfo( fieldNumber, @@ -224,8 +226,8 @@ static ReadFieldInfo readFieldInfo(IndexInput meta) throws IOException { groupOrd, vectorCount, ordToDoc, - ordToVecOffset, - ordToVecSize); + fieldOrdToGroupOrdOffset, + fieldOrdToGroupOrdSize); } static long hashBytes(byte[] bytes) { @@ -242,46 +244,52 @@ static long alignBytes(IndexOutput output, VectorEncoding encoding) throws IOExc } /** On-heap map used during a flush, backed directly by the buffered ordinals. */ - record OrdToVecOrdArrayList(IntArrayList ordToVecOrd) implements OrdToVecOrd { + record FieldOrdToGroupOrdArrayList(IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + @Override public int get(int ord) { - return ordToVecOrd.get(ord); + return fieldOrdToGroupOrd.get(ord); } @Override - public OrdToVecOrd copy() { - return new OrdToVecOrdArrayList(ordToVecOrd); + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); } } /** On-heap map used during a sorted flush, indirecting through a new-to-old ordinal map. */ - record OrdToVecOrdMappedArrayList(int[] map, IntArrayList ordToVecOrd) implements OrdToVecOrd { + record FieldOrdToGroupOrdMappedArrayList(int[] map, IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + @Override public int get(int ord) { - return ordToVecOrd.get(map[ord]); + return fieldOrdToGroupOrd.get(map[ord]); } @Override - public OrdToVecOrd copy() { - return new OrdToVecOrdMappedArrayList(map, ordToVecOrd); + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdMappedArrayList(map, fieldOrdToGroupOrd); } } /** Off-heap map used while reading, backed by a {@link DirectReader}. */ - static final class OrdToVecOrdOffHeap implements OrdToVecOrd { + static final class FieldOrdToGroupOrdOffHeap implements FieldOrdToGroupOrd { private final IndexInput vectorData; - private final long ordToVecOffset; - private final long ordToVecSize; + private final long fieldOrdToGroupOrdOffset; + private final long fieldOrdToGroupOrdSize; private final LongValues values; - OrdToVecOrdOffHeap(IndexInput vectorData, long ordToVecOffset, long ordToVecSize) + FieldOrdToGroupOrdOffHeap( + IndexInput vectorData, long fieldOrdToGroupOrdOffset, long fieldOrdToGroupOrdSize) throws IOException { this.vectorData = vectorData; - this.ordToVecOffset = ordToVecOffset; - this.ordToVecSize = ordToVecSize; + this.fieldOrdToGroupOrdOffset = fieldOrdToGroupOrdOffset; + this.fieldOrdToGroupOrdSize = fieldOrdToGroupOrdSize; - RandomAccessInput slice = vectorData.randomAccessSlice(ordToVecOffset, ordToVecSize); - this.values = DirectReader.getInstance(slice, ORD_TO_VEC_BITS_PER_VALUE); + RandomAccessInput slice = + vectorData.randomAccessSlice(fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + this.values = DirectReader.getInstance(slice, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); } @Override @@ -290,8 +298,9 @@ public int get(int v) { } @Override - public OrdToVecOrd copy() throws IOException { - return new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + public FieldOrdToGroupOrd copy() throws IOException { + return new FieldOrdToGroupOrdOffHeap( + vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); } } @@ -304,8 +313,8 @@ static ByteVectorValues loadDedupBytes( IndexInput vectorData, long vectorDataOffset, long vectorDataSize, - long ordToVecOffset, - long ordToVecSize) + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) throws IOException { final OffHeapByteVectorValues fieldView = @@ -321,10 +330,10 @@ static ByteVectorValues loadDedupBytes( vectorsScorer, function); - final OrdToVecOrd ordToVecOrd = - new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - return new ByteImpl(vectorsScorer, function, fieldView, groupView, ordToVecOrd); + return new ByteImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); } /** {@link DedupVectorValues} over byte vectors. */ @@ -333,7 +342,7 @@ private static final class ByteImpl extends ByteVectorValues implements DedupVec private final VectorSimilarityFunction function; private final ByteVectorValues fieldView; private final ByteVectorValues groupView; - private final OrdToVecOrd ordToVecOrd; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; private int[] scratch; ByteImpl( @@ -341,13 +350,13 @@ private static final class ByteImpl extends ByteVectorValues implements DedupVec VectorSimilarityFunction function, ByteVectorValues fieldView, ByteVectorValues groupView, - OrdToVecOrd ordToVecOrd) { + FieldOrdToGroupOrd fieldOrdToGroupOrd) { this.vectorsScorer = vectorsScorer; this.function = function; this.fieldView = fieldView; this.groupView = groupView; - this.ordToVecOrd = ordToVecOrd; - this.scratch = new int[SCRATCH_SIZE]; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; } @Override @@ -356,8 +365,8 @@ public ByteVectorValues getGroupView() { } @Override - public OrdToVecOrd getOrdToVecOrd() { - return ordToVecOrd; + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; } @Override @@ -371,14 +380,14 @@ public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); } for (int i = 0; i < numOrds; i++) { - scratch[i] = ordToVecOrd.get(ordsToPrefetch[i]); + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); } groupView.prefetch(scratch, numOrds); } @Override public byte[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(ordToVecOrd.get(ord)); + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); } @Override @@ -394,7 +403,7 @@ public int size() { @Override public ByteImpl copy() throws IOException { return new ByteImpl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); } @Override @@ -438,8 +447,8 @@ static FloatVectorValues loadDedupFloats( IndexInput vectorData, long vectorDataOffset, long vectorDataSize, - long ordToVecOffset, - long ordToVecSize) + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) throws IOException { final OffHeapFloatVectorValues fieldView = @@ -455,10 +464,10 @@ static FloatVectorValues loadDedupFloats( vectorsScorer, function); - final OrdToVecOrd ordToVecOrd = - new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - return new FloatImpl(vectorsScorer, function, fieldView, groupView, ordToVecOrd); + return new FloatImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); } /** {@link DedupVectorValues} over float vectors. */ @@ -467,7 +476,7 @@ private static final class FloatImpl extends FloatVectorValues implements DedupV private final VectorSimilarityFunction function; private final FloatVectorValues fieldView; private final FloatVectorValues groupView; - private final OrdToVecOrd ordToVecOrd; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; private int[] scratch; FloatImpl( @@ -475,13 +484,13 @@ private static final class FloatImpl extends FloatVectorValues implements DedupV VectorSimilarityFunction function, FloatVectorValues fieldView, FloatVectorValues groupView, - OrdToVecOrd ordToVecOrd) { + FieldOrdToGroupOrd fieldOrdToGroupOrd) { this.vectorsScorer = vectorsScorer; this.function = function; this.fieldView = fieldView; this.groupView = groupView; - this.ordToVecOrd = ordToVecOrd; - this.scratch = new int[SCRATCH_SIZE]; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; } @Override @@ -490,8 +499,8 @@ public FloatVectorValues getGroupView() { } @Override - public OrdToVecOrd getOrdToVecOrd() { - return ordToVecOrd; + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; } @Override @@ -505,14 +514,14 @@ public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); } for (int i = 0; i < numOrds; i++) { - scratch[i] = ordToVecOrd.get(ordsToPrefetch[i]); + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); } groupView.prefetch(scratch, numOrds); } @Override public float[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(ordToVecOrd.get(ord)); + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); } @Override @@ -528,7 +537,7 @@ public int size() { @Override public FloatImpl copy() throws IOException { return new FloatImpl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); } @Override @@ -572,8 +581,8 @@ static Float16VectorValues loadDedupFloat16s( IndexInput vectorData, long vectorDataOffset, long vectorDataSize, - long ordToVecOffset, - long ordToVecSize) + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) throws IOException { final OffHeapFloat16VectorValues fieldView = @@ -589,10 +598,10 @@ static Float16VectorValues loadDedupFloat16s( vectorsScorer, function); - final OrdToVecOrd ordToVecOrd = - new OrdToVecOrdOffHeap(vectorData, ordToVecOffset, ordToVecSize); + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - return new Float16Impl(vectorsScorer, function, fieldView, groupView, ordToVecOrd); + return new Float16Impl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); } /** {@link DedupVectorValues} over float16 vectors. */ @@ -601,7 +610,7 @@ private static final class Float16Impl extends Float16VectorValues implements De private final VectorSimilarityFunction function; private final Float16VectorValues fieldView; private final Float16VectorValues groupView; - private final OrdToVecOrd ordToVecOrd; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; private int[] scratch; Float16Impl( @@ -609,13 +618,13 @@ private static final class Float16Impl extends Float16VectorValues implements De VectorSimilarityFunction function, Float16VectorValues fieldView, Float16VectorValues groupView, - OrdToVecOrd ordToVecOrd) { + FieldOrdToGroupOrd fieldOrdToGroupOrd) { this.vectorsScorer = vectorsScorer; this.function = function; this.fieldView = fieldView; this.groupView = groupView; - this.ordToVecOrd = ordToVecOrd; - this.scratch = new int[SCRATCH_SIZE]; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; } @Override @@ -624,8 +633,8 @@ public Float16VectorValues getGroupView() { } @Override - public OrdToVecOrd getOrdToVecOrd() { - return ordToVecOrd; + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; } @Override @@ -639,14 +648,14 @@ public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); } for (int i = 0; i < numOrds; i++) { - scratch[i] = ordToVecOrd.get(ordsToPrefetch[i]); + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); } groupView.prefetch(scratch, numOrds); } @Override public short[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(ordToVecOrd.get(ord)); + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); } @Override @@ -662,7 +671,7 @@ public int size() { @Override public Float16Impl copy() throws IOException { return new Float16Impl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), ordToVecOrd.copy()); + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); } @Override 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 index d18d6bbf76fe..066bae973c69 100644 --- 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 @@ -44,9 +44,20 @@ *

      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. + * all documents that reference it. + * + *

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

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

      If you customize this format, be sure to share the same instance of the underlying + * {@link DedupFlatVectorsFormat} to de-duplicate raw vectors correctly. * * @lucene.experimental */ 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 index af94f8da0d0e..97be6c89e051 100644 --- 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 @@ -167,7 +167,7 @@ public void testOffHeapSize() throws Exception { knnVectorsReader = knnVectorsReader.unwrapReaderForField("f"); long expectedOffHeapSize = - (docVectors.length * Integer.BYTES) // ordToVec mapping + (docVectors.length * Integer.BYTES) // fieldOrdToGroupOrd mapping + (a.length + b.length) * Float.BYTES; // raw vector size assertEquals( From 4970f93a7dcd7c302cb6f12868a28211bf4ba73e Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Tue, 28 Jul 2026 19:37:30 +0000 Subject: [PATCH 07/12] fix failing tests --- .../dedup/TestDedupFlatVectorsFormat.java | 89 ++++++++++++++----- 1 file changed, 68 insertions(+), 21 deletions(-) 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 index 97be6c89e051..31a5e590064e 100644 --- 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 @@ -17,6 +17,7 @@ package org.apache.lucene.codecs.lucene106.dedup; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.hamcrest.Matchers.arrayContainingInAnyOrder; import static org.hamcrest.Matchers.instanceOf; import org.apache.lucene.codecs.KnnVectorsReader; @@ -25,6 +26,7 @@ import org.apache.lucene.document.KnnByteVectorField; import org.apache.lucene.document.KnnFloat16VectorField; import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.NumericDocValuesField; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DirectoryReader; @@ -35,6 +37,7 @@ import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NumericDocValues; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.TestUtil; @@ -58,19 +61,30 @@ public void testFloatDuplicatesWithinField() throws Exception { 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) { + for (int ord = 0; ord < docVectors.length; ord++) { Document doc = new Document(); - doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloatVectorField("f", docVectors[ord], EUCLIDEAN)); w.addDocument(doc); } w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { - FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.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); + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord), 0f); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); } } } @@ -82,19 +96,30 @@ public void testFloat16DuplicatesWithinField() throws Exception { short[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, config())) { - for (short[] vector : docVectors) { + for (int ord = 0; ord < docVectors.length; ord++) { Document doc = new Document(); - doc.add(new KnnFloat16VectorField("f", vector, EUCLIDEAN)); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloat16VectorField("f", docVectors[ord], EUCLIDEAN)); w.addDocument(doc); } w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { - Float16VectorValues values = getOnlyLeafReader(reader).getFloat16VectorValues("f"); + LeafReader leafReader = getOnlyLeafReader(reader); + Float16VectorValues values = leafReader.getFloat16VectorValues("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)); + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord)); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); } } } @@ -106,19 +131,30 @@ public void testByteDuplicatesWithinField() throws Exception { byte[][] docVectors = {a, a, b, a, b}; try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, config())) { - for (byte[] vector : docVectors) { + for (int ord = 0; ord < docVectors.length; ord++) { Document doc = new Document(); - doc.add(new KnnByteVectorField("f", vector, EUCLIDEAN)); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnByteVectorField("f", docVectors[ord], EUCLIDEAN)); w.addDocument(doc); } w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { - ByteVectorValues values = getOnlyLeafReader(reader).getByteVectorValues("f"); + LeafReader leafReader = getOnlyLeafReader(reader); + ByteVectorValues values = leafReader.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)); + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord)); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); } } } @@ -230,20 +266,31 @@ public void testDuplicatesAcrossSegmentsDedupOnMerge() throws Exception { 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) { + for (int ord = 0; ord < docVectors.length; ord++) { Document doc = new Document(); - doc.add(new KnnFloatVectorField("f", vector, EUCLIDEAN)); + doc.add(new NumericDocValuesField("id", ord)); + doc.add(new KnnFloatVectorField("f", docVectors[ord], EUCLIDEAN)); w.addDocument(doc); w.commit(); // one segment per document } w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { - FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); - assertEquals(3, values.size()); + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(docVectors.length, 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); + NumericDocValues docValues = leafReader.getNumericDocValues("id"); + Integer[] expectedOrds = new Integer[docVectors.length]; + Integer[] ordsSeen = new Integer[docVectors.length]; + for (int ord = 0; ord < docVectors.length; ord++) { + int docId = values.ordToDoc(ord); + assertTrue("id does not exist for docId=" + docId, docValues.advanceExact(docId)); + int originalOrd = (int) docValues.longValue(); + assertArrayEquals(docVectors[originalOrd], values.vectorValue(ord), 0f); + expectedOrds[ord] = ord; + ordsSeen[ord] = originalOrd; } + assertThat("all vectors not seen", ordsSeen, arrayContainingInAnyOrder(expectedOrds)); } } } From d1294e3df72dbb153e99e784dd6648b99037a53d Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Tue, 28 Jul 2026 20:00:14 +0000 Subject: [PATCH 08/12] move format to sandbox --- lucene/CHANGES.txt | 2 +- .../org.apache.lucene.codecs.KnnVectorsFormat | 1 - .../dedup/DedupFlatFieldVectorsWriter.java | 2 +- .../codecs}/dedup/DedupFlatVectorsFormat.java | 2 +- .../codecs}/dedup/DedupFlatVectorsReader.java | 30 +++++++++---------- .../codecs}/dedup/DedupFlatVectorsScorer.java | 8 ++--- .../codecs}/dedup/DedupFlatVectorsWriter.java | 12 ++++---- .../codecs}/dedup/DedupFlushContext.java | 26 ++++++++-------- .../sandbox/codecs}/dedup/DedupGroup.java | 2 +- .../codecs/dedup/DedupHnswVectorsFormat.java} | 16 +++++----- .../codecs}/dedup/DedupMergeContext.java | 26 ++++++++-------- .../sandbox/codecs}/dedup/DedupUtil.java | 2 +- .../sandbox/codecs}/dedup/package-info.java | 6 ++-- .../org.apache.lucene.codecs.KnnVectorsFormat | 1 + .../dedup/TestDedupFlatVectorsFormat.java | 12 ++++---- .../dedup/TestDedupHnswVectorsFormat.java} | 12 +++----- 16 files changed, 78 insertions(+), 82 deletions(-) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlatFieldVectorsWriter.java (98%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlatVectorsFormat.java (99%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlatVectorsReader.java (91%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlatVectorsScorer.java (96%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlatVectorsWriter.java (89%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupFlushContext.java (90%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupGroup.java (98%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java => sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java} (95%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupMergeContext.java (93%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/DedupUtil.java (99%) rename lucene/{core/src/java/org/apache/lucene/codecs/lucene106 => sandbox/src/java/org/apache/lucene/sandbox/codecs}/dedup/package-info.java (80%) rename lucene/{core/src/test/org/apache/lucene/codecs/lucene106 => sandbox/src/test/org/apache/lucene/sandbox/codecs}/dedup/TestDedupFlatVectorsFormat.java (96%) rename lucene/{core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java => sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java} (94%) diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index f67fdb42e2e4..beaf872211ad 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -330,7 +330,7 @@ New Features * GITHUB#16383: Add fp16 vector encoding support. (Pulkit Gupta) -* GITHUB#15979: Add a de-duplicating HNSW vector format (Lucene106DedupHnswVectorsFormat) that stores +* GITHUB#15979: Add a sandboxed de-duplicating HNSW vector format (DedupHnswVectorsFormat) that stores each distinct full-precision vector once, shared across all documents and fields that reference it. (Kaival Parikh) 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 9d562567d0da..3ac106d11c84 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,4 +16,3 @@ 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/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java similarity index 98% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java index a4ffbbc57608..b4bf395538ea 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatFieldVectorsWriter.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import java.io.IOException; import java.util.ArrayList; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java similarity index 99% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java index ba5e94db1481..7d640b0cc83f 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsFormat.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import java.io.IOException; import org.apache.lucene.codecs.hnsw.FlatVectorsFormat; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java similarity index 91% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java index 14d12d513f6a..7e89e9a0fcaa 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java @@ -14,22 +14,22 @@ * 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.loadDedupFloat16s; -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; +package org.apache.lucene.sandbox.codecs.dedup; + import static org.apache.lucene.index.VectorEncoding.BYTE; import static org.apache.lucene.index.VectorEncoding.FLOAT16; import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_START; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloat16s; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloats; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.readFieldInfo; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.readGroupInfo; import java.io.IOException; import java.util.ArrayList; @@ -39,8 +39,6 @@ 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; @@ -51,6 +49,8 @@ import org.apache.lucene.index.MergePolicy; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.ReadFieldInfo; import org.apache.lucene.store.ChecksumIndexInput; import org.apache.lucene.store.DataAccessHint; import org.apache.lucene.store.FileDataHint; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java similarity index 96% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java index 6885567e37b3..1b142f8cefcc 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsScorer.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java @@ -14,17 +14,17 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; -import static org.apache.lucene.codecs.lucene106.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; import java.io.IOException; import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil; import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.DedupVectorValues; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrd; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.Bits; import org.apache.lucene.util.hnsw.RandomVectorScorer; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java similarity index 89% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java index 3d0ed67345ea..d8020fb2e855 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsWriter.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsWriter.java @@ -14,13 +14,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.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.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.META_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_CODEC_NAME; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; +import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; import java.io.IOException; import org.apache.lucene.codecs.CodecUtil; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java similarity index 90% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java index 68273b9f3c77..0d23c82be046 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlushContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java @@ -14,15 +14,15 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static java.nio.ByteOrder.LITTLE_ENDIAN; -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 static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfFields; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeGroupInfo; import java.io.IOException; import java.nio.ByteBuffer; @@ -35,17 +35,17 @@ 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.FieldOrdToGroupOrd; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdMappedArrayList; -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.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.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdMappedArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.WriteFieldInfo; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.RamUsageEstimator; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java similarity index 98% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java index 4662be8ceb38..4fa9fd473b48 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupGroup.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import java.io.IOException; import java.util.ArrayList; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java similarity index 95% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java index 066bae973c69..bff40f470558 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/Lucene106DedupHnswVectorsFormat.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupHnswVectorsFormat.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; @@ -61,8 +61,8 @@ * * @lucene.experimental */ -public final class Lucene106DedupHnswVectorsFormat extends KnnVectorsFormat { - private static final String NAME = "Lucene106DedupHnswVectorsFormat"; +public final class DedupHnswVectorsFormat extends KnnVectorsFormat { + private static final String NAME = "DedupHnswVectorsFormat"; /** * Controls how many of the nearest neighbor candidates are connected to the new node. Defaults to @@ -99,7 +99,7 @@ public final class Lucene106DedupHnswVectorsFormat extends KnnVectorsFormat { private final int tinySegmentsThreshold; /** Constructs a format using default graph construction parameters */ - public Lucene106DedupHnswVectorsFormat() { + public DedupHnswVectorsFormat() { this( DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD); } @@ -110,7 +110,7 @@ public Lucene106DedupHnswVectorsFormat() { * @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) { + public DedupHnswVectorsFormat(int maxConn, int beamWidth) { this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, HNSW_GRAPH_THRESHOLD); } @@ -122,7 +122,7 @@ public Lucene106DedupHnswVectorsFormat(int maxConn, int beamWidth) { * @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) { + public DedupHnswVectorsFormat(int maxConn, int beamWidth, int tinySegmentsThreshold) { this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null, tinySegmentsThreshold); } @@ -137,7 +137,7 @@ public Lucene106DedupHnswVectorsFormat(int maxConn, int beamWidth, int tinySegme * generated by this format to do the merge. If null, the configured {@link * MergeScheduler#getIntraMergeExecutor(MergePolicy.OneMerge)} is used. */ - public Lucene106DedupHnswVectorsFormat( + public DedupHnswVectorsFormat( int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) { this(maxConn, beamWidth, numMergeWorkers, mergeExec, HNSW_GRAPH_THRESHOLD); } @@ -155,7 +155,7 @@ public Lucene106DedupHnswVectorsFormat( * @param tinySegmentsThreshold the expected number of vector operations to return k nearest * neighbors of the current graph size */ - public Lucene106DedupHnswVectorsFormat( + public DedupHnswVectorsFormat( int maxConn, int beamWidth, int numMergeWorkers, diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java similarity index 93% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java index 4e25900740e9..ceee9770d289 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupMergeContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java @@ -14,15 +14,15 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static java.nio.ByteOrder.LITTLE_ENDIAN; -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 static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfFields; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeGroupInfo; import java.io.IOException; import java.nio.ByteBuffer; @@ -34,12 +34,6 @@ 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.FieldOrdToGroupOrd; -import org.apache.lucene.codecs.lucene106.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; -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.WriteFieldInfo; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocIDMerger; import org.apache.lucene.index.DocsWithFieldSet; @@ -51,6 +45,12 @@ import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.internal.hppc.IntArrayList; import org.apache.lucene.internal.hppc.ObjectCursor; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.WriteFieldInfo; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.IOSupplier; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java similarity index 99% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java index e10bc93e9837..9f1ac6b07fc1 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupUtil.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static org.apache.lucene.index.VectorEncoding.BYTE; import static org.apache.lucene.index.VectorEncoding.FLOAT16; diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java similarity index 80% rename from lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java rename to lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java index 1b732db6a79b..31892094183d 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/package-info.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/package-info.java @@ -20,7 +20,7 @@ * *

      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. + * org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat} for the entry point and {@link + * org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat} for the on-disk layout. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; diff --git a/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat b/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat index 29a44d2ecfa8..4e3e3fc1fca0 100644 --- a/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat +++ b/lucene/sandbox/src/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat @@ -14,3 +14,4 @@ # limitations under the License. org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat +org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java similarity index 96% rename from lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java rename to lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java index 31a5e590064e..5d91ae991240 100644 --- a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestDedupFlatVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java @@ -14,14 +14,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.hamcrest.Matchers.arrayContainingInAnyOrder; import static org.hamcrest.Matchers.instanceOf; import org.apache.lucene.codecs.KnnVectorsReader; -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.KnnFloat16VectorField; @@ -38,20 +37,21 @@ import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; 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. + * Tests that {@link DedupHnswVectorsFormat} stores each distinct vector once. De-duplication is + * observed through the group view size: the number of distinct vectors physically stored, + * regardless of how many documents reference them. */ public class TestDedupFlatVectorsFormat extends LuceneTestCase { private IndexWriterConfig config() { return newIndexWriterConfig() - .setCodec(TestUtil.alwaysKnnVectorsFormat(new Lucene106DedupHnswVectorsFormat())); + .setCodec(TestUtil.alwaysKnnVectorsFormat(new DedupHnswVectorsFormat())); } /** Repeated float vectors within a field are stored once but still read back per document. */ diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java similarity index 94% rename from lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java rename to lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java index 8715e9638593..c8c3143d17d9 100644 --- a/lucene/core/src/test/org/apache/lucene/codecs/lucene106/dedup/TestLucene106DedupHnswVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.lucene.codecs.lucene106.dedup; +package org.apache.lucene.sandbox.codecs.dedup; import static org.hamcrest.Matchers.greaterThan; @@ -30,7 +30,6 @@ import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; -import org.apache.lucene.codecs.simpletext.SimpleTextKnnVectorsReader; import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DocValuesSkipIndexType; import org.apache.lucene.index.DocValuesType; @@ -53,9 +52,9 @@ * 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 { +public class TestDedupHnswVectorsFormat extends BaseKnnVectorsFormatTestCase { - private final KnnVectorsFormat format = new Lucene106DedupHnswVectorsFormat(); + private final KnnVectorsFormat format = new DedupHnswVectorsFormat(); @Override protected Codec getCodec() { @@ -76,10 +75,7 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx 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 (knnVectorsReader instanceof DedupFlatVectorsReader) { if (getNumVectors(knnVectorsReader, fieldInfo) == 0) { assertEquals(0L, totalByteSize); } else { From c20d5016bdf16126fbbab9b3baf47ac0b3048d07 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Wed, 29 Jul 2026 03:10:53 +0000 Subject: [PATCH 09/12] fixes from AI review --- lucene/sandbox/src/java/module-info.java | 4 +- .../codecs/dedup/DedupFlatVectorsScorer.java | 2 +- .../sandbox/codecs/dedup/DedupUtil.java | 89 ++++++++----------- .../dedup/TestDedupHnswVectorsFormat.java | 3 +- 4 files changed, 44 insertions(+), 54 deletions(-) diff --git a/lucene/sandbox/src/java/module-info.java b/lucene/sandbox/src/java/module-info.java index ee9be3227de2..8a3a60d7c544 100644 --- a/lucene/sandbox/src/java/module-info.java +++ b/lucene/sandbox/src/java/module-info.java @@ -22,6 +22,7 @@ requires org.apache.lucene.facet; exports org.apache.lucene.payloads; + exports org.apache.lucene.sandbox.codecs.dedup; exports org.apache.lucene.sandbox.codecs.faiss; exports org.apache.lucene.sandbox.codecs.idversion; exports org.apache.lucene.sandbox.codecs.quantization; @@ -41,5 +42,6 @@ provides org.apache.lucene.codecs.PostingsFormat with org.apache.lucene.sandbox.codecs.idversion.IDVersionPostingsFormat; provides org.apache.lucene.codecs.KnnVectorsFormat with - org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat; + org.apache.lucene.sandbox.codecs.faiss.FaissKnnVectorsFormat, + org.apache.lucene.sandbox.codecs.dedup.DedupHnswVectorsFormat; } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java index 1b142f8cefcc..a85d1f3fbdf7 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java @@ -114,7 +114,7 @@ public UpdateableRandomVectorScorer scorer() throws IOException { @Override public RandomVectorScorerSupplier copy() throws IOException { return new RandomVectorScorerSupplierImpl( - fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd); + fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); } } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java index 9f1ac6b07fc1..ef90916971ce 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java @@ -19,6 +19,8 @@ import static org.apache.lucene.index.VectorEncoding.BYTE; import static org.apache.lucene.index.VectorEncoding.FLOAT16; import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerDense; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerSparse; import static org.apache.lucene.util.StringHelper.GOOD_FAST_HASH_SEED; import static org.apache.lucene.util.StringHelper.murmurhash3_x64_128; @@ -417,24 +419,10 @@ public VectorScorer scorer(byte[] target) throws IOException { return null; } ByteImpl copy = copy(); - DocIndexIterator iterator = copy.iterator(); + DocIndexIterator indexIterator = 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); - } - }; + boolean isDense = copy.fieldView instanceof OffHeapByteVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); } } @@ -551,24 +539,10 @@ public VectorScorer scorer(float[] target) throws IOException { return null; } FloatImpl copy = copy(); - DocIndexIterator iterator = copy.iterator(); + DocIndexIterator indexIterator = 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); - } - }; + boolean isDense = copy.fieldView instanceof OffHeapFloatVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); } } @@ -685,24 +659,37 @@ public VectorScorer scorer(short[] target) throws IOException { return null; } Float16Impl copy = copy(); - DocIndexIterator iterator = copy.iterator(); + DocIndexIterator indexIterator = 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); - } - }; + boolean isDense = + copy.fieldView instanceof OffHeapFloat16VectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + private record DedupVectorScorer( + KnnVectorValues.DocIndexIterator indexIterator, + RandomVectorScorer vectorScorer, + boolean isDense) + implements VectorScorer { + + @Override + public float score() throws IOException { + return vectorScorer.score(indexIterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return indexIterator; + } + + @Override + public Bulk bulk(DocIdSetIterator matchingDocs) { + if (isDense) { + return fromRandomScorerDense(vectorScorer, indexIterator, matchingDocs); + } else { + return fromRandomScorerSparse(vectorScorer, indexIterator, matchingDocs); + } } } } diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java index c8c3143d17d9..6f4c67989ef7 100644 --- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupHnswVectorsFormat.java @@ -30,6 +30,7 @@ import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DocValuesSkipIndexType; import org.apache.lucene.index.DocValuesType; @@ -75,7 +76,7 @@ protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOEx knnVectorsReader = knnVectorsReader.unwrapReaderForField(fieldName); var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo); long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum(); - if (knnVectorsReader instanceof DedupFlatVectorsReader) { + if (knnVectorsReader instanceof Lucene99HnswVectorsReader) { if (getNumVectors(knnVectorsReader, fieldInfo) == 0) { assertEquals(0L, totalByteSize); } else { From 6de567e96e2f993e5bd9ac696db0a75764120fed Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Thu, 30 Jul 2026 19:46:09 +0000 Subject: [PATCH 10/12] address comments, harden tests --- .../lucene99/Lucene99HnswVectorsReader.java | 4 + .../codecs/dedup/DedupFlatVectorsReader.java | 17 +- .../codecs/dedup/DedupFlushContext.java | 64 ++++--- .../sandbox/codecs/dedup/DedupGroup.java | 2 +- .../codecs/dedup/DedupMergeContext.java | 52 +++--- .../sandbox/codecs/dedup/DedupUtil.java | 156 +++++++++--------- .../dedup/TestDedupFlatVectorsFormat.java | 77 ++++++--- 7 files changed, 200 insertions(+), 172 deletions(-) diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java index 59b766948740..f196df3fe3e6 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java @@ -151,6 +151,10 @@ public void finishMerge() throws IOException { flatVectorsReader.finishMerge(); } + public FlatVectorsReader getFlatVectorsReader() { + return flatVectorsReader; + } + private static IndexInput openDataInput( SegmentReadState state, int versionMeta, diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java index 7e89e9a0fcaa..92c610dbb3e9 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java @@ -28,8 +28,6 @@ import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupBytes; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloat16s; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloats; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.readFieldInfo; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.readGroupInfo; import java.io.IOException; import java.util.ArrayList; @@ -111,7 +109,7 @@ final class DedupFlatVectorsReader extends FlatVectorsReader { private void readMetaBody(ChecksumIndexInput meta, FieldInfos fieldInfos) throws IOException { List groupInfos = new ArrayList<>(); while (true) { - GroupInfo groupInfo = readGroupInfo(meta); + GroupInfo groupInfo = GroupInfo.readFromMeta(meta); if (groupInfo == null) { break; } @@ -119,7 +117,7 @@ private void readMetaBody(ChecksumIndexInput meta, FieldInfos fieldInfos) throws } while (true) { - ReadFieldInfo fieldInfo = readFieldInfo(meta); + ReadFieldInfo fieldInfo = ReadFieldInfo.read(meta); if (fieldInfo == null) { break; } @@ -225,7 +223,8 @@ public FlatVectorsScorer getFlatVectorScorer(String field) { return vectorsScorer; } - private FieldEntry getEntry(String field, VectorEncoding expected) { + // package-private for testing + FieldEntry getEntry(String field, VectorEncoding expected) { FieldEntry entry = fields.get(field); if (entry == null) { throw new IllegalArgumentException("field=" + field + " not found"); @@ -267,7 +266,7 @@ private FloatVectorValues getFloatVectorValues(FieldEntry entry) throws IOExcept entry.fieldInfo.function(), entry.fieldInfo.ordToDoc(), entry.fieldInfo.dimension(), - entry.groupInfo.groupSize(), + entry.groupInfo.groupNumVectors(), vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), @@ -286,7 +285,7 @@ private ByteVectorValues getByteVectorValues(FieldEntry entry) throws IOExceptio entry.fieldInfo.function(), entry.fieldInfo.ordToDoc(), entry.fieldInfo.dimension(), - entry.groupInfo.groupSize(), + entry.groupInfo.groupNumVectors(), vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), @@ -305,7 +304,7 @@ private Float16VectorValues getFloat16VectorValues(FieldEntry entry) throws IOEx entry.fieldInfo.function(), entry.fieldInfo.ordToDoc(), entry.fieldInfo.dimension(), - entry.groupInfo.groupSize(), + entry.groupInfo.groupNumVectors(), vectorData, entry.groupInfo.vectorDataOffset(), entry.groupInfo.vectorDataSize(), @@ -351,7 +350,7 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { entry.fieldInfo.fieldOrdToGroupOrdSize() + entry.groupInfo.vectorDataSize()); } - private record FieldEntry(ReadFieldInfo fieldInfo, GroupInfo groupInfo) { + record FieldEntry(ReadFieldInfo fieldInfo, GroupInfo groupInfo) { private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(FieldEntry.class) + RamUsageEstimator.shallowSizeOfInstance(ReadFieldInfo.class) diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java index 0d23c82be046..eae409613178 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java @@ -19,10 +19,8 @@ import static java.nio.ByteOrder.LITTLE_ENDIAN; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfFields; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndMarker; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeGroupInfo; import java.io.IOException; import java.nio.ByteBuffer; @@ -45,7 +43,6 @@ import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdMappedArrayList; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.WriteFieldInfo; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.RamUsageEstimator; @@ -68,10 +65,11 @@ final class DedupFlushContext implements Accountable { } private static DedupGroup getGroup(GroupKey groupKey) { + int dimension = groupKey.dimension(); return switch (groupKey.encoding()) { - case BYTE -> new ByteGroup(groupKey.dimension()); - case FLOAT32 -> new FloatGroup(groupKey.dimension()); - case FLOAT16 -> new Float16Group(groupKey.dimension()); + case BYTE -> new ByteGroup(dimension); + case FLOAT32 -> new FloatGroup(dimension); + case FLOAT16 -> new Float16Group(dimension); }; } @@ -104,30 +102,30 @@ void flush(IndexOutput meta, IndexOutput vectorData, int maxDoc, Sorter.DocMap s int groupOrd = 0; for (Map.Entry> entry : groups.entrySet()) { GroupKey groupKey = entry.getKey(); - DedupGroup group = entry.getValue(); + int dimension = groupKey.dimension(); + VectorEncoding encoding = groupKey.encoding(); - int groupSize = group.size(); - long vectorDataOffset = alignBytes(vectorData, groupKey.encoding()); + DedupGroup group = entry.getValue(); + int groupNumVectors = group.numVectors(); + long vectorDataOffset = alignBytes(vectorData, encoding); // TODO: Write in sorted order for faster merge? (with sequential IO) - for (int ord = 0; ord < groupSize; ord++) { + for (int ord = 0; ord < groupNumVectors; 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); + new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + groupInfo.write(meta); groupOrds.put(groupKey, groupOrd); groupOrd++; } - writeEndOfGroups(meta); + writeEndMarker(meta); for (FieldData fieldData : fieldDataList) { fieldData.fieldWriter.finish(); @@ -149,21 +147,21 @@ void flush(IndexOutput meta, IndexOutput vectorData, int maxDoc, Sorter.DocMap s new FieldOrdToGroupOrdMappedArrayList(new2OldOrd, fieldOrdToGroupOrd); } - WriteFieldInfo fieldInfo = - new WriteFieldInfo( - fieldData.fieldInfo.number, - fieldData.fieldInfo.getVectorSimilarityFunction(), - fieldData.fieldInfo.getVectorDimension(), - fieldData.fieldInfo.getVectorEncoding(), - groupOrds.get(fieldData.groupKey), - vectorCount, - maxDoc, - docs, - fieldOrdToGroupOrdFinal); - writeFieldInfo(meta, vectorData, fieldInfo); + writeFieldInfo( + meta, + vectorData, + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + vectorCount, + maxDoc, + docs, + fieldOrdToGroupOrdFinal); } - writeEndOfFields(meta); + writeEndMarker(meta); } static final class ByteGroup extends DedupGroup { @@ -200,7 +198,7 @@ byte[] serialize(int ord) { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * ramBytesPerVector; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; } } @@ -246,7 +244,7 @@ byte[] serialize(int ord) { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * ramBytesPerVector; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; } } @@ -292,7 +290,7 @@ byte[] serialize(int ord) { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * ramBytesPerVector; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ramBytesPerVector; } } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java index 4fa9fd473b48..e2b83cf67767 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupGroup.java @@ -63,7 +63,7 @@ abstract sealed class DedupGroup implements Accountable abstract byte[] serialize(int ord) throws IOException; - int size() { + int numVectors() { return vectors.size(); } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java index ceee9770d289..cf2510ef24ae 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java @@ -19,10 +19,8 @@ import static java.nio.ByteOrder.LITTLE_ENDIAN; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.alignBytes; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.hashBytes; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfFields; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndOfGroups; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeEndMarker; import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeFieldInfo; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.writeGroupInfo; import java.io.IOException; import java.nio.ByteBuffer; @@ -50,7 +48,6 @@ import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.WriteFieldInfo; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.IOSupplier; @@ -116,35 +113,36 @@ void finish(IndexOutput meta, IndexOutput vectorData) throws IOException { mergeGroup.processField(fieldData, vectorData); } - int groupSize = mergeGroup.size(); + int groupNumVectors = mergeGroup.numVectors(); long vectorDataSize = vectorData.getFilePointer() - vectorDataOffset; GroupInfo groupInfo = - new GroupInfo(groupOrd, dimension, encoding, groupSize, vectorDataOffset, vectorDataSize); - writeGroupInfo(meta, groupInfo); + new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + groupInfo.write(meta); groupOrds.put(groupKey, groupOrd); groupOrd++; } - writeEndOfGroups(meta); + writeEndMarker(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.fieldOrdToGroupOrd.elementsCount, - fieldData.maxDoc, - fieldData.docsWithFieldSet, - new FieldOrdToGroupOrdArrayList(fieldData.fieldOrdToGroupOrd)); - writeFieldInfo(meta, vectorData, fieldInfo); + writeFieldInfo( + meta, + vectorData, + fieldData.fieldInfo.number, + fieldData.fieldInfo.getVectorSimilarityFunction(), + fieldData.fieldInfo.getVectorDimension(), + fieldData.fieldInfo.getVectorEncoding(), + groupOrds.get(fieldData.groupKey), + fieldData.fieldOrdToGroupOrd.elementsCount, + fieldData.maxDoc, + fieldData.docsWithFieldSet, + new FieldOrdToGroupOrdArrayList(fieldData.fieldOrdToGroupOrd)); } - writeEndOfFields(meta); + writeEndMarker(meta); } abstract static sealed class DedupMergeGroup extends DedupGroup { @@ -157,13 +155,13 @@ void processField(FieldData fieldData, IndexOutput vectorData) throws IOExceptio // iterate merged docs one-by-one for (Sub next = merger.next(); next != null; next = merger.next()) { T vector = vectorFrom(next); - int groupSize = size(); + int groupNumVectors = numVectors(); // add vector to group ObjectCursor cursor = addUnique(vector); - if (cursor.index == groupSize) { // new addition + if (cursor.index == groupNumVectors) { // new addition // already on-heap, write immediately to avoid another IO read - byte[] bytes = serialize(groupSize); + byte[] bytes = serialize(groupNumVectors); vectorData.writeBytes(bytes, bytes.length); } @@ -225,7 +223,7 @@ byte[] serialize(int ord) throws IOException { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * ByteVector.SHALLOW_SIZE; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * ByteVector.SHALLOW_SIZE; } } @@ -291,7 +289,7 @@ byte[] serialize(int ord) throws IOException { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * FloatVector.SHALLOW_SIZE; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * FloatVector.SHALLOW_SIZE; } } @@ -358,7 +356,7 @@ byte[] serialize(int ord) throws IOException { @Override public long ramBytesUsed() { - return SHALLOW_SIZE + super.ramBytesUsed() + size() * Float16Vector.SHALLOW_SIZE; + return SHALLOW_SIZE + super.ramBytesUsed() + numVectors() * Float16Vector.SHALLOW_SIZE; } } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java index ef90916971ce..3a8ccdbdcc24 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java @@ -110,40 +110,39 @@ record GroupInfo( int groupOrd, int dimension, VectorEncoding encoding, - int groupSize, + int groupNumVectors, 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); - } + long vectorDataSize) { - static GroupInfo readGroupInfo(IndexInput meta) throws IOException { - int groupOrd = meta.readInt(); - if (groupOrd == END_MARKER) { - return null; + void write(IndexOutput meta) throws IOException { + meta.writeInt(groupOrd); + meta.writeInt(dimension); + meta.writeInt(encoding.ordinal()); + meta.writeInt(groupNumVectors); + meta.writeLong(vectorDataOffset); + meta.writeLong(vectorDataSize); } - int dimension = meta.readInt(); - VectorEncoding encoding = VectorEncoding.values()[meta.readInt()]; - int groupSize = meta.readInt(); - long vectorDataOffset = meta.readLong(); - long vectorDataSize = meta.readLong(); + static GroupInfo readFromMeta(IndexInput meta) throws IOException { + int groupOrd = meta.readInt(); + if (groupOrd == END_MARKER) { + return null; + } - return new GroupInfo( - groupOrd, dimension, encoding, groupSize, vectorDataOffset, vectorDataSize); + int dimension = meta.readInt(); + VectorEncoding encoding = VectorEncoding.values()[meta.readInt()]; + int groupNumVectors = meta.readInt(); + long vectorDataOffset = meta.readLong(); + long vectorDataSize = meta.readLong(); + + return new GroupInfo( + groupOrd, dimension, encoding, groupNumVectors, vectorDataOffset, vectorDataSize); + } } - record WriteFieldInfo( + static void writeFieldInfo( + IndexOutput meta, + IndexOutput vectorData, int fieldNumber, VectorSimilarityFunction function, int dimension, @@ -152,34 +151,26 @@ record WriteFieldInfo( int vectorCount, int maxDoc, DocsWithFieldSet docs, - FieldOrdToGroupOrd fieldOrdToGroupOrd) {} - - static void writeFieldInfo(IndexOutput meta, IndexOutput vectorData, WriteFieldInfo fieldInfo) + FieldOrdToGroupOrd fieldOrdToGroupOrd) 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); + meta.writeInt(fieldNumber); + meta.writeInt(function.ordinal()); + meta.writeInt(dimension); + meta.writeInt(encoding.ordinal()); + meta.writeInt(groupOrd); + meta.writeInt(vectorCount); // write ordToDoc OrdToDocDISIReaderConfiguration.writeStoredMeta( - ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT, - meta, - vectorData, - fieldInfo.vectorCount, - fieldInfo.maxDoc, - fieldInfo.docs); + ORD_TO_DOC_DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, vectorCount, maxDoc, docs); // write fieldOrdToGroupOrd long fieldOrdToGroupOrdOffset = vectorData.alignFilePointer(FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES); DirectWriter writer = - DirectWriter.getInstance( - vectorData, fieldInfo.vectorCount, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); - for (int i = 0; i < fieldInfo.vectorCount; i++) { - writer.add(fieldInfo.fieldOrdToGroupOrd.get(i)); + DirectWriter.getInstance(vectorData, vectorCount, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); + for (int i = 0; i < vectorCount; i++) { + writer.add(fieldOrdToGroupOrd.get(i)); } writer.finish(); long fieldOrdToGroupOrdSize = vectorData.getFilePointer() - fieldOrdToGroupOrdOffset; @@ -188,7 +179,7 @@ static void writeFieldInfo(IndexOutput meta, IndexOutput vectorData, WriteFieldI meta.writeLong(fieldOrdToGroupOrdSize); } - static void writeEndOfFields(IndexOutput meta) throws IOException { + static void writeEndMarker(IndexOutput meta) throws IOException { meta.writeInt(END_MARKER); } @@ -201,35 +192,36 @@ record ReadFieldInfo( int vectorCount, OrdToDocDISIReaderConfiguration ordToDoc, long fieldOrdToGroupOrdOffset, - long fieldOrdToGroupOrdSize) {} - - 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 fieldOrdToGroupOrdOffset = meta.readLong(); - long fieldOrdToGroupOrdSize = meta.readLong(); - - return new ReadFieldInfo( - fieldNumber, - function, - dimension, - encoding, - groupOrd, - vectorCount, - ordToDoc, - fieldOrdToGroupOrdOffset, - fieldOrdToGroupOrdSize); + long fieldOrdToGroupOrdSize) { + + static ReadFieldInfo read(IndexInput meta) throws IOException { + + int fieldNumber = meta.readInt(); + if (fieldNumber == END_MARKER) { + return null; + } + + VectorSimilarityFunction function = VectorSimilarityFunction.values()[meta.readInt()]; + int dimension = meta.readInt(); + VectorEncoding encoding = VectorEncoding.values()[meta.readInt()]; + int groupOrd = meta.readInt(); + int vectorCount = meta.readInt(); + OrdToDocDISIReaderConfiguration ordToDoc = + OrdToDocDISIReaderConfiguration.fromStoredMeta(meta, vectorCount); + long fieldOrdToGroupOrdOffset = meta.readLong(); + long fieldOrdToGroupOrdSize = meta.readLong(); + + return new ReadFieldInfo( + fieldNumber, + function, + dimension, + encoding, + groupOrd, + vectorCount, + ordToDoc, + fieldOrdToGroupOrdOffset, + fieldOrdToGroupOrdSize); + } } static long hashBytes(byte[] bytes) { @@ -311,7 +303,7 @@ static ByteVectorValues loadDedupBytes( VectorSimilarityFunction function, OrdToDocDISIReaderConfiguration configuration, int dimension, - int groupSize, + int groupNumVectors, IndexInput vectorData, long vectorDataOffset, long vectorDataSize, @@ -326,7 +318,7 @@ static ByteVectorValues loadDedupBytes( final OffHeapByteVectorValues groupView = new OffHeapByteVectorValues.DenseOffHeapVectorValues( dimension, - groupSize, + groupNumVectors, vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), fieldView.getVectorByteLength(), vectorsScorer, @@ -431,7 +423,7 @@ static FloatVectorValues loadDedupFloats( VectorSimilarityFunction function, OrdToDocDISIReaderConfiguration configuration, int dimension, - int groupSize, + int groupNumVectors, IndexInput vectorData, long vectorDataOffset, long vectorDataSize, @@ -446,7 +438,7 @@ static FloatVectorValues loadDedupFloats( final OffHeapFloatVectorValues groupView = new OffHeapFloatVectorValues.DenseOffHeapVectorValues( dimension, - groupSize, + groupNumVectors, vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), fieldView.getVectorByteLength(), vectorsScorer, @@ -551,7 +543,7 @@ static Float16VectorValues loadDedupFloat16s( VectorSimilarityFunction function, OrdToDocDISIReaderConfiguration configuration, int dimension, - int groupSize, + int groupNumVectors, IndexInput vectorData, long vectorDataOffset, long vectorDataSize, @@ -566,7 +558,7 @@ static Float16VectorValues loadDedupFloat16s( final OffHeapFloat16VectorValues groupView = new OffHeapFloat16VectorValues.DenseOffHeapVectorValues( dimension, - groupSize, + groupNumVectors, vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), fieldView.getVectorByteLength(), vectorsScorer, diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java index 5d91ae991240..ae8c76445788 100644 --- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java @@ -16,11 +16,14 @@ */ package org.apache.lucene.sandbox.codecs.dedup; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.hamcrest.Matchers.arrayContainingInAnyOrder; import static org.hamcrest.Matchers.instanceOf; import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.FlatVectorsReader; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; import org.apache.lucene.document.Document; import org.apache.lucene.document.KnnByteVectorField; import org.apache.lucene.document.KnnFloat16VectorField; @@ -49,7 +52,7 @@ */ public class TestDedupFlatVectorsFormat extends LuceneTestCase { - private IndexWriterConfig config() { + private static IndexWriterConfig config() { return newIndexWriterConfig() .setCodec(TestUtil.alwaysKnnVectorsFormat(new DedupHnswVectorsFormat())); } @@ -72,7 +75,7 @@ public void testFloatDuplicatesWithinField() throws Exception { LeafReader leafReader = getOnlyLeafReader(reader); FloatVectorValues values = leafReader.getFloatVectorValues("f"); assertEquals(docVectors.length, values.size()); // one entry per document - assertEquals(2, groupSize(values)); // only two distinct vectors stored + assertEquals(2, groupNumVectors(values)); // only two distinct vectors stored NumericDocValues docValues = leafReader.getNumericDocValues("id"); Integer[] expectedOrds = new Integer[docVectors.length]; Integer[] ordsSeen = new Integer[docVectors.length]; @@ -107,7 +110,7 @@ public void testFloat16DuplicatesWithinField() throws Exception { LeafReader leafReader = getOnlyLeafReader(reader); Float16VectorValues values = leafReader.getFloat16VectorValues("f"); assertEquals(docVectors.length, values.size()); // one entry per document - assertEquals(2, groupSize(values)); // only two distinct vectors stored + assertEquals(2, groupNumVectors(values)); // only two distinct vectors stored NumericDocValues docValues = leafReader.getNumericDocValues("id"); Integer[] expectedOrds = new Integer[docVectors.length]; Integer[] ordsSeen = new Integer[docVectors.length]; @@ -128,7 +131,7 @@ public void testFloat16DuplicatesWithinField() throws Exception { public void testByteDuplicatesWithinField() throws Exception { byte[] a = {1, 2, 3, 4}; byte[] b = {5, 6, 7, 8}; - byte[][] docVectors = {a, a, b, a, b}; + byte[][] docVectors = {a, b, a, b, a, b}; // 3 copies each of 2 vectors try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, config())) { for (int ord = 0; ord < docVectors.length; ord++) { @@ -142,7 +145,7 @@ public void testByteDuplicatesWithinField() throws Exception { LeafReader leafReader = getOnlyLeafReader(reader); ByteVectorValues values = leafReader.getByteVectorValues("f"); assertEquals(docVectors.length, values.size()); - assertEquals(2, groupSize(values)); + assertEquals(2, groupNumVectors(values)); NumericDocValues docValues = leafReader.getNumericDocValues("id"); Integer[] expectedOrds = new Integer[docVectors.length]; Integer[] ordsSeen = new Integer[docVectors.length]; @@ -176,7 +179,7 @@ public void testDistinctVectorsAllStored() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { FloatVectorValues values = getOnlyLeafReader(reader).getFloatVectorValues("f"); assertEquals(distinctDocVectors.length, values.size()); - assertEquals(distinctDocVectors.length, groupSize(values)); + assertEquals(distinctDocVectors.length, groupNumVectors(values)); } } } @@ -196,11 +199,8 @@ public void testOffHeapSize() throws Exception { w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader leafReader = getOnlyLeafReader(reader); - assertThat(leafReader, instanceOf(CodecReader.class)); - + DedupFlatVectorsReader dedupReader = getDedupReader(leafReader, "f"); FieldInfo fieldInfo = leafReader.getFieldInfos().fieldInfo("f"); - KnnVectorsReader knnVectorsReader = ((CodecReader) leafReader).getVectorReader(); - knnVectorsReader = knnVectorsReader.unwrapReaderForField("f"); long expectedOffHeapSize = (docVectors.length * Integer.BYTES) // fieldOrdToGroupOrd mapping @@ -208,7 +208,7 @@ public void testOffHeapSize() throws Exception { assertEquals( expectedOffHeapSize, - knnVectorsReader + dedupReader .getOffHeapByteSize(fieldInfo) .get("vdd") // vector data extension .longValue()); @@ -228,11 +228,21 @@ public void testDuplicatesAcrossFieldsShareGroup() throws Exception { w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader leaf = getOnlyLeafReader(reader); + + DedupFlatVectorsReader dedupReader1 = getDedupReader(leaf, "f1"); + DedupFlatVectorsReader dedupReader2 = getDedupReader(leaf, "f2"); + assertEquals(dedupReader1, dedupReader2); // de-duplication happened correctly + + assertEquals( // both fields DO resolve to the same group + dedupReader1.getEntry("f1", FLOAT32).groupInfo(), + dedupReader2.getEntry("f2", FLOAT32).groupInfo()); + FloatVectorValues v1 = leaf.getFloatVectorValues("f1"); - FloatVectorValues v2 = leaf.getFloatVectorValues("f2"); - assertEquals(1, groupSize(v1)); // both fields resolve to the same one-vector group - assertEquals(1, groupSize(v2)); + assertEquals(1, groupNumVectors(v1)); // the group has one vector assertArrayEquals(shared, v1.vectorValue(0), 0f); + + FloatVectorValues v2 = leaf.getFloatVectorValues("f2"); + assertEquals(1, groupNumVectors(v2)); // the group has one vector assertArrayEquals(shared, v2.vectorValue(0), 0f); } } @@ -251,10 +261,22 @@ public void testDifferentDimensionsUseSeparateGroups() throws Exception { 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(vector1, leaf.getFloatVectorValues("f2d").vectorValue(0), 0f); - assertArrayEquals(vector2, leaf.getFloatVectorValues("f3d").vectorValue(0), 0f); + + DedupFlatVectorsReader dedupReader1 = getDedupReader(leaf, "f2d"); + DedupFlatVectorsReader dedupReader2 = getDedupReader(leaf, "f3d"); + assertEquals(dedupReader1, dedupReader2); // de-duplication happened correctly + + assertNotEquals( // both fields DO NOT resolve to the same group + dedupReader1.getEntry("f2d", FLOAT32).groupInfo(), + dedupReader2.getEntry("f3d", FLOAT32).groupInfo()); + + FloatVectorValues v1 = leaf.getFloatVectorValues("f2d"); + assertEquals(1, groupNumVectors(v1)); // the group has one vector + assertArrayEquals(vector1, v1.vectorValue(0), 0f); + + FloatVectorValues v2 = leaf.getFloatVectorValues("f3d"); + assertEquals(1, groupNumVectors(v2)); // the group has one vector + assertArrayEquals(vector2, v2.vectorValue(0), 0f); } } } @@ -278,7 +300,7 @@ public void testDuplicatesAcrossSegmentsDedupOnMerge() throws Exception { LeafReader leafReader = getOnlyLeafReader(reader); FloatVectorValues values = leafReader.getFloatVectorValues("f"); assertEquals(docVectors.length, values.size()); - assertEquals(2, groupSize(values)); // a's duplicate collapsed across segments + assertEquals(2, groupNumVectors(values)); // a's duplicate collapsed across segments NumericDocValues docValues = leafReader.getNumericDocValues("id"); Integer[] expectedOrds = new Integer[docVectors.length]; Integer[] ordsSeen = new Integer[docVectors.length]; @@ -296,7 +318,22 @@ public void testDuplicatesAcrossSegmentsDedupOnMerge() throws Exception { } /** Number of distinct vectors physically stored for a field's group. */ - private static int groupSize(KnnVectorValues values) { + private static int groupNumVectors(KnnVectorValues values) { + assertThat(values, instanceOf(DedupVectorValues.class)); return ((DedupVectorValues) values).getGroupView().size(); } + + /** Get underlying dedup vector reader instance. */ + private static DedupFlatVectorsReader getDedupReader(LeafReader leafReader, String fieldName) { + assertThat(leafReader, instanceOf(CodecReader.class)); + KnnVectorsReader knnVectorsReader = ((CodecReader) leafReader).getVectorReader(); + knnVectorsReader = knnVectorsReader.unwrapReaderForField(fieldName); + + assertThat(knnVectorsReader, instanceOf(Lucene99HnswVectorsReader.class)); + FlatVectorsReader flatReader = + ((Lucene99HnswVectorsReader) knnVectorsReader).getFlatVectorsReader(); + + assertThat(flatReader, instanceOf(DedupFlatVectorsReader.class)); + return (DedupFlatVectorsReader) flatReader; + } } From 784da4faad1ff7363f5adb66091eee93c7f7cbc0 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Thu, 30 Jul 2026 19:53:37 +0000 Subject: [PATCH 11/12] iter --- .../sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java index ae8c76445788..0123d7153b3c 100644 --- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java @@ -17,6 +17,7 @@ package org.apache.lucene.sandbox.codecs.dedup; import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.index.VectorSimilarityFunction.DOT_PRODUCT; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.hamcrest.Matchers.arrayContainingInAnyOrder; import static org.hamcrest.Matchers.instanceOf; @@ -223,7 +224,7 @@ public void testDuplicatesAcrossFieldsShareGroup() throws Exception { IndexWriter w = new IndexWriter(dir, config())) { Document doc = new Document(); doc.add(new KnnFloatVectorField("f1", shared, EUCLIDEAN)); - doc.add(new KnnFloatVectorField("f2", shared, EUCLIDEAN)); + doc.add(new KnnFloatVectorField("f2", shared, DOT_PRODUCT)); // different function w.addDocument(doc); w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { From a67cf1280841a8f2e5f91d3e8c239947ccfc7403 Mon Sep 17 00:00:00 2001 From: Kaival Parikh Date: Fri, 31 Jul 2026 19:50:17 +0000 Subject: [PATCH 12/12] address comments --- .../codecs/dedup/DedupFlatVectorsReader.java | 10 +- .../codecs/dedup/DedupFlatVectorsScorer.java | 3 +- .../codecs/dedup/DedupFlushContext.java | 6 +- .../codecs/dedup/DedupMergeContext.java | 5 +- .../sandbox/codecs/dedup/DedupUtil.java | 500 +---------------- .../codecs/dedup/DedupVectorValues.java | 519 ++++++++++++++++++ .../dedup/TestDedupFlatVectorsFormat.java | 102 +++- 7 files changed, 635 insertions(+), 510 deletions(-) create mode 100644 lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java index 92c610dbb3e9..dfe5454cd0da 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsReader.java @@ -25,9 +25,9 @@ import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VECTOR_DATA_EXTENSION; import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_CURRENT; import static org.apache.lucene.sandbox.codecs.dedup.DedupFlatVectorsFormat.VERSION_START; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupBytes; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloat16s; -import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.loadDedupFloats; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupBytes; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupFloat16s; +import static org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.loadDedupFloats; import java.io.IOException; import java.util.ArrayList; @@ -319,7 +319,9 @@ public Float16VectorValues getFloat16VectorValues(String field) throws IOExcepti @Override public FlatVectorsReader getMergeInstance() { - // TODO: Can we improve performance using strictly sequential IO? + // TODO: Can we improve performance using sequential IO + avoiding read-backs? One way is to + // de-duplicate only within a document, allowing for cleaner sort of vectors during flush + + // avoid read-backs for full equality checks during merge. return this; } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java index a85d1f3fbdf7..63da8b2d950b 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlatVectorsScorer.java @@ -23,8 +23,7 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.Bits; import org.apache.lucene.util.hnsw.RandomVectorScorer; diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java index eae409613178..59ad5465b83e 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupFlushContext.java @@ -38,11 +38,11 @@ import org.apache.lucene.index.Sorter; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.internal.hppc.IntArrayList; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdMappedArrayList; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdArrayList; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdMappedArrayList; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.RamUsageEstimator; diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java index cf2510ef24ae..55ca0e2cb1ac 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupMergeContext.java @@ -43,11 +43,10 @@ import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.internal.hppc.IntArrayList; import org.apache.lucene.internal.hppc.ObjectCursor; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrd; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FieldOrdToGroupOrdArrayList; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupInfo; import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.GroupKey; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; +import org.apache.lucene.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrdArrayList; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.IOSupplier; diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java index 3a8ccdbdcc24..594fb6abe730 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupUtil.java @@ -16,38 +16,18 @@ */ package org.apache.lucene.sandbox.codecs.dedup; -import static org.apache.lucene.index.VectorEncoding.BYTE; -import static org.apache.lucene.index.VectorEncoding.FLOAT16; -import static org.apache.lucene.index.VectorEncoding.FLOAT32; -import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerDense; -import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerSparse; 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.OffHeapFloat16VectorValues; -import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues; import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; -import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.Float16VectorValues; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.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.sandbox.codecs.dedup.DedupVectorValues.FieldOrdToGroupOrd; 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; /** @@ -62,12 +42,13 @@ final class DedupUtil { private static final int END_MARKER = -1; + /** Alignment bytes on disk for fieldOrdToGroupOrd. */ private static final int FIELD_ORD_TO_GROUP_ORD_ALIGN_BYTES = 4; // TODO: This is the number of bits used to write each group ordinal in the index-backed per-field // FieldOrdToGroupOrd mapping. Evaluate using fewer bits to reduce index size, at the expense of // costlier lookups. - private static final int FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE = 32; + static final int FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE = 32; /** Initial allocation size for internal re-used int[] scratch buffers. */ static final int SCRATCH_INITIAL_SIZE = 16; @@ -79,33 +60,6 @@ record GroupKey(int dimension, VectorEncoding encoding) { } } - /** - * Vector values that share a single copy of each distinct vector across the documents and fields - * that reference it. - * - *

      Every instance is backed by two views: the {@code fieldView} maps ordinals to docs and - * drives iteration (one entry per document), while the {@code groupView} holds the de-duplicated - * vectors (one entry per distinct vector). {@code fieldOrdToGroupOrd} translates a document - * ordinal in the field into its group ordinal. - */ - sealed interface DedupVectorValues { - /** The dense view over distinct vectors, indexed by group ordinal. */ - KnnVectorValues getGroupView(); - - /** Maps a per-document ordinal to its group ordinal in {@link #getGroupView()}. */ - FieldOrdToGroupOrd getFieldOrdToGroupOrd(); - } - - /** - * Maps a field's per-document ordinal to the ordinal of its (shared) vector within the group. - * Backed on-heap while writing and off-heap while reading. - */ - sealed interface FieldOrdToGroupOrd { - int get(int ord); - - FieldOrdToGroupOrd copy() throws IOException; - } - record GroupInfo( int groupOrd, int dimension, @@ -236,452 +190,4 @@ static long alignBytes(IndexOutput output, VectorEncoding encoding) throws IOExc }; return output.alignFilePointer(alignBytes); } - - /** On-heap map used during a flush, backed directly by the buffered ordinals. */ - record FieldOrdToGroupOrdArrayList(IntArrayList fieldOrdToGroupOrd) - implements FieldOrdToGroupOrd { - - @Override - public int get(int ord) { - return fieldOrdToGroupOrd.get(ord); - } - - @Override - public FieldOrdToGroupOrd copy() { - return new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); - } - } - - /** On-heap map used during a sorted flush, indirecting through a new-to-old ordinal map. */ - record FieldOrdToGroupOrdMappedArrayList(int[] map, IntArrayList fieldOrdToGroupOrd) - implements FieldOrdToGroupOrd { - - @Override - public int get(int ord) { - return fieldOrdToGroupOrd.get(map[ord]); - } - - @Override - public FieldOrdToGroupOrd copy() { - return new FieldOrdToGroupOrdMappedArrayList(map, fieldOrdToGroupOrd); - } - } - - /** Off-heap map used while reading, backed by a {@link DirectReader}. */ - static final class FieldOrdToGroupOrdOffHeap implements FieldOrdToGroupOrd { - private final IndexInput vectorData; - private final long fieldOrdToGroupOrdOffset; - private final long fieldOrdToGroupOrdSize; - private final LongValues values; - - FieldOrdToGroupOrdOffHeap( - IndexInput vectorData, long fieldOrdToGroupOrdOffset, long fieldOrdToGroupOrdSize) - throws IOException { - this.vectorData = vectorData; - this.fieldOrdToGroupOrdOffset = fieldOrdToGroupOrdOffset; - this.fieldOrdToGroupOrdSize = fieldOrdToGroupOrdSize; - - RandomAccessInput slice = - vectorData.randomAccessSlice(fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - this.values = DirectReader.getInstance(slice, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); - } - - @Override - public int get(int v) { - return (int) values.get(v); - } - - @Override - public FieldOrdToGroupOrd copy() throws IOException { - return new FieldOrdToGroupOrdOffHeap( - vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - } - } - - static ByteVectorValues loadDedupBytes( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - OrdToDocDISIReaderConfiguration configuration, - int dimension, - int groupNumVectors, - IndexInput vectorData, - long vectorDataOffset, - long vectorDataSize, - long fieldOrdToGroupOrdOffset, - long fieldOrdToGroupOrdSize) - throws IOException { - - final OffHeapByteVectorValues fieldView = - OffHeapByteVectorValues.load( - function, vectorsScorer, configuration, BYTE, dimension, 0, 0, vectorData); - - final OffHeapByteVectorValues groupView = - new OffHeapByteVectorValues.DenseOffHeapVectorValues( - dimension, - groupNumVectors, - vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), - fieldView.getVectorByteLength(), - vectorsScorer, - function); - - final FieldOrdToGroupOrd fieldOrdToGroupOrd = - new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - - return new ByteImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); - } - - /** {@link DedupVectorValues} over byte vectors. */ - private static final class ByteImpl extends ByteVectorValues implements DedupVectorValues { - private final FlatVectorsScorer vectorsScorer; - private final VectorSimilarityFunction function; - private final ByteVectorValues fieldView; - private final ByteVectorValues groupView; - private final FieldOrdToGroupOrd fieldOrdToGroupOrd; - private int[] scratch; - - ByteImpl( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - ByteVectorValues fieldView, - ByteVectorValues groupView, - FieldOrdToGroupOrd fieldOrdToGroupOrd) { - this.vectorsScorer = vectorsScorer; - this.function = function; - this.fieldView = fieldView; - this.groupView = groupView; - this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; - this.scratch = new int[SCRATCH_INITIAL_SIZE]; - } - - @Override - public ByteVectorValues getGroupView() { - return groupView; - } - - @Override - public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { - return fieldOrdToGroupOrd; - } - - @Override - public int ordToDoc(int ord) { - return fieldView.ordToDoc(ord); - } - - @Override - public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { - if (scratch.length < ordsToPrefetch.length) { // grow if needed - scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); - } - for (int i = 0; i < numOrds; i++) { - scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); - } - groupView.prefetch(scratch, numOrds); - } - - @Override - public byte[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); - } - - @Override - public int dimension() { - return fieldView.dimension(); - } - - @Override - public int size() { - return fieldView.size(); - } - - @Override - public ByteImpl copy() throws IOException { - return new ByteImpl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); - } - - @Override - public DocIndexIterator iterator() { - return fieldView.iterator(); - } - - @Override - public VectorScorer scorer(byte[] target) throws IOException { - if (size() == 0) { - return null; - } - ByteImpl copy = copy(); - DocIndexIterator indexIterator = copy.iterator(); - RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); - boolean isDense = copy.fieldView instanceof OffHeapByteVectorValues.DenseOffHeapVectorValues; - return new DedupVectorScorer(indexIterator, vectorScorer, isDense); - } - } - - static FloatVectorValues loadDedupFloats( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - OrdToDocDISIReaderConfiguration configuration, - int dimension, - int groupNumVectors, - IndexInput vectorData, - long vectorDataOffset, - long vectorDataSize, - long fieldOrdToGroupOrdOffset, - long fieldOrdToGroupOrdSize) - throws IOException { - - final OffHeapFloatVectorValues fieldView = - OffHeapFloatVectorValues.load( - function, vectorsScorer, configuration, FLOAT32, dimension, 0, 0, vectorData); - - final OffHeapFloatVectorValues groupView = - new OffHeapFloatVectorValues.DenseOffHeapVectorValues( - dimension, - groupNumVectors, - vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), - fieldView.getVectorByteLength(), - vectorsScorer, - function); - - final FieldOrdToGroupOrd fieldOrdToGroupOrd = - new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - - return new FloatImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); - } - - /** {@link DedupVectorValues} over float vectors. */ - private static final class FloatImpl extends FloatVectorValues implements DedupVectorValues { - private final FlatVectorsScorer vectorsScorer; - private final VectorSimilarityFunction function; - private final FloatVectorValues fieldView; - private final FloatVectorValues groupView; - private final FieldOrdToGroupOrd fieldOrdToGroupOrd; - private int[] scratch; - - FloatImpl( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - FloatVectorValues fieldView, - FloatVectorValues groupView, - FieldOrdToGroupOrd fieldOrdToGroupOrd) { - this.vectorsScorer = vectorsScorer; - this.function = function; - this.fieldView = fieldView; - this.groupView = groupView; - this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; - this.scratch = new int[SCRATCH_INITIAL_SIZE]; - } - - @Override - public FloatVectorValues getGroupView() { - return groupView; - } - - @Override - public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { - return fieldOrdToGroupOrd; - } - - @Override - public int ordToDoc(int ord) { - return fieldView.ordToDoc(ord); - } - - @Override - public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { - if (scratch.length < ordsToPrefetch.length) { // grow if needed - scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); - } - for (int i = 0; i < numOrds; i++) { - scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); - } - groupView.prefetch(scratch, numOrds); - } - - @Override - public float[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); - } - - @Override - public int dimension() { - return fieldView.dimension(); - } - - @Override - public int size() { - return fieldView.size(); - } - - @Override - public FloatImpl copy() throws IOException { - return new FloatImpl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); - } - - @Override - public DocIndexIterator iterator() { - return fieldView.iterator(); - } - - @Override - public VectorScorer scorer(float[] target) throws IOException { - if (size() == 0) { - return null; - } - FloatImpl copy = copy(); - DocIndexIterator indexIterator = copy.iterator(); - RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); - boolean isDense = copy.fieldView instanceof OffHeapFloatVectorValues.DenseOffHeapVectorValues; - return new DedupVectorScorer(indexIterator, vectorScorer, isDense); - } - } - - static Float16VectorValues loadDedupFloat16s( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - OrdToDocDISIReaderConfiguration configuration, - int dimension, - int groupNumVectors, - IndexInput vectorData, - long vectorDataOffset, - long vectorDataSize, - long fieldOrdToGroupOrdOffset, - long fieldOrdToGroupOrdSize) - throws IOException { - - final OffHeapFloat16VectorValues fieldView = - OffHeapFloat16VectorValues.load( - function, vectorsScorer, configuration, FLOAT16, dimension, 0, 0, vectorData); - - final OffHeapFloat16VectorValues groupView = - new OffHeapFloat16VectorValues.DenseOffHeapVectorValues( - dimension, - groupNumVectors, - vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), - fieldView.getVectorByteLength(), - vectorsScorer, - function); - - final FieldOrdToGroupOrd fieldOrdToGroupOrd = - new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); - - return new Float16Impl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); - } - - /** {@link DedupVectorValues} over float16 vectors. */ - private static final class Float16Impl extends Float16VectorValues implements DedupVectorValues { - private final FlatVectorsScorer vectorsScorer; - private final VectorSimilarityFunction function; - private final Float16VectorValues fieldView; - private final Float16VectorValues groupView; - private final FieldOrdToGroupOrd fieldOrdToGroupOrd; - private int[] scratch; - - Float16Impl( - FlatVectorsScorer vectorsScorer, - VectorSimilarityFunction function, - Float16VectorValues fieldView, - Float16VectorValues groupView, - FieldOrdToGroupOrd fieldOrdToGroupOrd) { - this.vectorsScorer = vectorsScorer; - this.function = function; - this.fieldView = fieldView; - this.groupView = groupView; - this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; - this.scratch = new int[SCRATCH_INITIAL_SIZE]; - } - - @Override - public Float16VectorValues getGroupView() { - return groupView; - } - - @Override - public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { - return fieldOrdToGroupOrd; - } - - @Override - public int ordToDoc(int ord) { - return fieldView.ordToDoc(ord); - } - - @Override - public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { - if (scratch.length < ordsToPrefetch.length) { // grow if needed - scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); - } - for (int i = 0; i < numOrds; i++) { - scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); - } - groupView.prefetch(scratch, numOrds); - } - - @Override - public short[] vectorValue(int ord) throws IOException { - return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); - } - - @Override - public int dimension() { - return fieldView.dimension(); - } - - @Override - public int size() { - return fieldView.size(); - } - - @Override - public Float16Impl copy() throws IOException { - return new Float16Impl( - vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); - } - - @Override - public DocIndexIterator iterator() { - return fieldView.iterator(); - } - - @Override - public VectorScorer scorer(short[] target) throws IOException { - if (size() == 0) { - return null; - } - Float16Impl copy = copy(); - DocIndexIterator indexIterator = copy.iterator(); - RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); - boolean isDense = - copy.fieldView instanceof OffHeapFloat16VectorValues.DenseOffHeapVectorValues; - return new DedupVectorScorer(indexIterator, vectorScorer, isDense); - } - } - - private record DedupVectorScorer( - KnnVectorValues.DocIndexIterator indexIterator, - RandomVectorScorer vectorScorer, - boolean isDense) - implements VectorScorer { - - @Override - public float score() throws IOException { - return vectorScorer.score(indexIterator.index()); - } - - @Override - public DocIdSetIterator iterator() { - return indexIterator; - } - - @Override - public Bulk bulk(DocIdSetIterator matchingDocs) { - if (isDense) { - return fromRandomScorerDense(vectorScorer, indexIterator, matchingDocs); - } else { - return fromRandomScorerSparse(vectorScorer, indexIterator, matchingDocs); - } - } - } } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java new file mode 100644 index 000000000000..802f15d72404 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/codecs/dedup/DedupVectorValues.java @@ -0,0 +1,519 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.sandbox.codecs.dedup; + +import static org.apache.lucene.index.VectorEncoding.BYTE; +import static org.apache.lucene.index.VectorEncoding.FLOAT16; +import static org.apache.lucene.index.VectorEncoding.FLOAT32; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE; +import static org.apache.lucene.sandbox.codecs.dedup.DedupUtil.SCRATCH_INITIAL_SIZE; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerDense; +import static org.apache.lucene.search.VectorScorer.Bulk.fromRandomScorerSparse; + +import java.io.IOException; +import org.apache.lucene.codecs.hnsw.FlatVectorsScorer; +import org.apache.lucene.codecs.lucene95.OffHeapByteVectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloat16VectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues; +import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; +import org.apache.lucene.index.ByteVectorValues; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.internal.hppc.IntArrayList; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.RandomAccessInput; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.LongValues; +import org.apache.lucene.util.hnsw.RandomVectorScorer; +import org.apache.lucene.util.packed.DirectReader; + +/** + * Vector values that share a single copy of each distinct vector across the documents and fields + * that reference it. + * + *

      Every instance is backed by two views: the {@code fieldView} maps ordinals to docs and drives + * iteration (one entry per document), while the {@code groupView} holds the de-duplicated vectors + * (one entry per distinct vector). {@code fieldOrdToGroupOrd} translates a document ordinal in the + * field into its group ordinal. + */ +sealed interface DedupVectorValues { + /** The dense view over distinct vectors, indexed by group ordinal. */ + KnnVectorValues getGroupView(); + + /** Maps a per-document ordinal to its group ordinal in {@link #getGroupView()}. */ + FieldOrdToGroupOrd getFieldOrdToGroupOrd(); + + /** + * Maps a field's per-document ordinal to the ordinal of its (shared) vector within the group. + * Backed on-heap while writing and off-heap while reading. + */ + sealed interface FieldOrdToGroupOrd { + int get(int ord); + + FieldOrdToGroupOrd copy() throws IOException; + } + + static ByteVectorValues loadDedupBytes( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapByteVectorValues fieldView = + OffHeapByteVectorValues.load( + function, vectorsScorer, configuration, BYTE, dimension, 0, 0, vectorData); + + final OffHeapByteVectorValues groupView = + new OffHeapByteVectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new ByteImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over byte vectors. */ + final class ByteImpl extends ByteVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final ByteVectorValues fieldView; + private final ByteVectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + ByteImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + ByteVectorValues fieldView, + ByteVectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public ByteVectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public byte[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public ByteImpl copy() throws IOException { + return new ByteImpl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(byte[] target) throws IOException { + if (size() == 0) { + return null; + } + ByteImpl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = copy.fieldView instanceof OffHeapByteVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + static FloatVectorValues loadDedupFloats( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapFloatVectorValues fieldView = + OffHeapFloatVectorValues.load( + function, vectorsScorer, configuration, FLOAT32, dimension, 0, 0, vectorData); + + final OffHeapFloatVectorValues groupView = + new OffHeapFloatVectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new FloatImpl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over float vectors. */ + final class FloatImpl extends FloatVectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final FloatVectorValues fieldView; + private final FloatVectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + FloatImpl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + FloatVectorValues fieldView, + FloatVectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public FloatVectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public float[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public FloatImpl copy() throws IOException { + return new FloatImpl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(float[] target) throws IOException { + if (size() == 0) { + return null; + } + FloatImpl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = copy.fieldView instanceof OffHeapFloatVectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + static Float16VectorValues loadDedupFloat16s( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + OrdToDocDISIReaderConfiguration configuration, + int dimension, + int groupNumVectors, + IndexInput vectorData, + long vectorDataOffset, + long vectorDataSize, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize) + throws IOException { + + final OffHeapFloat16VectorValues fieldView = + OffHeapFloat16VectorValues.load( + function, vectorsScorer, configuration, FLOAT16, dimension, 0, 0, vectorData); + + final OffHeapFloat16VectorValues groupView = + new OffHeapFloat16VectorValues.DenseOffHeapVectorValues( + dimension, + groupNumVectors, + vectorData.slice("group-slice", vectorDataOffset, vectorDataSize), + fieldView.getVectorByteLength(), + vectorsScorer, + function); + + final FieldOrdToGroupOrd fieldOrdToGroupOrd = + new FieldOrdToGroupOrdOffHeap(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + + return new Float16Impl(vectorsScorer, function, fieldView, groupView, fieldOrdToGroupOrd); + } + + /** {@link DedupVectorValues} over float16 vectors. */ + final class Float16Impl extends Float16VectorValues implements DedupVectorValues { + private final FlatVectorsScorer vectorsScorer; + private final VectorSimilarityFunction function; + private final Float16VectorValues fieldView; + private final Float16VectorValues groupView; + private final FieldOrdToGroupOrd fieldOrdToGroupOrd; + private int[] scratch; + + Float16Impl( + FlatVectorsScorer vectorsScorer, + VectorSimilarityFunction function, + Float16VectorValues fieldView, + Float16VectorValues groupView, + FieldOrdToGroupOrd fieldOrdToGroupOrd) { + this.vectorsScorer = vectorsScorer; + this.function = function; + this.fieldView = fieldView; + this.groupView = groupView; + this.fieldOrdToGroupOrd = fieldOrdToGroupOrd; + this.scratch = new int[SCRATCH_INITIAL_SIZE]; + } + + @Override + public Float16VectorValues getGroupView() { + return groupView; + } + + @Override + public FieldOrdToGroupOrd getFieldOrdToGroupOrd() { + return fieldOrdToGroupOrd; + } + + @Override + public int ordToDoc(int ord) { + return fieldView.ordToDoc(ord); + } + + @Override + public void prefetch(int[] ordsToPrefetch, int numOrds) throws IOException { + if (scratch.length < ordsToPrefetch.length) { // grow if needed + scratch = ArrayUtil.grow(scratch, ordsToPrefetch.length); + } + for (int i = 0; i < numOrds; i++) { + scratch[i] = fieldOrdToGroupOrd.get(ordsToPrefetch[i]); + } + groupView.prefetch(scratch, numOrds); + } + + @Override + public short[] vectorValue(int ord) throws IOException { + return groupView.vectorValue(fieldOrdToGroupOrd.get(ord)); + } + + @Override + public int dimension() { + return fieldView.dimension(); + } + + @Override + public int size() { + return fieldView.size(); + } + + @Override + public Float16Impl copy() throws IOException { + return new Float16Impl( + vectorsScorer, function, fieldView.copy(), groupView.copy(), fieldOrdToGroupOrd.copy()); + } + + @Override + public DocIndexIterator iterator() { + return fieldView.iterator(); + } + + @Override + public VectorScorer scorer(short[] target) throws IOException { + if (size() == 0) { + return null; + } + Float16Impl copy = copy(); + DocIndexIterator indexIterator = copy.iterator(); + RandomVectorScorer vectorScorer = vectorsScorer.getRandomVectorScorer(function, copy, target); + boolean isDense = + copy.fieldView instanceof OffHeapFloat16VectorValues.DenseOffHeapVectorValues; + return new DedupVectorScorer(indexIterator, vectorScorer, isDense); + } + } + + record DedupVectorScorer( + KnnVectorValues.DocIndexIterator indexIterator, + RandomVectorScorer vectorScorer, + boolean isDense) + implements VectorScorer { + + @Override + public float score() throws IOException { + return vectorScorer.score(indexIterator.index()); + } + + @Override + public DocIdSetIterator iterator() { + return indexIterator; + } + + @Override + public Bulk bulk(DocIdSetIterator matchingDocs) { + if (isDense) { + return fromRandomScorerDense(vectorScorer, indexIterator, matchingDocs); + } else { + return fromRandomScorerSparse(vectorScorer, indexIterator, matchingDocs); + } + } + } + + /** On-heap map used during a flush, backed directly by the buffered ordinals. */ + record FieldOrdToGroupOrdArrayList(IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + + @Override + public int get(int ord) { + return fieldOrdToGroupOrd.get(ord); + } + + @Override + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdArrayList(fieldOrdToGroupOrd); + } + } + + /** On-heap map used during a sorted flush, indirecting through a new-to-old ordinal map. */ + record FieldOrdToGroupOrdMappedArrayList(int[] map, IntArrayList fieldOrdToGroupOrd) + implements FieldOrdToGroupOrd { + + @Override + public int get(int ord) { + return fieldOrdToGroupOrd.get(map[ord]); + } + + @Override + public FieldOrdToGroupOrd copy() { + return new FieldOrdToGroupOrdMappedArrayList(map, fieldOrdToGroupOrd); + } + } + + /** Off-heap map used while reading, backed by a {@link DirectReader}. */ + record FieldOrdToGroupOrdOffHeap( + IndexInput vectorData, + long fieldOrdToGroupOrdOffset, + long fieldOrdToGroupOrdSize, + LongValues values) + implements FieldOrdToGroupOrd { + + FieldOrdToGroupOrdOffHeap( + IndexInput vectorData, long fieldOrdToGroupOrdOffset, long fieldOrdToGroupOrdSize) + throws IOException { + RandomAccessInput slice = + vectorData.randomAccessSlice(fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + LongValues values = DirectReader.getInstance(slice, FIELD_ORD_TO_GROUP_ORD_BITS_PER_VALUE); + this(vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize, values); + } + + @Override + public int get(int v) { + return (int) values.get(v); + } + + @Override + public FieldOrdToGroupOrd copy() throws IOException { + return new FieldOrdToGroupOrdOffHeap( + vectorData, fieldOrdToGroupOrdOffset, fieldOrdToGroupOrdSize); + } + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java index 0123d7153b3c..bf8b1d4e9fc2 100644 --- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/codecs/dedup/TestDedupFlatVectorsFormat.java @@ -22,6 +22,9 @@ import static org.hamcrest.Matchers.arrayContainingInAnyOrder; import static org.hamcrest.Matchers.instanceOf; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.hnsw.FlatVectorsReader; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader; @@ -41,7 +44,8 @@ import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.NumericDocValues; -import org.apache.lucene.sandbox.codecs.dedup.DedupUtil.DedupVectorValues; +import org.apache.lucene.internal.hppc.LongArrayList; +import org.apache.lucene.search.Query; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.TestUtil; @@ -318,6 +322,102 @@ public void testDuplicatesAcrossSegmentsDedupOnMerge() throws Exception { } } + /** Test that vectors not referenced are deleted from the group. */ + public void testDeletes() throws Exception { + float[] a = {1, 1, 1, 1}; + float[] b = {2, 2, 2, 2}; + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + + LongArrayList docsWithArrayB = new LongArrayList(); + boolean aIndexed = false, bIndexed = false; + for (int i = 0; i < 50; i++) { // many documents + Document doc = new Document(); + doc.add(new NumericDocValuesField("id", i)); + if (random().nextBoolean()) { // index either a or b + doc.add(new KnnFloatVectorField("f", a)); + aIndexed = true; + } else { + doc.add(new KnnFloatVectorField("f", b)); + docsWithArrayB.add(i); + bIndexed = true; + } + w.addDocument(doc); + } + + assumeTrue("Both vectors a and b indexed", aIndexed && bIndexed); + + w.forceMerge(1); // de-duplicate everything + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(2, groupNumVectors(values)); // the group has both vectors + } + + Query matchDocsWithArrayB = + NumericDocValuesField.newSlowSetQuery("id", docsWithArrayB.toArray()); + w.deleteDocuments(matchDocsWithArrayB); // delete all docs with vector b + + w.forceMerge(1); // de-duplicate everything + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("f"); + assertEquals(1, groupNumVectors(values)); // the group now has one vector + } + } + } + + /** Test many duplicates spread across fields, documents, segments. */ + public void testManyDuplicate() throws Exception { + float[] shared = {1, 2, 3, 4}; + List fields = new ArrayList<>(List.of("a", "b", "c", "d", "e")); + boolean atLeastOne = false; + + try (Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, config())) { + for (int i = 0; i < 50; i++) { // many documents + Document doc = new Document(); + + // randomly pick [0, N) fields to index the same vector + int numFields = random().nextInt(fields.size()); + Collections.shuffle(fields, random()); + for (int j = 0; j < numFields; j++) { + doc.add(new KnnFloatVectorField(fields.get(j), shared)); + atLeastOne = true; + } + + w.addDocument(doc); + + if (random().nextFloat() < 0.2f) { // randomly create segments + w.commit(); + } + } + + w.forceMerge(1); // de-duplicate everything + + assumeTrue("At least one vector indexed", atLeastOne); + + try (DirectoryReader reader = DirectoryReader.open(w)) { + LeafReader leaf = getOnlyLeafReader(reader); + + DedupFlatVectorsReader dedupReader = getDedupReader(leaf, "a"); + for (String field : fields) { // all fields DO resolve to the same group + DedupFlatVectorsReader other = getDedupReader(leaf, field); + assertEquals(dedupReader, other); // de-duplication happened correctly + assertEquals( // both fields DO resolve to the same group + dedupReader.getEntry("a", FLOAT32).groupInfo(), + other.getEntry(field, FLOAT32).groupInfo()); + } + + FloatVectorValues values = leaf.getFloatVectorValues("a"); + assertEquals(1, groupNumVectors(values)); // the group has one vector + assertArrayEquals(shared, values.vectorValue(0), 0f); + } + } + } + /** Number of distinct vectors physically stored for a field's group. */ private static int groupNumVectors(KnnVectorValues values) { assertThat(values, instanceOf(DedupVectorValues.class));