From 41f256cead7ed3e8dd76d7a2b573689679fc15f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=A0=9C=ED=98=B8?= <135206467+jeho-rpls@users.noreply.github.com> Date: Mon, 13 Jul 2026 22:04:46 +0900 Subject: [PATCH] Check if merge is aborted during HNSW graph construction (#16368) (cherry picked from commit 2b237286b38e9b62d6e017a1d746a66c707e8462) --- lucene/CHANGES.txt | 4 + .../lucene99/Lucene99HnswVectorsWriter.java | 13 +- .../util/hnsw/ConcurrentHnswMerger.java | 20 ++- .../apache/lucene/util/hnsw/HnswBuilder.java | 13 ++ .../util/hnsw/HnswConcurrentMergeBuilder.java | 8 + .../lucene/util/hnsw/HnswGraphBuilder.java | 14 ++ .../util/hnsw/IncrementalHnswGraphMerger.java | 21 +++ .../lucene/index/TestHnswMergeAbort.java | 163 ++++++++++++++++++ .../lucene/util/hnsw/HnswGraphTestCase.java | 8 + 9 files changed, 258 insertions(+), 6 deletions(-) create mode 100644 lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 245d97efe8c7..0414abf02717 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -83,6 +83,10 @@ Bug Fixes * GITHUB#16389: Fix double-counting of the underlying Automaton in AutomatonQuery#ramBytesUsed. (Sasilekha R) +* 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) + Other --------------------- * GITHUB#16266: Remove deprecated search(Query, Collector) calls in QueryUtils by replacing 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 ebfd3ec5470d..8beb791560eb 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 @@ -477,7 +477,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( @@ -629,10 +630,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( @@ -641,9 +643,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 9421dfb1dc87..f540d9f406e1 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 d34a3c68bbf3..6a63ba596a64 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; /** @@ -136,6 +137,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 43cc3488fe1a..72ee257eaf75 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; /** @@ -230,6 +232,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) { @@ -301,6 +312,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 083c8f8c7c04..c94b1610ddc8 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 @@ -34,6 +34,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; /** @@ -47,6 +48,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; @@ -72,10 +74,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; } /** @@ -215,6 +233,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 2d1ef75c3c8b..f9b67210f217 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 @@ -868,6 +868,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(() -> {})); + } + public void testRamUsageEstimate() throws IOException { int size = atLeast(2000); int dim = randomIntBetween(100, 1024);