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 diff --git a/docs/reference/elasticsearch/mapping-reference/dense-vector.md b/docs/reference/elasticsearch/mapping-reference/dense-vector.md index ac362145c86ca..1ceabad90d993 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 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 fade2384441cf..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 @@ -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(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 context routes those reads to a merge-sized buffer. + return new MergeReaderWrapper(createReader(directIOState), () -> createReader(mergeDirectIOState)); } else { return createReader(state); } @@ -56,16 +65,22 @@ public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectI protected static class DirectIOContext implements IOContext { + private final Context context; final Set hints; public DirectIOContext(Set hints) { + this(Context.DEFAULT, hints); + } + + 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 @@ -85,7 +100,7 @@ public Set hints() { @Override public IOContext withHints(FileOpenHint... hints) { - return new DirectIOContext(Set.of(hints)); + return new DirectIOContext(context, Set.of(hints)); } } } 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..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 @@ -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,21 @@ 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 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 +84,19 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Accep } @Override - public FlatVectorsReader getMergeInstance() { - return mergeReader; + public FlatVectorsReader getMergeInstance() throws IOException { + // 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 @@ -98,6 +118,10 @@ public Map getOffHeapByteSize(FieldInfo fieldInfo) { @Override public void close() throws IOException { + 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/FsDirectoryFactory.java b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java index 792d54fb82b4c..b1b6fe15364ed 100644 --- a/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java +++ b/server/src/main/java/org/elasticsearch/index/store/FsDirectoryFactory.java @@ -177,21 +177,36 @@ 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 { - // 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 + ); + // 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 @@ -202,8 +217,9 @@ public IndexInput openInput(String name, IOContext context) throws IOException { ensureCanRead(name); try { Log.debug("Opening {} with direct IO", name); - return directIODelegate.openInput(name, context); - } catch (FileSystemException e) { + 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; // and fallthrough to normal opening below @@ -306,13 +322,18 @@ 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 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; } @@ -325,8 +346,9 @@ protected boolean useDirectIO(String name, IOContext context, OptionalLong fileL public IndexInput openInput(String name, IOContext context) throws IOException { ensureOpen(); if (asyncPrefetchLimit > 0) { - return new AsyncDirectIOIndexInput(getDirectory().resolve(name), blockSize, 8192, asyncPrefetchLimit); + 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 new file mode 100644 index 0000000000000..81fcc0168f380 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormatTests.java @@ -0,0 +1,178 @@ +/* + * 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.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.misc.store.DirectIODirectory; +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.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +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.List; +import java.util.concurrent.CopyOnWriteArrayList; + +/** + * 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 ESTestCase { + + @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.context() == IOContext.Context.MERGE && context.hints().contains(DirectIOHint.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 = randomIntBetween(30, 120); + float[][] vectors = new float[docsPerSegment * 2][]; + for (int i = 0; i < vectors.length; i++) { + vectors[i] = BaseKnnVectorsFormatTestCase.randomNormalizedVector(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(); + + 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-context direct IO open", + dir.vecOpens.stream().anyMatch(VecOpen::mergeDirectIO) + ); + } else { + // documents the current gap: this chain does not propagate getMergeInstance, so no + // MERGE-context open may occur; this fails loudly when propagation lands + assertTrue( + "unexpected MERGE-context open for a chain that does not propagate getMergeInstance", + dir.vecOpens.stream().noneMatch(VecOpen::mergeDirectIO) + ); + } + } + } + +}