goankur opened a new pull request, #16871:
URL: https://github.com/apache/lucene/pull/16871
**Title:** Quantize merged vectors on the merge workers in
Lucene104HnswScalarQuantizedVectorsFormat
### Problem
When `Lucene104HnswScalarQuantizedVectorsFormat` merges a field, it builds
the merged graph on its merge workers, but
`Lucene104ScalarQuantizedVectorsWriter` quantizes every merged vector on the
merging thread alone. Asymmetric encodings quantize each vector twice (the
document code and the query record that the graph merge scores with). On large
segments, this serial pass becomes a large part of the merge: in one profile,
the merge spent most of its time in `OptimizedScalarQuantizer` on a single
thread while the box sat idle.
### Change
- The writer quantizes merged vectors in blocks of 8,192. Each block is
copied (the quantizer centers its input in place), split into `numMergeWorkers`
ranges that run in parallel, and written in document order. The output is
byte-for-byte the same as the serial path.
- It uses the same executor and opt-in as the graph merge in
`Lucene99HnswVectorsWriter`: the format's `mergeExec` if set, else
`MergeState#intraMergeTaskExecutor`, and only when `numMergeWorkers > 1`.
Fields of one block or less, and formats with one merge worker, keep the serial
path. Default behavior does not change.
- `Lucene104HnswScalarQuantizedVectorsFormat` passes its `numMergeWorkers`
and `mergeExec` to its flat format, through a package-private constructor.
`Lucene104ScalarQuantizedVectorsWriter` gets a public constructor that takes
both; the old one delegates with one worker.
### Benchmark
Force merge of 2M random 768-dim unit vectors (dot product, 1-bit codes with
4-bit queries, maxConn 16, beamWidth 100) from 4 segments to 1, with 8 merge
workers on both arms, so both build the graph in parallel and only the
quantization differs. m8gd.16xlarge (Graviton4, 64 vCPU), Corretto 25, 2 runs
per arm in the order main, PR, PR, main:
| Arm | Run 1 | Run 2 |
|---|---|---|
| main | 252.6 s | 246.8 s |
| this PR | 185.7 s | 186.0 s |
The merge takes 26% less time. The driver uses only public API and is below,
so the numbers can be reproduced.
<details><summary>MergeBench.java</summary>
```java
import java.nio.file.Path;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.lucene104.Lucene104Codec;
import
org.apache.lucene.codecs.lucene104.Lucene104HnswScalarQuantizedVectorsFormat;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.KnnFloatVectorField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.NoMergePolicy;
import org.apache.lucene.index.TieredMergePolicy;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.VectorUtil;
import
org.apache.lucene.util.quantization.QuantizedByteVectorValues.ScalarEncoding;
/**
* Times a force merge of a 1-bit quantized HNSW field.
*
* <p>build DIR DOCS DIM SEGMENTS: writes DOCS random unit vectors (seed 42)
as SEGMENTS segments.
* merge DIR WORKERS: force-merges DIR to one segment with WORKERS merge
workers, prints the time.
*/
public class MergeBench {
static final ScalarEncoding ENCODING =
ScalarEncoding.SINGLE_BIT_QUERY_NIBBLE;
static Lucene104Codec codec(KnnVectorsFormat format) {
return new Lucene104Codec() {
@Override
public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
return format;
}
};
}
public static void main(String[] args) throws Exception {
Path dir = Path.of(args[1]);
if (args[0].equals("build")) {
int docs = Integer.parseInt(args[2]), dim = Integer.parseInt(args[3]);
int segments = Integer.parseInt(args[4]);
IndexWriterConfig iwc =
new IndexWriterConfig()
.setCodec(codec(new
Lucene104HnswScalarQuantizedVectorsFormat(ENCODING, 16, 100, 1, null)))
.setMergePolicy(NoMergePolicy.INSTANCE)
// large enough that only the explicit flushes below cut
segments
.setRAMBufferSizeMB(8192);
Random random = new Random(42);
try (FSDirectory d = FSDirectory.open(dir);
IndexWriter w = new IndexWriter(d, iwc)) {
int perSegment = (docs + segments - 1) / segments;
for (int i = 0; i < docs; i++) {
float[] v = new float[dim];
for (int j = 0; j < dim; j++) {
v[j] = random.nextFloat() * 2 - 1;
}
Document doc = new Document();
doc.add(new KnnFloatVectorField("v", VectorUtil.l2normalize(v),
VectorSimilarityFunction.DOT_PRODUCT));
w.addDocument(doc);
if ((i + 1) % perSegment == 0) {
w.flush();
}
}
w.commit();
}
try (DirectoryReader r = DirectoryReader.open(FSDirectory.open(dir))) {
System.out.println("BUILT docs=" + r.numDocs() + " segments=" +
r.leaves().size());
}
} else {
int workers = Integer.parseInt(args[2]);
ExecutorService exec = workers > 1 ?
Executors.newFixedThreadPool(workers) : null;
IndexWriterConfig iwc =
new IndexWriterConfig()
.setCodec(codec(new
Lucene104HnswScalarQuantizedVectorsFormat(ENCODING, 16, 100, workers, exec)))
.setMergePolicy(new TieredMergePolicy());
try (FSDirectory d = FSDirectory.open(dir);
IndexWriter w = new IndexWriter(d, iwc)) {
long t0 = System.nanoTime();
w.forceMerge(1);
long ms = (System.nanoTime() - t0) / 1_000_000;
w.commit();
System.out.println("MERGE workers=" + workers + " ms=" + ms);
} finally {
if (exec != null) {
exec.shutdown();
}
}
try (DirectoryReader r = DirectoryReader.open(FSDirectory.open(dir))) {
System.out.println("MERGED docs=" + r.numDocs() + " segments=" +
r.leaves().size());
}
}
}
}
```
</details>
### Tests
- `TestParallelMergeQuantization`: the parallel writers produce the same
bytes as the serial ones for every `ScalarEncoding` and similarity, for
document codes and for query records, with random block sizes and worker counts.
- An end-to-end test force-merges the same documents with one merge worker
and with four (through the format's executor or the merge's own) and compares
every merged quantized vector and its corrective terms. With the parallel path
deliberately broken, all three tests fail.
- `./gradlew tidy check -x test` and the lucene104, lucene99, HNSW and KNN
query suites pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
