Parallel O_DIRECT rerank reads for larger-than-RAM KNN search

[goankur]

Title:

(Implemented with AI, with Human in the loop)

Parallel O_DIRECT full-precision rerank reads for larger-than-RAM KNN search

Description (Reviewed and Edited by Human)

For a quantized-plus-rerank KNN search over an index larger than RAM, cost is dominated by fetching
the full-precision vectors of the candidate shortlist. RescoreTopNQuery scores one document at a
time, so a shortlist of M candidates becomes M dependent blocking reads at a device queue depth of
~1.

Batching per input is not enough either. A shortlist is spread across segments, so M candidates over
N segments leaves only ~M/N reads per .vec input: a per-input batch tops out at M/N in flight and
the query still pays N sequential rounds. The depth that matters is M, which requires gathering reads
across the inputs of the whole index.

This PR adds that, plus an example Directory in misc that serves it with O_DIRECT.

VectorBatchCapable (o.a.l.store) — optional capability of an IndexInput, in the spirit of
RandomAccessInput:

VectorBatch newBatch() throws IOException;

VectorBatch (o.a.l.store) — the cross-input accumulator; nothing is read until execute():

boolean add(IndexInput in, long[] positions, int dim, int count, float[] out) throws IOException;
void execute() throws IOException;

add returns false for an input the batch does not recognise, so the caller falls back for that input
alone.

FlatVectorsReader gains newRawVectorBatch(field) and addRawVectors(...), defaulting to "not
supported", so a codec can contribute raw float32 vectors without exposing its layout.
Lucene99FlatVectorsReader implements them only when its vector data input offers the capability.

RescoreTopNQuery.rewrite queues every segment's shortlist, issues one execute() per query,
then scores. No core Directory implements the capability, so the default path is unchanged.

SelectiveDirectIODirectory (misc) routes only .vec through O_DIRECT — the HNSW graph,
quantized codes and metadata stay on the mmap delegate and stay page-cached — and serves a batch
through a caller-owned Executor, sized independently of the searcher's executor.

Lucene99FlatVectorsWriter page-aligns FLOAT32 vector data to 4 KB, so a vector is one block
rather than straddling two.

O_DIRECT is opt-in and confined to .vec for two measured reasons: read-ahead is waste on a random
shortlist (27.6 KB fetched per 4 KB vector, below), and one-shot rerank pages evict the graph and
codes that are worth caching.

How the numbers were obtained

AWS G6 instance, x86_64, AL2023, kernel 6.1.182, local NVMe. JDK 25 (Corretto), Panama Vector API.

25M Cohere v3 embeddings, 1024-dim fp32, DOT_PRODUCT; HNSW maxConn 64 / beamWidth 250, 1-bit BBQ
plus full-precision rerank via RescoreTopNQuery; default maxMergedSegmentMB; ~104 GB on disk,
~101 GB of it vector data.

Every run is confined to a systemd scope with MemoryMax=10G against that ~101 GB of vectors, page
cache dropped before every run. Driver is lucene-util's KnnGraphTester, 10,000 queries, topK 100, against precomputed exact-NN ground truth so recall is comparable across runs. Every row below uses overSample 5 / fanout 100, so only the read path varies.

Single-stream is one query at a time. Concurrent is open-loop: Poisson arrivals, bounded 32 server
threads with a 32-deep queue, load shedding when full; SLA-QPS is the highest offered rate holding
p99 ≤ 50 ms with ≤ 0.1% shed.

I/O figures are iostat -x averages over the measured phase (r/s, rkB/s, rareq-sz, aqu-sz,
r_await), excluding idle samples. %util is not used — it pins near 100% on NVMe well before the
device is busy.

Bottleneck attribution before optimizing: PSI io.pressure full ≈ 19% — 19% of wall time with every
runnable task stalled on I/O — against ~50% idle CPU, corroborated by eBPF offcputime in
uninterruptible sleep.

Results

Single-stream, 10,000 queries:

Read path	Recall	p99	QPS	Read IOPS	Read GB/s	Avg read	aqu-sz	r_await
Stock HNSW + mmap	0.969	326.0 ms	7	18.0k	0.46	27.6 KB	5.9	0.33 ms
O_DIRECT, unaligned vectors	0.969	16.5 ms	75	76.9k	0.29	4.0 KB	11.2	0.15 ms
O_DIRECT, 4 KB-aligned vectors	0.967	13.0 ms	99	51.9k	0.20	~4 KB	7.2	0.14 ms

• mmap spends 0.46 GB/s of device bandwidth, over 2× the aligned path's 0.20 GB/s, while being
  25× slower at p99 — it moves 27.6 KB to deliver 4 KB.
• Alignment costs operations, not bytes: unaligned needs 76.9k IOPS against 51.9k (1.48×) at the
  same ~4 KB per read. Its higher aqu-sz (11.2 vs 7.2) is extra outstanding work, not useful depth.
• r_await is 0.14–0.15 ms, so at a 13 ms p99 the device accounts for under 2% of latency; the rest
  was queueing and serialization.

The aligned row was measured twice from a dropped cache: p99 13.02 and 12.91 ms, recall 0.967 both,
IOPS within 0.2%, aqu-sz within 2%. The mmap baseline varies more across cold 40-minute passes (an
earlier run of the same stock code gave 235 ms).

Concurrent, aligned O_DIRECT, 32 read threads — SLA-QPS ~110:

Offered QPS	Achieved	p50	p99	Shed	Within SLA
100	99.1	13.3 ms	44.1 ms	0	✓
110	111.1	14.0 ms	49.4 ms	0	✓
120	122.1	15.0 ms	52.2 ms	0	✗
140	140.3	16.4 ms	70.2 ms	0	✗
160	159.6	17.5 ms	76.0 ms	0	✗

Every point, including the failures, achieved its offered rate with zero shedding — the ceiling is
latency from queueing, not admission control.

Recall is tunable against throughput independently of this change: at overSample 3.5 / fanout 25
the same code sustains ~190 SLA-QPS at recall 0.937.

Scope and known gaps

• Behind a capability check throughout; no core Directory implements it, so nothing changes by
  default.
• The new core API has no test in core, for that same reason — coverage lives in misc
  (TestSelectiveDirectIODirectory: serial and parallel batch reads, and a query where the first
  segment cannot supply a batch while later ones can). Happy to add a test-only batch-capable
  FilterDirectory under core's test tree if that is the preferred shape.
• Recall differs slightly between rows (0.967 vs 0.969) because 4 KB alignment changes the index, not
  because of the read path.

Companion PR

#16666 adds an io_uring-backed Directory in sandbox on top of this API, where a single submission
spans every segment's .vec — reaching aqu-sz 16.3 against 7.2 here, for 9.53 ms single-stream p99
and ~180 SLA-QPS at the same operating point and recall. Kept separate: this PR is pure JDK
(ExtendedOpenOption.DIRECT) with no native dependency.

[jimczi]

Thanks for this.
The conclusion that the time is lost waiting rather than computing looks right.
I also think the structure is right: queue the candidates, submit them together, then score them.
My point is that Lucene already has a hook for this: IndexInput#prefetch(offset, length), “start this read, I’ll consume it later.” So the three phases don’t need a new SPI.

RescoreTopNQuery never calls it today. That means the “stock mmap” result is really “mmap at queue depth 1.” It’s slow because nothing tells the kernel about the next candidate, not because it’s mmap.

I’d also revisit the two reasons for routing .vec around the page cache. Lucene already applies MADV_RANDOM to that file through DataAccessHint.RANDOM in Lucene99FlatVectorsReader, and it addresses both:

• Read amplification. VM_RAND_READ disables readahead in the mmap fault path. Both do_sync_mmap_readahead and do_async_mmap_readahead return early, so a fault reads only the pages it needs. That makes the 27.6 KB read per 4 KB vector surprising. I’d double-check that --enable-native-access was enabled for that run. Without it, madvise never runs and MemorySegmentIndexInput#prefetch always returns false.

• Eviction of the graph and codes. The mmap fault path doesn’t mark folios as accessed. A page faulted once under a random mapping goes onto the inactive list and gets reclaimed first. The graph and quantized codes are touched on every query, so they stay active. As long as the rerank working set is manageable, the page cache already acts like the bounded rerank buffer you’re building by hand. It would also be worth testing with MGLRU enabled and disabled. MGLRU marks mapped folios active on fault, and it’s disabled by default on 6.1, so your current runs use classic LRU.

This is why mmap + MADV_RANDOM + batched WILLNEED is hard to beat here. There’s already a lot of relevant work:

• Introduce a pread Directory based on Panama-FFI ? #16044 started from the same observation: mmap page-fault storms in cgroup-limited containers with indices larger than RAM. It proposed a native pread Directory and grew into a broad study across three platforms, four cache regimes, 1–16 threads, mmap NORMAL, MADV_RANDOM, batched WILLNEED, FFI pread, FileChannel, and pread + O_DIRECT.

  Batched prefetch won every memory-pressured case. The biggest margin was exactly where you are: one thread, cold, and file ≫ RAM.

  16 KB random reads, cold, file > RAM, NVMe (ops/ms)	T01	T08	T16
  pread	0.58	4.02	4.65
  mmap, no prefetch (MADV_RANDOM)	0.15	1.10	1.94
  mmap + batched prefetch	4.19	5.62	5.98

  That T01 result is around 67k IOPS / 1.25 GB/s from one thread, which is what fio gets from that device at iodepth 16. It’s a different box from your g6, so it’s only indicative, but it’s the same order as your 67.3k without a ring or native dependency.

• Add JMH benchmarks comparing read I/O strategies under memory pressure #16279 is the JMH harness behind it — basically fio in Java over Lucene’s store primitives. It’s the easiest way to compare prefetch with your O_DIRECT results on the same hardware without touching the search path.

• MemorySegmentIndexInput: always prefetch on RANDOM mode #16145 and Remove mmap isLoaded check before madvise #14156 cover the weak spot in mmap prefetch. The power-of-two backoff in MemorySegmentIndexInput#prefetch suppresses madvise when the index barely fits in RAM, exactly when it matters most. That’s a live tuning issue in one method.

Here’s how we do it in Elasticsearch using APIs Lucene already has. The ring sits outside the leaf loop, so a candidate in segment N+1 can be prefetched while segment N is still being scored. The window isn’t tied to one input, so there’s no per-file ceiling:

PrefetchRing ring = new PrefetchRing(WINDOW);   // WINDOW = 100 for us; doc ids, never vectors

for (LeafReaderContext leaf : reader.leaves()) {
  FloatVectorValues values = leaf.reader().getFloatVectorValues(field);
  Scorer inner = weight.scorer(leaf);
  if (values == null || inner == null) continue;

  VectorScorer rescorer = values.rescorer(queryVector);      // full precision
  KnnVectorValues.DocIndexIterator vectorIter = values.iterator();
  DocIdSetIterator conj =
      ConjunctionUtils.intersectIterators(List.of(vectorIter, inner.iterator()));

  for (int doc = conj.nextDoc(); doc != NO_MORE_DOCS; doc = conj.nextDoc()) {
    values.prefetch(vectorIter.index());                     // fire and forget
    if (ring.isFull()) {
      // oldest entry was prefetched WINDOW candidates ago; scores via VectorScorer.Bulk
      ring.advance(scoreOldest(ring, buffer, results));
    }
    ring.append(doc, leaf.docBase, rescorer);
  }
}
while (ring.size() > 0) ring.advance(scoreOldest(ring, buffer, results));

Heap stays flat because the ring holds WINDOW doc IDs instead of count × dim floats across a barrier. Outstanding submissions from one thread provide the depth. There’s no read pool and no context-switch cost.

The current limitation is that values.prefetch(ord) doesn’t work. KnnVectorValues#prefetch only accepts an array and returns early below two ords. Also, Lucene104ScalarQuantizedVectorsReader.ScalarQuantizedVectorValues (what getFloatVectorValues() returns for a quantized field) forwards vectorValue and rescorer, but not prefetching. So it silently does nothing on the BBQ rerank path.

We hit the same issue in Elasticsearch and fixed it in our wrapper. The upstream change is small: add prefetch(int ord), make the array version loop over it, and forward it from the wrapper like the other methods. I’m happy to open that PR so you have something to build on.

This doesn’t rule out O_DIRECT. It should also be able to rescore efficiently, but it doesn’t need the new SPI either. The input can submit asynchronously during prefetch and reap during readFloats. Inputs from the same directory can share a ring. Submissions can also accumulate until the first dependent read forces a flush, which is what VectorBatch#execute expresses.

The query loop stays the same and the store chooses the mechanism.

Would you be up for measuring the prefetch path before adding the API, on the same box and at the same operating point? I still expect mmap to be better, but if O_DIRECT clearly wins, that’s a real result. The next step would be your Directory behind prefetch, not a new interface.