john-mlika opened a new issue, #16420:
URL: https://github.com/apache/lucene/issues/16420
### Description
`Lucene104ScalarQuantizedVectorsWriter.mergeAndRecalculateCentroids` has a
fast path that computes
the merged centroid as the vector-count-weighted mean of the source
segments' stored centroids,
skipping any data read. It is unreachable for HNSW-wrapped fields — that is,
for
`Lucene104HnswScalarQuantizedVectorsFormat`:
```java
// Lucene104ScalarQuantizedVectorsWriter (main:423-429)
static float[] getCentroid(KnnVectorsReader vectorsReader, String fieldName)
{
vectorsReader = vectorsReader.unwrapReaderForField(fieldName);
if (vectorsReader instanceof Lucene104ScalarQuantizedVectorsReader reader)
{
return reader.getCentroid(fieldName);
}
return null;
}
```
`unwrapReaderForField`'s default returns `this` and only
`PerFieldKnnVectorsFormat.FieldsReader`
overrides it (`KnnVectorsReader.java:58-60`), so for an HNSW-wrapped field
the unwrap stops at
`Lucene99HnswVectorsReader`, the `instanceof` fails, `getCentroid` returns
`null`, and
`mergeAndRecalculateCentroids` (main:431+) sets `recalculate = true` —
triggering
`calculateCentroid`: a **full pass over every float of every source
segment**, on **every** merge,
including the delete-free ones the fast path exists for.
That pass is one of **four** full reads the merge makes of each source
`.vec` — the `checkIntegrity`
checksum (`KnnVectorsWriter#merge:117`), the raw-vector copy
(`Lucene99FlatVectorsWriter#mergeOneFlatVectorField:279`), the centroid
recomputation
(`calculateCentroid:482`) and the re-quantization pass
(`Lucene104ScalarQuantizedVectorsWriter#mergeOneFlatVectorField:339`).
Measured on a delete-free `forceMerge(1)` of two 500-doc, 64-dim segments,
with byte counting at a
`FilterDirectory` and each `slice()` tagged by its calling frame,
eliminating the centroid pass cuts
source-segment raw-vector reads by **exactly one full pass** — 256,000
bytes, to the byte:
| metric | unpatched | with the fix below | delta |
|---|---|---|---|
| source-segment `.vec` reads | 1,024,516 B (4.002 passes) | 768,516 B
(3.002) | **−25.0%** |
| all reads from the source segments | 1,111,226 B | 855,226 B | −23.0% |
| total merge read volume (all files) | 7,814,144 B | 7,558,144 B | −3.3% |
The last row is the honest context: HNSW graph construction re-reading the
merged segment's `.veq`
is 6.5 MB of that total (83%), so this dominates only the *source-read*
side. What it removes is a
recomputation of metadata already sitting in the field entries, on every
delete-free merge.
### The obvious fix does not work
Overriding `unwrapReaderForField` on `Lucene99HnswVectorsReader` to return
its flat reader is not
viable. I checked the callers, and it breaks:
- `HnswUtil.graphIsRooted` (the unwrap no longer surfaces the
`HnswGraphProvider`),
- two CheckIndex paths,
- `SlowCodecReaderWrapper.getOffHeapByteSize`,
- `TestLucene104HnswScalarQuantizedVectorsFormat.testSimpleOffHeapSize`.
`unwrapReaderForField` is a general-purpose unwrap whose callers need the
HNSW layer; its semantics
shouldn't change for this.
### Proposed fix: route centroid access through `QuantizedVectorsReader`
`Lucene99HnswVectorsReader` **already implements** `QuantizedVectorsReader`,
delegating to its flat
reader when that reader is itself a `QuantizedVectorsReader` (main:74-75
declaration, :426/:434/:443
delegation guards). Riding that interface avoids touching the unwrap:
1. Add a centroid accessor to the `@lucene.experimental` interface — the
default keeps every
existing implementor source- and binary-compatible:
```java
// org.apache.lucene.util.quantization.QuantizedVectorsReader
/** The stored centroid for {@code fieldName}, or {@code null} if this
reader has none. */
default float[] getCentroid(String fieldName) {
return null;
}
```
2. `Lucene104ScalarQuantizedVectorsReader` marks its existing
`getCentroid(String)` as the interface
override — the method already exists; the writer calls it today after the
concrete-class cast.
3. `Lucene99HnswVectorsReader` forwards it in its existing delegation block:
```java
@Override
public float[] getCentroid(String fieldName) {
if (flatVectorsReader instanceof QuantizedVectorsReader qvr) {
return qvr.getCentroid(fieldName);
}
return null;
}
```
4. The writer's helper checks the interface after the unchanged,
PerField-only unwrap:
```java
static float[] getCentroid(KnnVectorsReader vectorsReader, String fieldName)
{
vectorsReader = vectorsReader.unwrapReaderForField(fieldName);
if (vectorsReader instanceof QuantizedVectorsReader reader) {
return reader.getCentroid(fieldName);
}
return null;
}
```
No change to `unwrapReaderForField` or its callers, and
`mergeAndRecalculateCentroids` is untouched —
its `centroid == null || liveDocs != null` recalc logic already does the
right thing once centroids
arrive. Non-quantized readers return `null` via the default and recalculate
exactly as today.
### Testing it
I verified the above with a same-package test
(`org.apache.lucene.codecs.lucene104`) that calls the
package-private `getCentroid` directly and measures merge I/O by counting
bytes at a
`FilterDirectory`, tagging each `slice()` with its calling frame so reads
attribute to individual
merge passes. Counting inside a wrapping `KnnVectorsReader` delegate turned
out to be the wrong
instrument — wrapping the reader changes what the `instanceof` /
`QuantizedVectorsReader` checks see
and perturbs the very path being measured.
On unpatched `main` it reports:
```
[centroid] top reader = PerFieldKnnVectorsFormat$FieldsReader
[centroid] unwrapReaderForField = Lucene99HnswVectorsReader
[centroid] Lucene104ScalarQuantizedVectorsWriter.getCentroid(...) = null
```
With the four-step sketch applied, the same line reads `= float[64]
0.520761`, and the merge's `.vec`
read volume drops by exactly one full pass over the source floats (1,280,774
→ 1,024,774 bytes for
1000 × 64-dim vectors). A run of `org.apache.lucene.codecs.lucene104.*` and
`org.apache.lucene.codecs.lucene99.*` with the sketch applied is clean (338
tests; the only failure
is that suite's own deliberately inverted `assertNull` on `getCentroid`).
Worth covering in a real PR beyond that: the delete case (`liveDocs !=
null`), where
`calculateCentroid` must still run, and a regression sweep over the
`unwrapReaderForField` callers
listed above, which this approach leaves untouched.
I'm happy to put up the PR, with the test, if the approach looks right.
--
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]
