From 9ec4305fe5fb5b1b99a1610e599442ab54ac1c93 Mon Sep 17 00:00:00 2001 From: John Maurice Date: Sat, 25 Jul 2026 04:27:21 +0300 Subject: [PATCH 1/3] Use direct I/O for merge-time reads of raw vectors --- .../mapping-reference/dense-vector.md | 2 +- .../DirectIOCapableFlatVectorsFormat.java | 27 ++- .../codec/vectors/DirectIOMergeHint.java | 20 ++ .../codec/vectors/MergeReaderWrapper.java | 40 +++- ...ectIOCapableLucene99FlatVectorsFormat.java | 23 -- .../index/store/AsyncDirectIOIndexInput.java | 8 + .../index/store/FsDirectoryFactory.java | 28 ++- ...DirectIOCapableFlatVectorsFormatTests.java | 222 ++++++++++++++++++ .../store/AsyncDirectIOIndexInputTests.java | 64 +++++ 9 files changed, 393 insertions(+), 41 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java create mode 100644 server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java diff --git a/docs/reference/elasticsearch/mapping-reference/dense-vector.md b/docs/reference/elasticsearch/mapping-reference/dense-vector.md index 7355d24a34e10..6e99968a31f60 100644 --- a/docs/reference/elasticsearch/mapping-reference/dense-vector.md +++ b/docs/reference/elasticsearch/mapping-reference/dense-vector.md @@ -602,7 +602,7 @@ $$$dense-vector-index-options$$$ :::: `on_disk_rescore` {applies_to}`stack: preview 9.3` {applies_to}`serverless: unavailable` -: (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`. +: (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. For `bbq_hnsw` indices, segment merges also read the raw vector data directly from disk where the platform supports it, so that merging is less likely to evict more frequently accessed data from the filesystem cache. {applies_to}`stack: ga 9.6` This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`. `auto_calibrate` {applies_to}`stack: ga 9.5` : (Optional, boolean) Only applicable to `bbq_disk`. When `true`, {{es}} automatically selects the optimal quantization encoding, oversampling factor, and preconditioning for each merged segment based on the actual recall characteristics of the merged corpus. Segments containing fewer than 10,000 vectors after merging are not calibrated and, when not otherwise specified in the mappings, use the default oversampling factor of 3.0x. Defaults to `false`. Cannot be changed after the field is created. Refer to [Auto-calibration for `bbq_disk`](/reference/elasticsearch/mapping-reference/bbq.md#bbq-auto-calibration) for details. diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java index fade2384441cf..0ab03af7c8dce 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java @@ -39,7 +39,8 @@ public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException { if (state.context.context() == IOContext.Context.DEFAULT && useDirectIO && canUseDirectIO(state)) { - // only override the context for the random-access use case + // only wrap readers opened for searching (DEFAULT context); the wrapper adds a + // lazily-created, merge-hinted direct I/O reader for merges SegmentReadState directIOState = new SegmentReadState( state.directory, state.segmentInfo, @@ -47,8 +48,16 @@ public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectI new DirectIOContext(state.context.hints()), state.segmentSuffix ); - // Use mmap for merges and direct I/O for searches. - return new MergeReaderWrapper(createReader(directIOState), createReader(state)); + SegmentReadState mergeDirectIOState = new SegmentReadState( + state.directory, + state.segmentInfo, + state.fieldInfos, + new DirectIOContext(state.context.hints(), Set.of(DirectIOHint.INSTANCE, DirectIOMergeHint.INSTANCE)), + state.segmentSuffix + ); + // Use direct I/O for merges too, so merge reads do not evict hotter data from the + // page cache. The merge hint makes those reads use a merge-sized buffer. + return new MergeReaderWrapper(createReader(directIOState), () -> createReader(mergeDirectIOState)); } else { return createReader(state); } @@ -56,11 +65,17 @@ public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectI protected static class DirectIOContext implements IOContext { + private final Set stickyHints; final Set hints; public DirectIOContext(Set hints) { - // always add DirectIOHint to the hints given - this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE)); + this(hints, Set.of(DirectIOHint.INSTANCE)); + } + + public DirectIOContext(Set hints, Set stickyHints) { + // the sticky hints are always added, and survive withHints() replacing the others + this.stickyHints = stickyHints; + this.hints = Sets.union(hints, stickyHints); } @Override @@ -85,7 +100,7 @@ public Set hints() { @Override public IOContext withHints(FileOpenHint... hints) { - return new DirectIOContext(Set.of(hints)); + return new DirectIOContext(Set.of(hints), stickyHints); } } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java new file mode 100644 index 0000000000000..a60cfc4c36a8d --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java @@ -0,0 +1,20 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors; + +import org.apache.lucene.store.IOContext; + +/** + * Hint that a file opened with direct I/O will be read as a long sequential stream by a merge, + * so reads should use a merge-sized buffer rather than the small buffer used for random access. + */ +public enum DirectIOMergeHint implements IOContext.FileOpenHint { + INSTANCE +} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java index ebbadddecef30..d4e4e84cea005 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java @@ -16,7 +16,9 @@ import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.search.AcceptDocs; import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.store.AlreadyClosedException; import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.IOSupplier; import org.apache.lucene.util.hnsw.RandomVectorScorer; import org.elasticsearch.core.IOUtils; @@ -24,14 +26,22 @@ import java.util.Collection; import java.util.Map; +/** + * A {@link FlatVectorsReader} that serves searches from one reader and merges from a second, + * lazily-created reader, so that the two can use different I/O strategies (such as different + * direct I/O buffer sizes) without sharing input state. + */ public class MergeReaderWrapper extends FlatVectorsReader { private final FlatVectorsReader mainReader; - private final FlatVectorsReader mergeReader; + private final IOSupplier mergeReaderSupplier; + private final Object lock = new Object(); + private FlatVectorsReader mergeReader; + private boolean closed; - public MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) { + public MergeReaderWrapper(FlatVectorsReader mainReader, IOSupplier mergeReaderSupplier) { this.mainReader = mainReader; - this.mergeReader = mergeReader; + this.mergeReaderSupplier = mergeReaderSupplier; } @Override @@ -75,8 +85,20 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep } @Override - public FlatVectorsReader getMergeInstance() { - return mergeReader; + public FlatVectorsReader getMergeInstance() throws IOException { + synchronized (lock) { + if (closed) { + throw new AlreadyClosedException("this MergeReaderWrapper is closed"); + } + // created lazily: most segments are never merged during a reader's lifetime, and the + // merge reader holds direct I/O resources + if (mergeReader == null) { + mergeReader = mergeReaderSupplier.get(); + } + // delegate so the reader can prepare itself for merging, e.g. Lucene99FlatVectorsReader + // switches its data input to sequential read advice + return mergeReader.getMergeInstance(); + } } @Override @@ -98,6 +120,12 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { @Override public void close() throws IOException { - IOUtils.close(mainReader, mergeReader); + synchronized (lock) { + if (closed) { + return; + } + closed = true; + IOUtils.close(mainReader, mergeReader); + } } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/DirectIOCapableLucene99FlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/DirectIOCapableLucene99FlatVectorsFormat.java index 06d303b8dc97e..ba7aa5398d4fc 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/DirectIOCapableLucene99FlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/es93/DirectIOCapableLucene99FlatVectorsFormat.java @@ -15,9 +15,7 @@ import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IOContext; import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat; -import org.elasticsearch.index.codec.vectors.MergeReaderWrapper; import java.io.IOException; @@ -47,25 +45,4 @@ protected FlatVectorsReader createReader(SegmentReadState state) throws IOExcept public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { return new Lucene99FlatVectorsWriter(state, vectorsScorer); } - - @Override - public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException { - if (state.context.context() == IOContext.Context.DEFAULT && useDirectIO && canUseDirectIO(state)) { - // only override the context for the random-access use case - SegmentReadState directIOState = new SegmentReadState( - state.directory, - state.segmentInfo, - state.fieldInfos, - new DirectIOContext(state.context.hints()), - state.segmentSuffix - ); - // Use mmap for merges and direct I/O for searches. - return new MergeReaderWrapper( - new Lucene99FlatVectorsReader(directIOState, vectorsScorer), - new Lucene99FlatVectorsReader(state, vectorsScorer) - ); - } else { - return new Lucene99FlatVectorsReader(state, vectorsScorer); - } - } } diff --git a/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java b/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java index 5af1ddaf5cc44..e39c8341eb4b3 100644 --- a/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java +++ b/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java @@ -401,6 +401,10 @@ public IndexInput slice(String sliceDescription, long offset, long length) throw } // pkg private for testing + int bufferCapacity() { + return buffer.capacity(); + } + int prefetchSlots() { return prefetcher.posToSlot.size(); } @@ -448,6 +452,10 @@ private static class DirectIOPrefetcher implements Closeable { */ void prefetch(long pos, long length) { assert pos % blockSize == 0 : "prefetch pos [" + pos + "] must be aligned to block size [" + blockSize + "]"; + if (maxConcurrentPrefetches == 0) { + // opened without prefetch slots, e.g. a merge-hinted input when async prefetch is disabled + return; + } // first determine how many slots we need given the length while (length > 0) { Map.Entry floor = this.posToSlot.floorEntry(pos); diff --git a/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java index 792d54fb82b4c..0eefdbb5bfa91 100644 --- a/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java +++ b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java @@ -30,6 +30,7 @@ import org.elasticsearch.index.IndexModule; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.StandardIOBehaviorHint; +import org.elasticsearch.index.codec.vectors.DirectIOMergeHint; import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; import org.elasticsearch.index.shard.ShardPath; import org.elasticsearch.logging.LogManager; @@ -184,8 +185,12 @@ public HybridDirectory(LockFactory lockFactory, MMapDirectory delegate, int asyn DirectIODirectory directIO; try { - // use 8kB buffer (two pages) to guarantee it can load all of an un-page-aligned 1024-dim float vector - directIO = new AlwaysDirectIODirectory(delegate, 8192, DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, asyncPrefetchLimit); + directIO = new AlwaysDirectIODirectory( + delegate, + AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, + DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, + asyncPrefetchLimit + ); } catch (Exception e) { // directio not supported Log.warn("Could not initialize DirectIO access", e); @@ -203,7 +208,7 @@ public IndexInput openInput(String name, IOContext context) throws IOException { try { Log.debug("Opening {} with direct IO", name); return directIODelegate.openInput(name, context); - } catch (FileSystemException e) { + } catch (FileSystemException | UnsupportedOperationException e) { Log.debug(() -> Strings.format("Could not open %s with direct IO", name), e); directIOException = e; // and fallthrough to normal opening below @@ -306,6 +311,9 @@ MMapDirectory getDelegate() { } public static final class AlwaysDirectIODirectory extends DirectIODirectory { + // two pages, guaranteeing a single buffer can load all of an un-page-aligned 1024-dim float vector + public static final int RANDOM_ACCESS_BUFFER_SIZE = 8192; + private final int blockSize; private final int asyncPrefetchLimit; @@ -324,8 +332,18 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL @Override public IndexInput openInput(String name, IOContext context) throws IOException { ensureOpen(); - if (asyncPrefetchLimit > 0) { - return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, 8192, asyncPrefetchLimit); + // merge reads are long sequential streams, so they use a merge-sized buffer rather + // than the two-page buffer that random-access rescore reads use. They get no prefetch + // slots: the large buffer is their readahead, and per-slot prefetch buffers are sized + // by the read buffer, so slots would multiply the prefetch-limit setting's documented + // memory bound by the merge buffer size. + boolean forMerge = context.hints().contains(DirectIOMergeHint.INSTANCE); + int bufferSize = forMerge ? DEFAULT_MERGE_BUFFER_SIZE : RANDOM_ACCESS_BUFFER_SIZE; + int prefetchLimit = forMerge ? 0 : asyncPrefetchLimit; + if (asyncPrefetchLimit > 0 || forMerge) { + // merge-hinted opens always take this path, since the delegate's buffer size is + // fixed at construction and cannot honor the merge hint + return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, bufferSize, prefetchLimit); } else { return super.openInput(name, context); } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java new file mode 100644 index 0000000000000..d66121621649d --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java @@ -0,0 +1,222 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.index.codec.vectors; + +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +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.index.VectorSimilarityFunction; +import org.apache.lucene.misc.store.DirectIODirectory; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.store.MMapDirectory; +import org.apache.lucene.store.NativeFSLockFactory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; +import org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; +import org.elasticsearch.index.codec.vectors.es94.ES94HnswScalarQuantizedVectorsFormat; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; +import org.elasticsearch.index.store.FsDirectoryFactory; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +/** + * Tests that formats based on {@link DirectIOCapableFlatVectorsFormat} open the raw vector data + * with direct I/O for both searches and merges when direct I/O is requested. + */ +public class DirectIOCapableFlatVectorsFormatTests extends LuceneTestCase { + + static { + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + @BeforeClass + public static void checkDirectIOSupported() throws IOException { + Path path = createTempDir("directIOProbe"); + try ( + Directory dir = new FsDirectoryFactory.AlwaysDirectIODirectory( + new MMapDirectory(path), + FsDirectoryFactory.AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, + DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, + 0 + ) + ) { + try (IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) { + out.writeString("test"); + } + try (IndexInput in = dir.openInput("out", IOContext.DEFAULT)) { + assertEquals("test", in.readString()); + } + } catch (IOException | UnsupportedOperationException e) { + assumeNoException("test requires a JDK and filesystem that support Direct IO", e); + } + } + + /** Records the IOContext of every open of a raw vector file. */ + private record VecOpen(String name, IOContext.Context context, boolean directIO, boolean mergeDirectIO) {} + + private static class VecOpenRecordingDirectory extends FilterDirectory { + final List vecOpens = new CopyOnWriteArrayList<>(); + + VecOpenRecordingDirectory(Directory in) { + super(in); + } + + @Override + public IndexInput openInput(String name, IOContext context) throws IOException { + if (name.endsWith(".vec")) { + vecOpens.add( + new VecOpen( + name, + context.context(), + context.hints().contains(DirectIOHint.INSTANCE), + context.hints().contains(DirectIOMergeHint.INSTANCE) + ) + ); + } + return super.openInput(name, context); + } + } + + public void testInt8HnswOpensRawVectorsWithDirectIO() throws IOException { + // the scalar-quantized reader chain does not currently propagate getMergeInstance down to + // the raw reader, so only open-time direct I/O behavior can be asserted for this format + int maxConn = 16; + int beamWidth = 100; + int bits = 7; + boolean useDirectIO = true; + runMergeTest( + new ES94HnswScalarQuantizedVectorsFormat(maxConn, beamWidth, DenseVectorFieldMapper.ElementType.FLOAT, bits, useDirectIO), + false + ); + } + + public void testBbqHnswMergeReaderUsesMergeSizedDirectIO() throws IOException { + // the binary-quantized reader chain propagates getMergeInstance to the raw reader, so the + // merge triggers the lazy creation of the merge reader, which must open the raw vectors + // with the merge hint + runMergeTest(new ES93HnswBinaryQuantizedVectorsFormat(DenseVectorFieldMapper.ElementType.FLOAT, true), true); + } + + private void runMergeTest(KnnVectorsFormat format, boolean expectMergeHintedOpen) throws IOException { + int dims = 64; + int docsPerSegment = 50; + float[][] vectors = new float[docsPerSegment * 2][]; + for (int i = 0; i < vectors.length; i++) { + vectors[i] = randomVector(dims); + } + + Path path = createTempDir("directIOMerge"); + IndexWriterConfig config = new IndexWriterConfig().setCodec(TestUtil.alwaysKnnVectorsFormat(format)); + // direct I/O only applies to non-compound segments; compound files would also hide + // the raw vector file opens behind the .cfs + config.setUseCompoundFile(false); + config.getMergePolicy().setNoCFSRatio(0.0); + + try ( + VecOpenRecordingDirectory dir = new VecOpenRecordingDirectory( + new FsDirectoryFactory.HybridDirectory(NativeFSLockFactory.INSTANCE, new MMapDirectory(path), 64) + ); + IndexWriter writer = new IndexWriter(dir, config) + ) { + for (int i = 0; i < vectors.length; i++) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField("v", vectors[i], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + if (i == docsPerSegment - 1) { + writer.commit(); // ensure at least two segments exist, so forceMerge does real work + } + } + writer.commit(); + + // hold a reader open across the merge, so that the merge reads through the pooled + // DEFAULT-context readers (as on a node serving searches) rather than readers opened + // with a MERGE context, which do not use direct I/O + try (DirectoryReader beforeMerge = DirectoryReader.open(writer)) { + assertEquals(vectors.length, beforeMerge.numDocs()); + writer.forceMerge(1); + } + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + LeafReader leafReader = getOnlyLeafReader(reader); + FloatVectorValues values = leafReader.getFloatVectorValues("v"); + KnnVectorValues.DocIndexIterator iterator = values.iterator(); + int count = 0; + while (iterator.nextDoc() != NO_MORE_DOCS) { + assertTrue(containsVector(vectors, values.vectorValue(iterator.index()))); + count++; + } + assertEquals(vectors.length, count); + + TopDocs topDocs = new IndexSearcher(reader).search(new KnnFloatVectorQuery("v", vectors[0], 5), 5); + assertEquals(5, topDocs.scoreDocs.length); + } + + assertTrue("expected at least one open of a raw vector file", dir.vecOpens.isEmpty() == false); + for (VecOpen open : dir.vecOpens) { + if (open.context() == IOContext.Context.DEFAULT) { + assertTrue("raw vector file [" + open.name() + "] was opened without requesting direct IO", open.directIO()); + } + } + if (expectMergeHintedOpen) { + assertTrue( + "expected the merge to open a raw vector file with a merge-hinted direct IO context", + dir.vecOpens.stream().anyMatch(VecOpen::mergeDirectIO) + ); + } else { + // documents the current gap: this chain does not propagate getMergeInstance, so no + // merge-hinted open may occur; this fails loudly when propagation lands + assertTrue( + "unexpected merge-hinted open for a chain that does not propagate getMergeInstance", + dir.vecOpens.stream().noneMatch(VecOpen::mergeDirectIO) + ); + } + } + } + + private static boolean containsVector(float[][] vectors, float[] candidate) { + for (float[] vector : vectors) { + if (Arrays.equals(vector, candidate)) { + return true; + } + } + return false; + } + + private static float[] randomVector(int dims) { + float[] vector = new float[dims]; + for (int i = 0; i < dims; i++) { + vector[i] = random().nextFloat(); + } + return vector; + } +} diff --git a/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java b/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java index 8189d834f727b..28e69644491e6 100644 --- a/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java +++ b/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java @@ -9,11 +9,15 @@ package org.elasticsearch.index.store; +import org.apache.lucene.misc.store.DirectIODirectory; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.MMapDirectory; import org.apache.lucene.store.NIOFSDirectory; import org.apache.lucene.util.hnsw.IntToIntFunction; import org.elasticsearch.core.SuppressForbidden; +import org.elasticsearch.index.codec.vectors.DirectIOMergeHint; +import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; import org.elasticsearch.test.ESTestCase; import java.io.IOException; @@ -71,6 +75,66 @@ public void testPrefetchEdgeCase() throws IOException { } } + public void testPrefetchWithNoPrefetchSlots() throws IOException { + byte[] bytes = new byte[BASE_BUFFER_SIZE * 4 + randomIntBetween(1, BASE_BUFFER_SIZE)]; + random().nextBytes(bytes); + Path path = createTempDir("testDirectIODirectory"); + int blockSize = getBlockSize(path); + try (Directory dir = new NIOFSDirectory(path)) { + try (var output = dir.createOutput("test", org.apache.lucene.store.IOContext.DEFAULT)) { + output.writeBytes(bytes, bytes.length); + } + try (var input = new AsyncDirectIOIndexInput(path.resolve("test"), blockSize, BASE_BUFFER_SIZE, 0)) { + assertEquals(0, input.prefetchSlots()); + // prefetch must be a no-op, not a failure, when the input has no prefetch slots + input.prefetch(0, bytes.length); + byte[] read = new byte[bytes.length]; + input.readBytes(read, 0, read.length); + assertArrayEquals(bytes, read); + } + } + } + + public void testMergeHintedOpenUsesMergeBufferAndNoPrefetchSlots() throws IOException { + Path path = createTempDir("testDirectIODirectory"); + FsDirectoryFactory.AlwaysDirectIODirectory dir; + try { + dir = new FsDirectoryFactory.AlwaysDirectIODirectory( + new MMapDirectory(path), + FsDirectoryFactory.AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, + DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, + 4 + ); + } catch (Exception e) { + assumeNoException("test requires a JDK and filesystem that support Direct IO", e); + return; + } + try (dir) { + try (var output = dir.createOutput("test", org.apache.lucene.store.IOContext.DEFAULT)) { + output.writeString("test"); + } + try (var input = dir.openInput("test", org.apache.lucene.store.IOContext.DEFAULT.withHints(DirectIOHint.INSTANCE))) { + AsyncDirectIOIndexInput asyncInput = (AsyncDirectIOIndexInput) input; + assertEquals(FsDirectoryFactory.AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, asyncInput.bufferCapacity()); + // prefetchSlots() reports live registrations: one prefetch call occupies a slot + asyncInput.prefetch(0, asyncInput.length()); + assertEquals(1, asyncInput.prefetchSlots()); + } + try ( + var input = dir.openInput( + "test", + org.apache.lucene.store.IOContext.DEFAULT.withHints(DirectIOHint.INSTANCE, DirectIOMergeHint.INSTANCE) + ) + ) { + AsyncDirectIOIndexInput asyncInput = (AsyncDirectIOIndexInput) input; + assertEquals(DirectIODirectory.DEFAULT_MERGE_BUFFER_SIZE, asyncInput.bufferCapacity()); + // merge-hinted inputs get no prefetch slots: prefetch is a no-op and registers nothing + asyncInput.prefetch(0, asyncInput.length()); + assertEquals(0, asyncInput.prefetchSlots()); + } + } + } + public void testLargePrefetch() throws IOException { byte[] bytes = new byte[BASE_BUFFER_SIZE * 10 + randomIntBetween(1, BASE_BUFFER_SIZE)]; int offset = randomIntBetween(1, BASE_BUFFER_SIZE); From ae005d11f76b99b76b002692b5db717e489a065b Mon Sep 17 00:00:00 2001 From: John Maurice Date: Wed, 5 Aug 2026 05:43:22 +0300 Subject: [PATCH 2/3] Add changelog --- docs/changelog/155919.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 docs/changelog/155919.yaml diff --git a/docs/changelog/155919.yaml b/docs/changelog/155919.yaml new file mode 100644 index 0000000000000..ee96aa91c85d0 --- /dev/null +++ b/docs/changelog/155919.yaml @@ -0,0 +1,6 @@ +pr: 155919 +summary: Use direct I/O for merge-time reads of raw vectors +area: Vector Search +type: enhancement +issues: + - 155021 From 202775c438bda6a55a0d40095650169ecdb4a101 Mon Sep 17 00:00:00 2001 From: John Maurice Date: Wed, 5 Aug 2026 08:49:26 +0000 Subject: [PATCH 3/3] Route merge reads by IOContext.MERGE instead of a merge hint --- .../mapping-reference/dense-vector.md | 2 +- .../DirectIOCapableFlatVectorsFormat.java | 20 +++--- .../codec/vectors/DirectIOMergeHint.java | 20 ------ .../codec/vectors/MergeReaderWrapper.java | 34 +++++----- .../index/store/AsyncDirectIOIndexInput.java | 8 --- .../index/store/FsDirectoryFactory.java | 36 ++++++----- ...DirectIOCapableFlatVectorsFormatTests.java | 62 +++--------------- .../store/AsyncDirectIOIndexInputTests.java | 64 ------------------- 8 files changed, 55 insertions(+), 191 deletions(-) delete mode 100644 server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java diff --git a/docs/reference/elasticsearch/mapping-reference/dense-vector.md b/docs/reference/elasticsearch/mapping-reference/dense-vector.md index 6e99968a31f60..532a1cbb15691 100644 --- a/docs/reference/elasticsearch/mapping-reference/dense-vector.md +++ b/docs/reference/elasticsearch/mapping-reference/dense-vector.md @@ -602,7 +602,7 @@ $$$dense-vector-index-options$$$ :::: `on_disk_rescore` {applies_to}`stack: preview 9.3` {applies_to}`serverless: unavailable` -: (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. For `bbq_hnsw` indices, segment merges also read the raw vector data directly from disk where the platform supports it, so that merging is less likely to evict more frequently accessed data from the filesystem cache. {applies_to}`stack: ga 9.6` This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`. +: (Optional, boolean) Only applicable to quantized HNSW and `bbq_disk` index types. When `true`, vector rescoring and merging will read the raw vector data directly from disk, and will not copy it in memory. This can improve performance when vector data is larger than the amount of available RAM. This setting only applies to newly-indexed vectors; after changing this setting, the vectors must be reindexed or force-merged to apply the new setting to the whole index. Defaults to `false`. `auto_calibrate` {applies_to}`stack: ga 9.5` : (Optional, boolean) Only applicable to `bbq_disk`. When `true`, {{es}} automatically selects the optimal quantization encoding, oversampling factor, and preconditioning for each merged segment based on the actual recall characteristics of the merged corpus. Segments containing fewer than 10,000 vectors after merging are not calibrated and, when not otherwise specified in the mappings, use the default oversampling factor of 3.0x. Defaults to `false`. Cannot be changed after the field is created. Refer to [Auto-calibration for `bbq_disk`](/reference/elasticsearch/mapping-reference/bbq.md#bbq-auto-calibration) for details. diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java index 0ab03af7c8dce..940e919505232 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java @@ -52,11 +52,11 @@ public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectI state.directory, state.segmentInfo, state.fieldInfos, - new DirectIOContext(state.context.hints(), Set.of(DirectIOHint.INSTANCE, DirectIOMergeHint.INSTANCE)), + new DirectIOContext(IOContext.Context.MERGE, state.context.hints()), state.segmentSuffix ); // Use direct I/O for merges too, so merge reads do not evict hotter data from the - // page cache. The merge hint makes those reads use a merge-sized buffer. + // page cache. The MERGE context routes those reads to a merge-sized buffer. return new MergeReaderWrapper(createReader(directIOState), () -> createReader(mergeDirectIOState)); } else { return createReader(state); @@ -65,22 +65,22 @@ public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectI protected static class DirectIOContext implements IOContext { - private final Set stickyHints; + private final Context context; final Set hints; public DirectIOContext(Set hints) { - this(hints, Set.of(DirectIOHint.INSTANCE)); + this(Context.DEFAULT, hints); } - public DirectIOContext(Set hints, Set stickyHints) { - // the sticky hints are always added, and survive withHints() replacing the others - this.stickyHints = stickyHints; - this.hints = Sets.union(hints, stickyHints); + public DirectIOContext(Context context, Set hints) { + this.context = context; + // always add DirectIOHint to the hints given + this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE)); } @Override public Context context() { - return Context.DEFAULT; + return context; } @Override @@ -100,7 +100,7 @@ public Set hints() { @Override public IOContext withHints(FileOpenHint... hints) { - return new DirectIOContext(Set.of(hints), stickyHints); + return new DirectIOContext(context, Set.of(hints)); } } } diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java deleted file mode 100644 index a60cfc4c36a8d..0000000000000 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOMergeHint.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the "Elastic License - * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side - * Public License v 1"; you may not use this file except in compliance with, at - * your election, the "Elastic License 2.0", the "GNU Affero General Public - * License v3.0 only", or the "Server Side Public License, v 1". - */ - -package org.elasticsearch.index.codec.vectors; - -import org.apache.lucene.store.IOContext; - -/** - * Hint that a file opened with direct I/O will be read as a long sequential stream by a merge, - * so reads should use a merge-sized buffer rather than the small buffer used for random access. - */ -public enum DirectIOMergeHint implements IOContext.FileOpenHint { - INSTANCE -} diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java index d4e4e84cea005..e7fa39e57e882 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/MergeReaderWrapper.java @@ -35,7 +35,6 @@ public class MergeReaderWrapper extends FlatVectorsReader { private final FlatVectorsReader mainReader; private final IOSupplier mergeReaderSupplier; - private final Object lock = new Object(); private FlatVectorsReader mergeReader; private boolean closed; @@ -86,19 +85,18 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep @Override public FlatVectorsReader getMergeInstance() throws IOException { - synchronized (lock) { - if (closed) { - throw new AlreadyClosedException("this MergeReaderWrapper is closed"); - } - // created lazily: most segments are never merged during a reader's lifetime, and the - // merge reader holds direct I/O resources - if (mergeReader == null) { - mergeReader = mergeReaderSupplier.get(); - } - // delegate so the reader can prepare itself for merging, e.g. Lucene99FlatVectorsReader - // switches its data input to sequential read advice - return mergeReader.getMergeInstance(); + // only the single thread running the merge calls this + if (closed) { + throw new AlreadyClosedException("this MergeReaderWrapper is closed"); } + // created lazily: most segments are never merged during a reader's lifetime, and the + // merge reader holds direct I/O resources + if (mergeReader == null) { + mergeReader = mergeReaderSupplier.get(); + } + // delegate so the reader can prepare itself for merging, e.g. Lucene99FlatVectorsReader + // switches its data input to sequential read advice + return mergeReader.getMergeInstance(); } @Override @@ -120,12 +118,10 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { @Override public void close() throws IOException { - synchronized (lock) { - if (closed) { - return; - } - closed = true; - IOUtils.close(mainReader, mergeReader); + if (closed) { + return; } + closed = true; + IOUtils.close(mainReader, mergeReader); } } diff --git a/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java b/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java index e39c8341eb4b3..5af1ddaf5cc44 100644 --- a/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java +++ b/server/src/main/java/org/elasticsearch/index/store/AsyncDirectIOIndexInput.java @@ -401,10 +401,6 @@ public IndexInput slice(String sliceDescription, long offset, long length) throw } // pkg private for testing - int bufferCapacity() { - return buffer.capacity(); - } - int prefetchSlots() { return prefetcher.posToSlot.size(); } @@ -452,10 +448,6 @@ private static class DirectIOPrefetcher implements Closeable { */ void prefetch(long pos, long length) { assert pos % blockSize == 0 : "prefetch pos [" + pos + "] must be aligned to block size [" + blockSize + "]"; - if (maxConcurrentPrefetches == 0) { - // opened without prefetch slots, e.g. a merge-hinted input when async prefetch is disabled - return; - } // first determine how many slots we need given the length while (length > 0) { Map.Entry floor = this.posToSlot.floorEntry(pos); diff --git a/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java index 0eefdbb5bfa91..b1b6fe15364ed 100644 --- a/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java +++ b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java @@ -30,7 +30,6 @@ import org.elasticsearch.index.IndexModule; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.StandardIOBehaviorHint; -import org.elasticsearch.index.codec.vectors.DirectIOMergeHint; import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; import org.elasticsearch.index.shard.ShardPath; import org.elasticsearch.logging.LogManager; @@ -178,12 +177,14 @@ private static int getBlockSize(Path path) throws IOException { public static final class HybridDirectory extends NIOFSDirectory { private final MMapDirectory delegate; private final DirectIODirectory directIODelegate; + private final DirectIODirectory mergeDirectIODelegate; public HybridDirectory(LockFactory lockFactory, MMapDirectory delegate, int asyncPrefetchLimit) throws IOException { super(delegate.getDirectory(), lockFactory); this.delegate = delegate; DirectIODirectory directIO; + DirectIODirectory mergeDirectIO; try { directIO = new AlwaysDirectIODirectory( delegate, @@ -191,12 +192,21 @@ public HybridDirectory(LockFactory lockFactory, MMapDirectory delegate, int asyn DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, asyncPrefetchLimit ); + // merge reads are long sequential streams: they get a merge-sized buffer and no async prefetch + mergeDirectIO = new AlwaysDirectIODirectory( + delegate, + DirectIODirectory.DEFAULT_MERGE_BUFFER_SIZE, + DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, + 0 + ); } catch (Exception e) { // directio not supported Log.warn("Could not initialize DirectIO access", e); directIO = null; + mergeDirectIO = null; } this.directIODelegate = directIO; + this.mergeDirectIODelegate = mergeDirectIO; } @Override @@ -207,7 +217,8 @@ public IndexInput openInput(String name, IOContext context) throws IOException { ensureCanRead(name); try { Log.debug("Opening {} with direct IO", name); - return directIODelegate.openInput(name, context); + DirectIODirectory dio = context.context() == IOContext.Context.MERGE ? mergeDirectIODelegate : directIODelegate; + return dio.openInput(name, context); } catch (FileSystemException | UnsupportedOperationException e) { Log.debug(() -> Strings.format("Could not open %s with direct IO", name), e); directIOException = e; @@ -315,12 +326,14 @@ public static final class AlwaysDirectIODirectory extends DirectIODirectory { public static final int RANDOM_ACCESS_BUFFER_SIZE = 8192; private final int blockSize; + private final int bufferSize; private final int asyncPrefetchLimit; - public AlwaysDirectIODirectory(FSDirectory delegate, int mergeBufferSize, long minBytesDirect, int asyncPrefetchLimit) + public AlwaysDirectIODirectory(FSDirectory delegate, int bufferSize, long minBytesDirect, int asyncPrefetchLimit) throws IOException { - super(delegate, mergeBufferSize, minBytesDirect); + super(delegate, bufferSize, minBytesDirect); blockSize = getBlockSize(delegate.getDirectory()); + this.bufferSize = bufferSize; this.asyncPrefetchLimit = asyncPrefetchLimit; } @@ -332,19 +345,10 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL @Override public IndexInput openInput(String name, IOContext context) throws IOException { ensureOpen(); - // merge reads are long sequential streams, so they use a merge-sized buffer rather - // than the two-page buffer that random-access rescore reads use. They get no prefetch - // slots: the large buffer is their readahead, and per-slot prefetch buffers are sized - // by the read buffer, so slots would multiply the prefetch-limit setting's documented - // memory bound by the merge buffer size. - boolean forMerge = context.hints().contains(DirectIOMergeHint.INSTANCE); - int bufferSize = forMerge ? DEFAULT_MERGE_BUFFER_SIZE : RANDOM_ACCESS_BUFFER_SIZE; - int prefetchLimit = forMerge ? 0 : asyncPrefetchLimit; - if (asyncPrefetchLimit > 0 || forMerge) { - // merge-hinted opens always take this path, since the delegate's buffer size is - // fixed at construction and cannot honor the merge hint - return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, bufferSize, prefetchLimit); + if (asyncPrefetchLimit > 0) { + return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, bufferSize, asyncPrefetchLimit); } else { + // no async prefetching: a plain direct-IO input at this instance's buffer size return super.openInput(name, context); } } diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java index d66121621649d..81fcc0168f380 100644 --- a/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java @@ -13,16 +13,10 @@ import org.apache.lucene.document.Document; import org.apache.lucene.document.KnnFloatVectorField; 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.index.VectorSimilarityFunction; import org.apache.lucene.misc.store.DirectIODirectory; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FilterDirectory; import org.apache.lucene.store.IOContext; @@ -30,33 +24,26 @@ import org.apache.lucene.store.IndexOutput; import org.apache.lucene.store.MMapDirectory; import org.apache.lucene.store.NativeFSLockFactory; -import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; import org.apache.lucene.tests.util.TestUtil; -import org.elasticsearch.common.logging.LogConfigurator; import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; import org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat; import org.elasticsearch.index.codec.vectors.es94.ES94HnswScalarQuantizedVectorsFormat; import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; import org.elasticsearch.index.store.FsDirectoryFactory; +import org.elasticsearch.test.ESTestCase; import org.junit.BeforeClass; import java.io.IOException; import java.nio.file.Path; -import java.util.Arrays; import java.util.List; import java.util.concurrent.CopyOnWriteArrayList; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - /** * Tests that formats based on {@link DirectIOCapableFlatVectorsFormat} open the raw vector data * with direct I/O for both searches and merges when direct I/O is requested. */ -public class DirectIOCapableFlatVectorsFormatTests extends LuceneTestCase { - - static { - LogConfigurator.configureESLogging(); // native access requires logging to be initialized - } +public class DirectIOCapableFlatVectorsFormatTests extends ESTestCase { @BeforeClass public static void checkDirectIOSupported() throws IOException { @@ -98,7 +85,7 @@ public IndexInput openInput(String name, IOContext context) throws IOException { name, context.context(), context.hints().contains(DirectIOHint.INSTANCE), - context.hints().contains(DirectIOMergeHint.INSTANCE) + context.context() == IOContext.Context.MERGE && context.hints().contains(DirectIOHint.INSTANCE) ) ); } @@ -128,10 +115,10 @@ public void testBbqHnswMergeReaderUsesMergeSizedDirectIO() throws IOException { private void runMergeTest(KnnVectorsFormat format, boolean expectMergeHintedOpen) throws IOException { int dims = 64; - int docsPerSegment = 50; + int docsPerSegment = randomIntBetween(30, 120); float[][] vectors = new float[docsPerSegment * 2][]; for (int i = 0; i < vectors.length; i++) { - vectors[i] = randomVector(dims); + vectors[i] = BaseKnnVectorsFormatTestCase.randomNormalizedVector(dims); } Path path = createTempDir("directIOMerge"); @@ -166,21 +153,6 @@ private void runMergeTest(KnnVectorsFormat format, boolean expectMergeHintedOpen } writer.commit(); - try (DirectoryReader reader = DirectoryReader.open(writer)) { - LeafReader leafReader = getOnlyLeafReader(reader); - FloatVectorValues values = leafReader.getFloatVectorValues("v"); - KnnVectorValues.DocIndexIterator iterator = values.iterator(); - int count = 0; - while (iterator.nextDoc() != NO_MORE_DOCS) { - assertTrue(containsVector(vectors, values.vectorValue(iterator.index()))); - count++; - } - assertEquals(vectors.length, count); - - TopDocs topDocs = new IndexSearcher(reader).search(new KnnFloatVectorQuery("v", vectors[0], 5), 5); - assertEquals(5, topDocs.scoreDocs.length); - } - assertTrue("expected at least one open of a raw vector file", dir.vecOpens.isEmpty() == false); for (VecOpen open : dir.vecOpens) { if (open.context() == IOContext.Context.DEFAULT) { @@ -189,34 +161,18 @@ private void runMergeTest(KnnVectorsFormat format, boolean expectMergeHintedOpen } if (expectMergeHintedOpen) { assertTrue( - "expected the merge to open a raw vector file with a merge-hinted direct IO context", + "expected the merge to open a raw vector file with a MERGE-context direct IO open", dir.vecOpens.stream().anyMatch(VecOpen::mergeDirectIO) ); } else { // documents the current gap: this chain does not propagate getMergeInstance, so no - // merge-hinted open may occur; this fails loudly when propagation lands + // MERGE-context open may occur; this fails loudly when propagation lands assertTrue( - "unexpected merge-hinted open for a chain that does not propagate getMergeInstance", + "unexpected MERGE-context open for a chain that does not propagate getMergeInstance", dir.vecOpens.stream().noneMatch(VecOpen::mergeDirectIO) ); } } } - private static boolean containsVector(float[][] vectors, float[] candidate) { - for (float[] vector : vectors) { - if (Arrays.equals(vector, candidate)) { - return true; - } - } - return false; - } - - private static float[] randomVector(int dims) { - float[] vector = new float[dims]; - for (int i = 0; i < dims; i++) { - vector[i] = random().nextFloat(); - } - return vector; - } } diff --git a/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java b/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java index 28e69644491e6..8189d834f727b 100644 --- a/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java +++ b/server/src/test/java/org/elasticsearch/index/store/AsyncDirectIOIndexInputTests.java @@ -9,15 +9,11 @@ package org.elasticsearch.index.store; -import org.apache.lucene.misc.store.DirectIODirectory; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.MMapDirectory; import org.apache.lucene.store.NIOFSDirectory; import org.apache.lucene.util.hnsw.IntToIntFunction; import org.elasticsearch.core.SuppressForbidden; -import org.elasticsearch.index.codec.vectors.DirectIOMergeHint; -import org.elasticsearch.index.codec.vectors.es818.DirectIOHint; import org.elasticsearch.test.ESTestCase; import java.io.IOException; @@ -75,66 +71,6 @@ public void testPrefetchEdgeCase() throws IOException { } } - public void testPrefetchWithNoPrefetchSlots() throws IOException { - byte[] bytes = new byte[BASE_BUFFER_SIZE * 4 + randomIntBetween(1, BASE_BUFFER_SIZE)]; - random().nextBytes(bytes); - Path path = createTempDir("testDirectIODirectory"); - int blockSize = getBlockSize(path); - try (Directory dir = new NIOFSDirectory(path)) { - try (var output = dir.createOutput("test", org.apache.lucene.store.IOContext.DEFAULT)) { - output.writeBytes(bytes, bytes.length); - } - try (var input = new AsyncDirectIOIndexInput(path.resolve("test"), blockSize, BASE_BUFFER_SIZE, 0)) { - assertEquals(0, input.prefetchSlots()); - // prefetch must be a no-op, not a failure, when the input has no prefetch slots - input.prefetch(0, bytes.length); - byte[] read = new byte[bytes.length]; - input.readBytes(read, 0, read.length); - assertArrayEquals(bytes, read); - } - } - } - - public void testMergeHintedOpenUsesMergeBufferAndNoPrefetchSlots() throws IOException { - Path path = createTempDir("testDirectIODirectory"); - FsDirectoryFactory.AlwaysDirectIODirectory dir; - try { - dir = new FsDirectoryFactory.AlwaysDirectIODirectory( - new MMapDirectory(path), - FsDirectoryFactory.AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, - DirectIODirectory.DEFAULT_MIN_BYTES_DIRECT, - 4 - ); - } catch (Exception e) { - assumeNoException("test requires a JDK and filesystem that support Direct IO", e); - return; - } - try (dir) { - try (var output = dir.createOutput("test", org.apache.lucene.store.IOContext.DEFAULT)) { - output.writeString("test"); - } - try (var input = dir.openInput("test", org.apache.lucene.store.IOContext.DEFAULT.withHints(DirectIOHint.INSTANCE))) { - AsyncDirectIOIndexInput asyncInput = (AsyncDirectIOIndexInput) input; - assertEquals(FsDirectoryFactory.AlwaysDirectIODirectory.RANDOM_ACCESS_BUFFER_SIZE, asyncInput.bufferCapacity()); - // prefetchSlots() reports live registrations: one prefetch call occupies a slot - asyncInput.prefetch(0, asyncInput.length()); - assertEquals(1, asyncInput.prefetchSlots()); - } - try ( - var input = dir.openInput( - "test", - org.apache.lucene.store.IOContext.DEFAULT.withHints(DirectIOHint.INSTANCE, DirectIOMergeHint.INSTANCE) - ) - ) { - AsyncDirectIOIndexInput asyncInput = (AsyncDirectIOIndexInput) input; - assertEquals(DirectIODirectory.DEFAULT_MERGE_BUFFER_SIZE, asyncInput.bufferCapacity()); - // merge-hinted inputs get no prefetch slots: prefetch is a no-op and registers nothing - asyncInput.prefetch(0, asyncInput.length()); - assertEquals(0, asyncInput.prefetchSlots()); - } - } - } - public void testLargePrefetch() throws IOException { byte[] bytes = new byte[BASE_BUFFER_SIZE * 10 + randomIntBetween(1, BASE_BUFFER_SIZE)]; int offset = randomIntBetween(1, BASE_BUFFER_SIZE);