Skip to content

Commit 23ed23a

Browse files
kaivalnpkaivalnp
committed
Add a de-duplicating vector format
1 parent 99fc6a1 commit 23ed23a

16 files changed

Lines changed: 2757 additions & 2 deletions

‎lucene/CHANGES.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,9 @@ New Features
321321
and document length instead of corpus statistics such as document frequency. It also supports k3
322322
query-term frequency saturation. (Tianxiao Wei)
323323

324+
* GITHUB#15979: Add a de-duplicating HNSW vector format (Lucene106DedupHnswVectorsFormat) that stores
325+
each distinct vector once, shared across all documents and fields that reference it. (Kaival Parikh)
326+
324327
Improvements
325328
---------------------
326329

‎lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatFieldVectorsWriter.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.lucene.codecs.lucene106.dedup;
18+
19+
import java.io.IOException;
20+
import java.util.ArrayList;
21+
import java.util.List;
22+
import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter;
23+
import org.apache.lucene.index.DocsWithFieldSet;
24+
import org.apache.lucene.internal.hppc.IntArrayList;
25+
import org.apache.lucene.internal.hppc.ObjectCursor;
26+
import org.apache.lucene.util.RamUsageEstimator;
27+
28+
/**
29+
* Buffers one field's vectors during flush, de-duplicating them through a shared {@link
30+
* DedupGroup}.
31+
*
32+
* @lucene.experimental
33+
*/
34+
final class DedupFlatFieldVectorsWriter<T> extends FlatFieldVectorsWriter<T> {
35+
private static final long SHALLOW_SIZE =
36+
RamUsageEstimator.shallowSizeOfInstance(DedupFlatFieldVectorsWriter.class);
37+
38+
private final DedupGroup<T> group;
39+
private final DocsWithFieldSet docsWithFieldSet;
40+
private final List<T> vectors;
41+
private final IntArrayList ordToVecOrd;
42+
private int lastDocID;
43+
private boolean finished;
44+
45+
DedupFlatFieldVectorsWriter(DedupGroup<T> group) {
46+
this.group = group;
47+
this.docsWithFieldSet = new DocsWithFieldSet();
48+
this.vectors = new ArrayList<>();
49+
this.ordToVecOrd = new IntArrayList();
50+
this.lastDocID = -1;
51+
this.finished = false;
52+
}
53+
54+
@Override
55+
public List<T> getVectors() {
56+
return vectors;
57+
}
58+
59+
@Override
60+
public DocsWithFieldSet getDocsWithFieldSet() {
61+
return docsWithFieldSet;
62+
}
63+
64+
IntArrayList getOrdToVecOrd() {
65+
return ordToVecOrd;
66+
}
67+
68+
@Override
69+
public void finish() {
70+
if (finished) {
71+
throw new IllegalStateException("already finished");
72+
}
73+
finished = true;
74+
}
75+
76+
@Override
77+
public boolean isFinished() {
78+
return finished;
79+
}
80+
81+
@Override
82+
public T copyValue(T vectorValue) {
83+
throw new UnsupportedOperationException(); // handled inside group
84+
}
85+
86+
@Override
87+
public void addValue(int docID, T vectorValue) throws IOException {
88+
if (finished) {
89+
throw new IllegalStateException("already finished");
90+
} else if (docID <= lastDocID) {
91+
throw new IllegalArgumentException(
92+
"docID=" + docID + " not going forwards, indexed lastDocID=" + lastDocID);
93+
}
94+
95+
lastDocID = docID;
96+
docsWithFieldSet.add(docID);
97+
98+
ObjectCursor<T> cursor = group.addUnique(vectorValue);
99+
vectors.add(cursor.value); // owned vector value
100+
ordToVecOrd.add(cursor.index); // index in group
101+
}
102+
103+
@Override
104+
public long ramBytesUsed() {
105+
return SHALLOW_SIZE
106+
+ docsWithFieldSet.ramBytesUsed()
107+
+ (long) vectors.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF
108+
+ ordToVecOrd.ramBytesUsed();
109+
}
110+
}

‎lucene/core/src/java/org/apache/lucene/codecs/lucene106/dedup/DedupFlatVectorsFormat.java‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.lucene.codecs.lucene106.dedup;
18+
19+
import java.io.IOException;
20+
import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
21+
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
22+
import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
23+
import org.apache.lucene.index.SegmentReadState;
24+
import org.apache.lucene.index.SegmentWriteState;
25+
26+
/**
27+
* Flat vector format that stores each distinct vector once.
28+
*
29+
* <p>Vectors that share the same dimension and encoding form a <i>group</i>. Within a group, an
30+
* identical vector is stored a single time regardless of how many documents (across all fields that
31+
* map to that group) reference it; each field then keeps a per-document {@code ordToVecOrd} map
32+
* from its document ordinal to the group ordinal of the shared vector. This is well suited to
33+
* indexes with repeated vectors, e.g. several fields derived from the same embedding, or heavily
34+
* duplicated content.
35+
*
36+
* <h2>.vdd (vector de-dup data) file</h2>
37+
*
38+
* <ul>
39+
* <li>For each group, its distinct vectors, aligned to 4 bytes (BYTE) or 64 bytes (FLOAT32).
40+
* <li>For each field:
41+
* <ul>
42+
* <li>The sparse-encoding data (only when some documents lack the field): DocIds encoded by
43+
* {@link
44+
* org.apache.lucene.codecs.lucene90.IndexedDISI#writeBitSet(org.apache.lucene.search.DocIdSetIterator,
45+
* org.apache.lucene.store.IndexOutput, byte)}, followed by the ordinal-to-doc mapping
46+
* encoded by {@link org.apache.lucene.util.packed.DirectMonotonicWriter}.
47+
* <li>The {@code ordToVecOrd} map (aligned to 4 bytes): one entry per document ordinal
48+
* giving the group ordinal of the shared vector, packed by {@link
49+
* org.apache.lucene.util.packed.DirectWriter}.
50+
* </ul>
51+
* </ul>
52+
*
53+
* <h2>.vdm (vector de-dup metadata) file</h2>
54+
*
55+
* <p>A list of groups, each:
56+
*
57+
* <ul>
58+
* <li><b>[int32]</b> group ordinal
59+
* <li><b>[int32]</b> vector dimension
60+
* <li><b>[int32]</b> vector encoding ordinal
61+
* <li><b>[int32]</b> group size (number of distinct vectors)
62+
* <li><b>[int64]</b> offset to this group's vectors in the .vdd file
63+
* <li><b>[int64]</b> length of this group's vectors, in bytes
64+
* </ul>
65+
*
66+
* <p>terminated by <b>[int32]</b> {@code -1}, then a list of fields, each:
67+
*
68+
* <ul>
69+
* <li><b>[int32]</b> field number
70+
* <li><b>[int32]</b> vector similarity function ordinal
71+
* <li><b>[int32]</b> vector dimension
72+
* <li><b>[int32]</b> vector encoding ordinal
73+
* <li><b>[int32]</b> ordinal of the group holding this field's vectors
74+
* <li><b>[int32]</b> the number of documents having values for this field
75+
* <li>the sparse-encoding metadata (docs-with-field offset/length and ordToDoc configuration), as
76+
* written by {@link
77+
* org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration#writeStoredMeta}
78+
* <li><b>[int64]</b> offset to this field's {@code ordToVecOrd} map in the .vdd file
79+
* <li><b>[int64]</b> length of this field's {@code ordToVecOrd} map, in bytes
80+
* </ul>
81+
*
82+
* <p>also terminated by <b>[int32]</b> {@code -1}.
83+
*
84+
* <h2>Complexity</h2>
85+
*
86+
* <p>Let {@code N} be the number of indexed vectors (one per document per field), {@code U} the
87+
* number of distinct vectors, and {@code d} the dimension. De-duplication interns each vector via a
88+
* hash lookup with linear probing, resolving hash collisions with a full equality check.
89+
*
90+
* <ul>
91+
* <li><b>Indexing (flush):</b> expected {@code O(N * d)} time (a hash plus occasional equality
92+
* check per vector). Heap is {@code O(U * d)} for the distinct vectors held in the group,
93+
* plus {@code O(N)} for the per-document references and {@code ordToVecOrd} entries.
94+
* <li><b>Merge:</b> expected {@code O(N * d)} time; distinct vectors are written to disk as soon
95+
* as they are first seen rather than buffered, so heap stays {@code O(N)} (the per-field
96+
* {@code ordToVecOrd} maps and light per-vector handles) with no {@code O(U * d)} term. When
97+
* a source segment is itself in this format, equality is decided by comparing group ordinals
98+
* in {@code O(1)} without reading the vectors back.
99+
* <li><b>Reading:</b> both the vectors and the {@code ordToVecOrd} map stay off-heap; a read
100+
* resolves a document ordinal to its vector via one extra {@code ordToVecOrd} lookup.
101+
* </ul>
102+
*
103+
* @lucene.experimental
104+
*/
105+
final class DedupFlatVectorsFormat extends FlatVectorsFormat {
106+
static final String NAME = "Lucene106DedupFlatVectorsFormat";
107+
108+
static final String META_CODEC_NAME = "Lucene106DedupFlatVectorsFormatMeta";
109+
static final String META_EXTENSION = "vdm";
110+
111+
static final String VECTOR_DATA_CODEC_NAME = "Lucene106DedupFlatVectorsFormatVectorData";
112+
static final String VECTOR_DATA_EXTENSION = "vdd";
113+
114+
static final int VERSION_START = 0;
115+
static final int VERSION_CURRENT = VERSION_START;
116+
117+
private static final DedupFlatVectorsScorer FLAT_VECTORS_SCORER = new DedupFlatVectorsScorer();
118+
119+
DedupFlatVectorsFormat() {
120+
super(NAME);
121+
}
122+
123+
@Override
124+
public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
125+
return new DedupFlatVectorsWriter(state, FLAT_VECTORS_SCORER);
126+
}
127+
128+
@Override
129+
public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
130+
return new DedupFlatVectorsReader(state, FLAT_VECTORS_SCORER);
131+
}
132+
}

0 commit comments

Comments
 (0)