diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 26729df99791..ad733266f16e 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -178,6 +178,9 @@ Optimizations Bug Fixes --------------------- +* GITHUB#16367: HNSW graph construction now periodically checks whether the surrounding merge has + been aborted, so IndexWriter#rollback and abortMerges no longer block until the entire graph is + built. (Jeho Jeong) * GITHUB#14049: Randomize KNN codec params in RandomCodec. Fixes scalar quantization div-by-zero when all values are identical. (Mike Sokolov) diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java index 1e330fee21b9..0b0bf5bda506 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java +++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java @@ -473,7 +473,8 @@ private void buildAndWriteGraph( mergeState.intraMergeTaskExecutor == null ? null : new TaskExecutor(mergeState.intraMergeTaskExecutor), - numMergeWorkers); + numMergeWorkers, + mergeState::checkAborted); for (int i = 0; i < mergeState.liveDocs.length; i++) { if (hasVectorValues(mergeState.fieldInfos[i], fieldInfo.name)) { merger.addReader( @@ -626,10 +627,11 @@ private HnswGraphMerger createGraphMerger( FieldInfo fieldInfo, RandomVectorScorerSupplier scorerSupplier, TaskExecutor parallelMergeTaskExecutor, - int numParallelMergeWorkers) { + int numParallelMergeWorkers, + IORunnable abortCheck) { if (mergeExec != null) { return new ConcurrentHnswMerger( - fieldInfo, scorerSupplier, M, beamWidth, mergeExec, numMergeWorkers); + fieldInfo, scorerSupplier, M, beamWidth, mergeExec, numMergeWorkers, abortCheck); } if (parallelMergeTaskExecutor != null && numParallelMergeWorkers > 1) { return new ConcurrentHnswMerger( @@ -638,9 +640,10 @@ private HnswGraphMerger createGraphMerger( M, beamWidth, parallelMergeTaskExecutor, - numParallelMergeWorkers); + numParallelMergeWorkers, + abortCheck); } - return new IncrementalHnswGraphMerger(fieldInfo, scorerSupplier, M, beamWidth); + return new IncrementalHnswGraphMerger(fieldInfo, scorerSupplier, M, beamWidth, abortCheck); } @Override diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java index 4b6244c18522..0803c1389ed2 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java @@ -29,6 +29,7 @@ import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.BitSet; import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.IORunnable; /** This merger merges graph in a concurrent manner, by using {@link HnswConcurrentMergeBuilder} */ public class ConcurrentHnswMerger extends IncrementalHnswGraphMerger { @@ -46,7 +47,24 @@ public ConcurrentHnswMerger( int beamWidth, TaskExecutor taskExecutor, int numWorker) { - super(fieldInfo, scorerSupplier, M, beamWidth); + this(fieldInfo, scorerSupplier, M, beamWidth, taskExecutor, numWorker, null); + } + + /** + * @param fieldInfo FieldInfo for the field being merged + * @param abortCheck optional check invoked before every node insertion during graph construction; + * may throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} to abort the + * build when the surrounding merge has been aborted, or null + */ + public ConcurrentHnswMerger( + FieldInfo fieldInfo, + RandomVectorScorerSupplier scorerSupplier, + int M, + int beamWidth, + TaskExecutor taskExecutor, + int numWorker, + IORunnable abortCheck) { + super(fieldInfo, scorerSupplier, M, beamWidth, abortCheck); this.taskExecutor = taskExecutor; this.numWorker = numWorker; } diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java index 38109c9c95e2..43cef2cde8ee 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java @@ -19,6 +19,7 @@ import java.io.IOException; import org.apache.lucene.internal.hppc.IntHashSet; +import org.apache.lucene.util.IORunnable; import org.apache.lucene.util.InfoStream; /** @@ -46,6 +47,18 @@ public interface HnswBuilder { /** Set info-stream to output debugging information */ void setInfoStream(InfoStream infoStream); + /** + * Sets a check that is invoked before every node insertion during graph construction. The check + * may throw an exception to abort the build promptly, e.g. {@link + * org.apache.lucene.index.MergePolicy.MergeAbortedException} when the merge that triggered the + * build has been aborted. + * + *

The check must be non-null and can be set at most once. + * + * @throws IllegalStateException if the check was already set + */ + void setAbortCheck(IORunnable abortCheck); + OnHeapHnswGraph getGraph(); /** diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java index f1a4666fd7e6..26fb7e420998 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java @@ -31,6 +31,7 @@ import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.BitSet; import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.IORunnable; import org.apache.lucene.util.InfoStream; /** @@ -134,6 +135,13 @@ public void setInfoStream(InfoStream infoStream) { } } + @Override + public void setAbortCheck(IORunnable abortCheck) { + for (HnswBuilder worker : workers) { + worker.setAbortCheck(abortCheck); + } + } + @Override public OnHeapHnswGraph getCompletedGraph() throws IOException { if (frozen == false) { diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java index e42b69ddb278..dbd92163f8d6 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java @@ -33,6 +33,7 @@ import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.knn.KnnSearchStrategy; import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.IORunnable; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswUtil.Component; @@ -83,6 +84,7 @@ public class HnswGraphBuilder implements HnswBuilder { protected final HnswLock hnswLock; protected InfoStream infoStream = InfoStream.getDefault(); + private IORunnable abortCheck; protected boolean frozen; /** @@ -209,6 +211,15 @@ public void setInfoStream(InfoStream infoStream) { this.infoStream = infoStream; } + @Override + public void setAbortCheck(IORunnable abortCheck) { + Objects.requireNonNull(abortCheck); + if (this.abortCheck != null) { + throw new IllegalStateException("abort check was already set"); + } + this.abortCheck = abortCheck; + } + @Override public OnHeapHnswGraph getCompletedGraph() throws IOException { if (!frozen) { @@ -280,6 +291,9 @@ private void addGraphNodeInternal(int node, UpdateableRandomVectorScorer scorer, if (frozen) { throw new IllegalStateException("Graph builder is already frozen"); } + if (abortCheck != null) { + abortCheck.run(); + } final int nodeLevel = getRandomGraphLevel(ml, random); // first add nodes to all levels for (int level = nodeLevel; level >= 0; level--) { diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java index 3caa9600ed1a..d9dc4144472e 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java @@ -33,6 +33,7 @@ import org.apache.lucene.util.BitSet; import org.apache.lucene.util.Bits; import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.IORunnable; import org.apache.lucene.util.InfoStream; /** @@ -46,6 +47,7 @@ public class IncrementalHnswGraphMerger implements HnswGraphMerger { protected final RandomVectorScorerSupplier scorerSupplier; protected final int M; protected final int beamWidth; + protected final IORunnable abortCheck; protected List graphReaders = new ArrayList<>(); protected GraphReader largestGraphReader; @@ -71,10 +73,26 @@ protected record GraphReader( */ public IncrementalHnswGraphMerger( FieldInfo fieldInfo, RandomVectorScorerSupplier scorerSupplier, int M, int beamWidth) { + this(fieldInfo, scorerSupplier, M, beamWidth, null); + } + + /** + * @param fieldInfo FieldInfo for the field being merged + * @param abortCheck optional check invoked before every node insertion during graph construction; + * may throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} to abort the + * build when the surrounding merge has been aborted, or null + */ + public IncrementalHnswGraphMerger( + FieldInfo fieldInfo, + RandomVectorScorerSupplier scorerSupplier, + int M, + int beamWidth, + IORunnable abortCheck) { this.fieldInfo = fieldInfo; this.scorerSupplier = scorerSupplier; this.M = M; this.beamWidth = beamWidth; + this.abortCheck = abortCheck; } /** @@ -211,6 +229,9 @@ public OnHeapHnswGraph merge( KnnVectorValues mergedVectorValues, InfoStream infoStream, int maxOrd) throws IOException { HnswBuilder builder = createBuilder(mergedVectorValues, maxOrd); builder.setInfoStream(infoStream); + if (abortCheck != null) { + builder.setAbortCheck(abortCheck); + } return builder.build(maxOrd); } diff --git a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java new file mode 100644 index 000000000000..70653dbc16d1 --- /dev/null +++ b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java @@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.index; + +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.NamedThreadFactory; + +/** + * Tests that aborting a merge (e.g. via {@link IndexWriter#rollback()}) promptly interrupts HNSW + * graph construction instead of blocking until the entire graph is built. + */ +public class TestHnswMergeAbort extends LuceneTestCase { + + private static final int DIM = 96; + private static final int SEGMENTS = 4; + private static final int DOCS_PER_SEGMENT = 12_000; + private static final int BEAM_WIDTH = 250; + + /** + * Every segment carries more than {@code IncrementalHnswGraphMerger#DELETE_PCT_THRESHOLD} + * deletions, so no source graph is eligible as a base and the merged graph is rebuilt from + * scratch via {@code HnswGraphBuilder#addVectors}. + */ + public void testRollbackDuringFullRebuildMerge() throws Exception { + doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); + } + + /** + * All segments are deletion-free, so the merged graph is produced by joining the source graphs + * via {@code MergingHnswGraphBuilder}. + */ + public void testRollbackDuringGraphJoinMerge() throws Exception { + doTestRollbackDuringMerge(false, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); + } + + /** + * The merged graph is built by {@code HnswConcurrentMergeBuilder} workers ({@code numMergeWorkers + * > 1}), which must forward the abort check to every worker. + */ + public void testRollbackDuringConcurrentMerge() throws Exception { + ExecutorService mergeExec = + Executors.newFixedThreadPool(2, new NamedThreadFactory("hnsw-merge-worker")); + try { + doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, 2, mergeExec)); + } finally { + mergeExec.shutdown(); + assertTrue(mergeExec.awaitTermination(30, TimeUnit.SECONDS)); + } + } + + private void doTestRollbackDuringMerge(boolean withDeletes, KnnVectorsFormat format) + throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig cfg = new IndexWriterConfig(); + cfg.setCodec(TestUtil.alwaysKnnVectorsFormat(format)); + cfg.setMergePolicy(NoMergePolicy.INSTANCE); + try (IndexWriter w = new IndexWriter(dir, cfg)) { + Random r = random(); + int id = 0; + for (int s = 0; s < SEGMENTS; s++) { + for (int i = 0; i < DOCS_PER_SEGMENT; i++, id++) { + Document doc = new Document(); + doc.add(new StringField("id", Integer.toString(id), Field.Store.NO)); + float[] v = new float[DIM]; + for (int j = 0; j < DIM; j++) { + v[j] = r.nextFloat(); + } + doc.add(new KnnFloatVectorField("v", v)); + w.addDocument(doc); + } + w.flush(); + } + if (withDeletes) { + // delete half of every segment, above the 40% base-graph eligibility threshold + for (int d = 0; d < id; d += 2) { + w.deleteDocuments(new Term("id", Integer.toString(d))); + } + } + w.commit(); + } + + CountDownLatch buildStarted = new CountDownLatch(1); + InfoStream latching = + new InfoStream() { + @Override + public void message(String component, String message) { + if ("HNSW".equals(component) && message.startsWith("build graph")) { + buildStarted.countDown(); + } + } + + @Override + public boolean isEnabled(String component) { + return "HNSW".equals(component); + } + + @Override + public void close() {} + }; + + IndexWriterConfig cfg2 = new IndexWriterConfig(); + cfg2.setCodec(TestUtil.alwaysKnnVectorsFormat(format)); + cfg2.setMergeScheduler(new ConcurrentMergeScheduler()); + cfg2.setInfoStream(latching); + IndexWriter w2 = new IndexWriter(dir, cfg2); + AtomicReference mergeFailure = new AtomicReference<>(); + Thread merger = + new Thread( + () -> { + try { + w2.forceMerge(1); + } catch (Throwable t) { + mergeFailure.set(t); + } + }); + merger.start(); + try { + assertTrue( + "HNSW graph construction never started", buildStarted.await(120, TimeUnit.SECONDS)); + long t0 = System.nanoTime(); + w2.rollback(); + long rollbackMillis = (System.nanoTime() - t0) / 1_000_000; + assertTrue( + "rollback() blocked for " + + rollbackMillis + + " ms waiting for HNSW graph construction to finish", + rollbackMillis < 10_000); + } finally { + merger.join(TimeUnit.MINUTES.toMillis(5)); + assertFalse("merge thread did not terminate", merger.isAlive()); + } + } + } +} diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java index db3c590d920b..0928b718f583 100644 --- a/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java +++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/HnswGraphTestCase.java @@ -860,6 +860,14 @@ public void testHnswGraphBuilderInvalid() throws IOException { IllegalArgumentException.class, () -> HnswGraphBuilder.create(scorerSupplier, 10, 0, 0)); } + public void testAbortCheckSetAtMostOnce() throws IOException { + RandomVectorScorerSupplier scorerSupplier = buildScorerSupplier(vectorValues(1, 1)); + HnswGraphBuilder builder = HnswGraphBuilder.create(scorerSupplier, 10, 10, 0); + expectThrows(NullPointerException.class, () -> builder.setAbortCheck(null)); + builder.setAbortCheck(() -> {}); + expectThrows(IllegalStateException.class, () -> builder.setAbortCheck(() -> {})); + } + // TODO: incredibly slow @Nightly public void testRamUsageEstimate() throws IOException {