michaeljmarshall commented on code in PR #16571: URL: https://github.com/apache/lucene/pull/16571#discussion_r3910808863
########## lucene/misc/src/java/org/apache/lucene/misc/index/BalancedSegmentsMergePolicy.java: ########## @@ -0,0 +1,231 @@ +/* + * 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.misc.index; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.FilterMergePolicy; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.util.Bits; + +/** + * Makes {@link org.apache.lucene.index.IndexWriter#forceMerge(int)} leave segments holding about + * the same number of documents. + * + * <p>{@code forceMerge(n)} says how many segments to leave but nothing about how the documents are + * shared between them, and in practice the share is uneven: merging a hundred equal segments of a + * thousand documents into four leaves them holding 97,000, 1,000, 1,000 and 1,000. Since a search + * is parallelised by segment, a query then waits on a segment holding 97% of the index. This policy + * leaves four segments of 25,000. + * + * <pre class="prettyprint"> + * iwc.setMergePolicy(new BalancedSegmentsMergePolicy(new TieredMergePolicy())); + * ... + * writer.forceMerge(numberOfSearchThreads); + * </pre> + * + * <p>The whole forced merge is planned at once: the segments are packed into groups and each group + * is given the number of outputs its size deserves, so the work is several independent merges the + * scheduler can run concurrently rather than one merge of everything. A group that is already a + * single segment of about the right size is left alone, so a nearly balanced index is balanced + * without being rewritten. The number of outputs is per merge rather than global, so a merge may + * also have a single input: {@code forceMerge(16)} on a one-segment index splits it into sixteen. + * + * <p>Only forced merges are affected; ordinary merging is left to the wrapped policy. No index sort + * is needed, since documents are shared out by position rather than by any key. A policy that + * placed its boundaries on a key instead would give each output a range of that key, which this one + * does not attempt. + * + * @lucene.experimental + */ +public class BalancedSegmentsMergePolicy extends FilterMergePolicy { + + /** Wraps {@link TieredMergePolicy}. */ + public BalancedSegmentsMergePolicy() { + this(new TieredMergePolicy()); + } + + /** + * @param in the policy that decides ordinary merges; only forced merges are changed + */ + public BalancedSegmentsMergePolicy(MergePolicy in) { + super(in); + } + + @Override + public MergeSpecification findForcedMerges( + SegmentInfos infos, + int maxSegmentCount, + Map<SegmentCommitInfo, Boolean> segmentsToMerge, + MergeContext context) + throws IOException { + if (maxSegmentCount == Integer.MAX_VALUE || maxSegmentCount < 1) { + // Not a request for a particular number of segments, so there is nothing to balance. + return super.findForcedMerges(infos, maxSegmentCount, segmentsToMerge, context); + } + final List<SegmentCommitInfo> eligible = new ArrayList<>(); + for (SegmentCommitInfo info : infos) { + if (segmentsToMerge.containsKey(info) == false) { + continue; + } + if (context.getMergingSegments().contains(info)) { + // A plan is already running. Since a plan is made for the whole index at once, planning + // again over what is left of it would give the outputs still being written a second share + // of the segment count, so wait for it to finish instead. + return null; + } + eligible.add(info); + } + if (eligible.isEmpty()) { + return null; + } + long total = 0; + for (SegmentCommitInfo info : eligible) { + total += liveDocs(info); + } + if (total == 0) { + return super.findForcedMerges(infos, maxSegmentCount, segmentsToMerge, context); + } + // Largest first, so a segment bigger than a share of its own takes as many outputs as it needs + // and the small ones pack into what is left. + eligible.sort(Comparator.comparingLong(BalancedSegmentsMergePolicy::liveDocs).reversed()); + + final double share = (double) total / maxSegmentCount; + final MergeSpecification spec = new MergeSpecification(); + final List<SegmentCommitInfo> group = new ArrayList<>(); + long groupDocs = 0; + int assigned = 0; + for (int i = 0; i < eligible.size(); i++) { + group.add(eligible.get(i)); + groupDocs += liveDocs(eligible.get(i)); + final int left = eligible.size() - i - 1; + if (left > 0) { + final double withNext = groupDocs + liveDocs(eligible.get(i + 1)); + if (Math.abs(withNext - share) < Math.abs(groupDocs - share)) { + continue; // taking the next segment too gets this group closer to a share + } + } + // What this group has earned, less one output for every group still to be formed. Each of + // those needs at least one segment, so at most `left` of them remain. + final int outputs = + left == 0 + ? maxSegmentCount - assigned + : Math.min((int) Math.round(groupDocs / share), maxSegmentCount - assigned - 1); + if (outputs < 1) { + continue; // no output left to give it; it joins the next group + } + if (group.size() > 1 || outputs > 1) { + // A lone segment that is already about the right size needs no merge at all. + spec.add(new Split(new ArrayList<>(group), outputs)); + } + assigned += outputs; + group.clear(); + groupDocs = 0; + } + return spec.merges.isEmpty() ? null : spec; + } + + private static long liveDocs(SegmentCommitInfo info) { + return info.info.maxDoc() - info.getDelCount(); Review Comment: Should this use `context.numDeletesToMerge(info)`? `getDelCount()` misses pending deletes and any policy-adjusted delete count, so the live-document estimates used for balancing may not match what the merge will actually output. ########## lucene/core/src/java/org/apache/lucene/index/MergeState.java: ########## @@ -287,6 +287,36 @@ static PackedLongValues removeDeletes(final int maxDoc, final Bits liveDocs) { return docMapBuilder.build(); } + /** + * A copy of {@code other} that maps documents differently: same inputs, same fields, same + * everything a format reads, but a different view of where each document ends up. + * + * <p>This exists for a merge writing several outputs from one pass, which needs the same readers + * addressed in the whole merged document space rather than in one output's. Copying keeps that + * caller out of the constructor below, whose argument list is long enough that a caller passing + * the fields through by hand is one reordering away from a silent bug. + */ + MergeState(MergeState other, DocMap[] docMaps, boolean needsIndexSort) { Review Comment: This constructor isn't used. ########## lucene/core/src/java/org/apache/lucene/index/IndexWriter.java: ########## @@ -5214,11 +5254,463 @@ private boolean assertSoftDeletesCount(CodecReader reader, int expectedCount) th return true; } + /** + * Turns one finished merge output into a segment on disk: packs its files into a compound file if + * the codec asks for one, writes its segment info, and warms it. + * + * <p>Shared by the single-output and the partitioned merge paths, which differ in how many + * outputs they produce and not in what packaging one of them means. The delete-on-failure and + * abort handling here is the delicate part, and is the reason this is one method rather than two + * similar ones. + * + * @return false if the merge was aborted while this ran, in which case the caller must abandon it + */ + private boolean packageMergedSegment( + MergePolicy.OneMerge merge, + MergePolicy mergePolicy, + SegmentCommitInfo info, + IOContext context) + throws IOException { + // Very important to do this before opening the reader + // because codec must know if prox was written for + // this segment: + boolean useCompoundFile; + synchronized (this) { // Guard segmentInfos + useCompoundFile = + info.info + .getCodec() + .compoundFormat() + .useCompoundFile(mergePolicy.size(info, this), mergePolicy); + } + + if (useCompoundFile) { + Collection<String> filesToRemove = info.files(); + // NOTE: Creation of the CFS file must be performed with the original + // directory rather than with the merging directory, so that it is not + // subject to merge throttling. + TrackingDirectoryWrapper trackingCFSDir = new TrackingDirectoryWrapper(directory); + try { + createCompoundFile(infoStream, trackingCFSDir, info.info, context, this::deleteNewFiles); + } catch (Throwable t) { + try { + synchronized (this) { + if (merge.isAborted()) { + // This can happen if rollback is called while we were building + // our CFS -- fall through to logic below to remove the non-CFS + // merged files: + if (infoStream.isEnabled("IW")) { + infoStream.message( + "IW", "hit merge abort exception creating compound file during merge: " + t); + } + return false; + } else { + handleMergeException(t, merge); + } + } + } finally { + if (infoStream.isEnabled("IW")) { + infoStream.message("IW", "hit exception creating compound file during merge: " + t); + } + // Safe: these files must exist + deleteNewFiles(info.files()); + } + } + + synchronized (this) { + + // delete new non cfs files directly: they were never + // registered with IFD + deleteNewFiles(filesToRemove); + + if (merge.isAborted()) { + if (infoStream.isEnabled("IW")) { + infoStream.message("IW", "abort merge after building CFS"); + } + // Safe: these files must exist + deleteNewFiles(info.files()); + return false; + } + } + + info.info.setUseCompoundFile(true); + } + + // Have codec write SegmentInfo. Must do this after + // creating CFS so that 1) .si isn't slurped into CFS, + // and 2) .si reflects useCompoundFile=true change + // above: + try { + config.getCodec().segmentInfoFormat().write(directory, info.info, context); + } catch (Throwable t) { + // Safe: these files must exist + deleteNewFiles(info.files()); + throw t; + } + + // TODO: ideally we would freeze info here!! + // because any changes after writing the .si will be + // lost... + + final IndexReaderWarmer mergedSegmentWarmer = config.getMergedSegmentWarmer(); + if (readerPool.isReaderPoolingEnabled() && mergedSegmentWarmer != null) { + final ReadersAndUpdates rld = getPooledInstance(info, true); + final SegmentReader sr = rld.getReader(IOContext.DEFAULT); + try { + mergedSegmentWarmer.warm(sr); + } finally { + synchronized (this) { + rld.release(sr); + release(rld); + } + } + } + return true; + } + + /** Opens a pooled reader for each of a merge's input segments, and holds their files open. */ + private void initMergeReaders(MergePolicy.OneMerge merge, IOContext context) throws IOException { + merge.initMergeReaders( + sci -> { + final ReadersAndUpdates rld = getPooledInstance(sci, true); + rld.setIsMerging(); + synchronized (this) { + return rld.getReaderForMerge( + context, mr -> deleter.incRef(mr.reader.getSegmentInfo().files())); + } + }); + } + + /** Fresh {@link SegmentCommitInfo} for one output of a merge. */ + private SegmentCommitInfo newMergeSegmentInfo(MergePolicy.OneMerge merge) { + boolean hasBlocks = false; + for (SegmentCommitInfo info : merge.segments) { + if (info.info.getHasBlocks()) { + hasBlocks = true; + break; + } + } + SegmentInfo si = + new SegmentInfo( + directoryOrig, + Version.LATEST, + null, + newSegmentName(), + -1, + false, + hasBlocks, + config.getCodec(), + Collections.emptyMap(), + StringHelper.randomId(), + Collections.emptyMap(), + config.getIndexSort()); + Map<String, String> details = new HashMap<>(); + details.put("mergeMaxNumSegments", "" + merge.maxNumSegments); + details.put("mergeFactor", Integer.toString(merge.segments.size())); + details.put("mergeOutputs", Integer.toString(merge.getOutputCount())); + setDiagnostics(si, SOURCE_MERGE, details); + return new SegmentCommitInfo(si, 0, 0, -1L, -1L, -1L, StringHelper.randomId()); + } + + /** + * Validate a doc-range partition spec. Only the SHAPE is checked -- one boundary array per input, + * all the same length, non-decreasing, starting at 0 and ending at that input's maxDoc. Given + * that shape, the two properties the merge depends on, disjointness and full coverage, hold by + * construction rather than by trusting the caller. + * + * <p>Plus the one thing that is not about the spec: an index sort. Correctness rests on the + * outputs being contiguous AND in order in the merged document space, which holds under a sort, + * where the merged order is key order and the outputs are key ranges. Without one, {@link + * DocIDMerger} concatenates input by input, so an output's documents land in one block per input + * rather than in a single run, and splitting a term's postings by document id would quietly hand + * documents to the wrong output. + */ + private void validateDocRangePartitions( + MergePolicy.OneMerge merge, int[][] partitions, List<CodecReader> readers) + throws IOException { + if (partitions.length != merge.segments.size()) { + throw new IllegalArgumentException( + "docRangePartitions has " + + partitions.length + + " entries but the merge has " + + merge.segments.size() + + " input segments"); + } + final int outputs = merge.getOutputCount(); + if (outputs < 1) { + throw new IllegalArgumentException("a partitioned merge must have at least one output"); + } + for (int i = 0; i < partitions.length; i++) { + final int[] b = partitions[i]; + final int maxDoc = merge.segments.get(i).info.maxDoc(); + if (b.length != outputs + 1) { + throw new IllegalArgumentException( + "docRangePartitions[" + i + "] has length " + b.length + ", expected " + (outputs + 1)); + } + if (b[0] != 0 || b[outputs] != maxDoc) { + throw new IllegalArgumentException( + "docRangePartitions[" + + i + + "] must span [0, " + + maxDoc + + "], got [" + + b[0] + + ", " + + b[outputs] + + "]"); + } + for (int o = 1; o <= outputs; o++) { + if (b[o] < b[o - 1]) { + throw new IllegalArgumentException( + "docRangePartitions[" + i + "] is not non-decreasing at " + o); + } + } + checkBoundariesRespectBlocks(readers.get(i), b, i); + } + } + + /** + * Refuses a boundary that falls inside a document block. + * + * <p>A block is a run of documents ending at its parent, and queries over it find the children by + * counting back from the parent. An index sort keeps a block contiguous, but says nothing about + * where a partition may cut, so a boundary landing between a child and its parent would put them + * in different segments and leave both halves quietly wrong. + * + * <p>Carrying the same partitioning value on the children is not sufficient on its own: the + * boundary is a document offset, and whether the offset a caller derives from that value lands + * before the children or between them and their parent depends on whether the children carry the + * value at all. So the invariant is enforced here, on the offsets themselves. A policy that wants + * to partition a block index has the readers when it chooses its boundaries, and can align them + * by reading the same parent field this does. + */ + private static void checkBoundariesRespectBlocks(CodecReader reader, int[] b, int input) + throws IOException { + final String parentField = reader.getFieldInfos().getParentField(); + if (parentField == null) { + return; + } + final NumericDocValues parents = reader.getNumericDocValues(parentField); + if (parents == null) { + return; + } + final int maxDoc = reader.maxDoc(); + int previous = -1; + for (int o = 1; o < b.length - 1; o++) { + final int boundary = b[o]; + // The ends of the space are block boundaries by construction, and a repeated boundary is an + // empty output, which cannot split anything the previous one did not. + if (boundary == 0 || boundary == maxDoc || boundary == previous) { + continue; + } + previous = boundary; + // A boundary is legal exactly when the document before it ends a block. + if (parents.advanceExact(boundary - 1) == false) { + throw new IllegalArgumentException( + "docRangePartitions[" + + input + + "] cuts inside a document block at " + + boundary + + ": document " + + (boundary - 1) + + " is not the last of its block, and a partitioned merge must not separate a " + + "block's documents from their parent"); + } + } + } + + /** + * Merge producing several output segments, each holding a contiguous doc range of every input. + * Kept separate from {@link #mergeMiddle} so the single-output path stays untouched. + */ + private int multiOutputMergeMiddle(MergePolicy.OneMerge merge, MergePolicy mergePolicy) + throws IOException { + testPoint("mergeMiddleStart"); + merge.checkAborted(); + + final Directory mergeDirectory = mergeScheduler.wrapForMerge(merge, directory); + final IOContext context = IOContext.merge(merge.getStoreMergeInfo()); + + boolean success = false; + int totalDocs = 0; + // Outlive the loop that creates them: the postings of every output are written between the + // phases either side of them, so each merger stays open across all three phases. + final List<SegmentMerger> mergers = new ArrayList<>(); + // The outputs share one set of input readers, so a reader that hands out itself as its own + // merge instance must still be finished exactly once, whichever path gets there first. + final Set<KnnVectorsReader> finishedVectorReaders = + Collections.newSetFromMap(new IdentityHashMap<>()); + try { + initMergeReaders(merge, context); + + // Resolved only now, so the policy can place boundaries on real key values + // from the actual readers rather than guessing from doc counts. They are + // passed in rather than read back off the merge, so that a wrapping merge + // can delegate this method. + final List<CodecReader> rawReaders = new ArrayList<>(); + for (MergePolicy.MergeReader mr : merge.getMergeReader()) { + rawReaders.add(mr.reader); + } + final int[][] partitions = merge.getDocRangePartitions(rawReaders); + if (partitions == null) { + throw new IllegalStateException( + "OneMerge.isPartitioned() returned true but getDocRangePartitions() returned null"); + } + // Before indexing into it: a policy that returns the wrong shape should be told so, rather + // than reaching this with an array index out of bounds. + if (partitions.length != merge.segments.size()) { + throw new IllegalArgumentException( + "docRangePartitions has " + + partitions.length + + " entries but the merge has " + + merge.segments.size() + + " input segments"); + } + merge.outputCount = partitions[0].length - 1; + validateDocRangePartitions(merge, partitions, rawReaders); + final int outputCount = merge.getOutputCount(); + + if (infoStream.isEnabled("IW")) { + infoStream.message( + "IW", "merging " + segString(merge.segments) + " into " + outputCount + " outputs"); + } + + // Wrapped once and shared by every output. Wrapping is where a caller filters a reader, and + // every output sees the same inputs, so doing it per output would repeat that work k times. + // This is also where the inputs are verified: each format checksums the files it is about to + // read when its merge begins, which costs a full read of them, and a partitioned merge runs + // those merges once per output. Verifying here instead leaves one check per input per merge, + // and does it before any output has written anything. + final List<CodecReader> wrappedReaders = new ArrayList<>(rawReaders.size()); + for (MergePolicy.MergeReader mergeReader : merge.getMergeReader()) { + merge.checkAborted(); + final CodecReader wrapped = merge.wrapForMerge(mergeReader.reader); + validateMergeReader(wrapped); + wrapped.checkIntegrity(merge); + wrappedReaders.add(wrapped); + } + // Only now, so that the checks just above are the ones that read the files. + merge.markInputsVerified(); + + final Executor intraMergeExecutor = mergeScheduler.getIntraMergeExecutor(merge); + final List<MergeState.DocMap[]> docMapsPerOutput = new ArrayList<>(outputCount); + merge.mergeStartNS = System.nanoTime(); + + final List<SegmentCommitInfo> outInfos = new ArrayList<>(outputCount); + final List<TrackingDirectoryWrapper> dirWrappers = new ArrayList<>(outputCount); + // Phase A: merge each output. + for (int output = 0; output < outputCount; output++) { + merge.checkAborted(); + final TrackingDirectoryWrapper dirWrapper = new TrackingDirectoryWrapper(mergeDirectory); + final SegmentCommitInfo outInfo = newMergeSegmentInfo(merge); + merge.setMergeInfo(output, outInfo); + + final List<CodecReader> mergeReaders = new ArrayList<>(); + final Counter softDeleteCount = Counter.newCounter(false); + int i = 0; + for (MergePolicy.MergeReader mergeReader : merge.getMergeReader()) { + // Everything outside this output's range looks deleted, so it maps to -1 in the + // resulting DocMap -- which is what routes concurrent deletes to the right output. + CodecReader wrapped = + new DocRangeCodecReader( + wrappedReaders.get(i), partitions[i][output], partitions[i][output + 1]); + if (softDeletesEnabled) { + // Count soft deletes that fall INSIDE this output's range. The + // single-output shortcut (softDelCount - numDeletedDocs) cannot be + // used here: numDeletedDocs on a range-restricted reader also counts + // every document belonging to the other outputs. hardLiveDocs may be + // null and countSoftDeletes handles that. + Counter hardDeleteCounter = Counter.newCounter(false); Review Comment: This counter's value isn't accessed. We need the same handling as the `if (hardDeleteCount > 0)` block in `mergeMiddle` before adding the wrapped reader to `mergeReaders`. Otherwise, a hard-deleted doc that is retained by the wrapper can be merged into the output as live. ########## lucene/misc/src/java/org/apache/lucene/misc/index/BalancedSegmentsMergePolicy.java: ########## @@ -0,0 +1,231 @@ +/* + * 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.misc.index; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.FilterMergePolicy; +import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.util.Bits; + +/** + * Makes {@link org.apache.lucene.index.IndexWriter#forceMerge(int)} leave segments holding about + * the same number of documents. + * + * <p>{@code forceMerge(n)} says how many segments to leave but nothing about how the documents are + * shared between them, and in practice the share is uneven: merging a hundred equal segments of a + * thousand documents into four leaves them holding 97,000, 1,000, 1,000 and 1,000. Since a search + * is parallelised by segment, a query then waits on a segment holding 97% of the index. This policy + * leaves four segments of 25,000. + * + * <pre class="prettyprint"> + * iwc.setMergePolicy(new BalancedSegmentsMergePolicy(new TieredMergePolicy())); + * ... + * writer.forceMerge(numberOfSearchThreads); + * </pre> + * + * <p>The whole forced merge is planned at once: the segments are packed into groups and each group + * is given the number of outputs its size deserves, so the work is several independent merges the + * scheduler can run concurrently rather than one merge of everything. A group that is already a + * single segment of about the right size is left alone, so a nearly balanced index is balanced + * without being rewritten. The number of outputs is per merge rather than global, so a merge may + * also have a single input: {@code forceMerge(16)} on a one-segment index splits it into sixteen. + * + * <p>Only forced merges are affected; ordinary merging is left to the wrapped policy. No index sort + * is needed, since documents are shared out by position rather than by any key. A policy that + * placed its boundaries on a key instead would give each output a range of that key, which this one + * does not attempt. + * + * @lucene.experimental + */ +public class BalancedSegmentsMergePolicy extends FilterMergePolicy { + + /** Wraps {@link TieredMergePolicy}. */ + public BalancedSegmentsMergePolicy() { + this(new TieredMergePolicy()); + } + + /** + * @param in the policy that decides ordinary merges; only forced merges are changed + */ + public BalancedSegmentsMergePolicy(MergePolicy in) { + super(in); + } + + @Override + public MergeSpecification findForcedMerges( + SegmentInfos infos, + int maxSegmentCount, + Map<SegmentCommitInfo, Boolean> segmentsToMerge, + MergeContext context) + throws IOException { + if (maxSegmentCount == Integer.MAX_VALUE || maxSegmentCount < 1) { + // Not a request for a particular number of segments, so there is nothing to balance. + return super.findForcedMerges(infos, maxSegmentCount, segmentsToMerge, context); + } + final List<SegmentCommitInfo> eligible = new ArrayList<>(); + for (SegmentCommitInfo info : infos) { + if (segmentsToMerge.containsKey(info) == false) { + continue; + } + if (context.getMergingSegments().contains(info)) { + // A plan is already running. Since a plan is made for the whole index at once, planning + // again over what is left of it would give the outputs still being written a second share + // of the segment count, so wait for it to finish instead. + return null; + } + eligible.add(info); + } + if (eligible.isEmpty()) { + return null; + } + long total = 0; + for (SegmentCommitInfo info : eligible) { + total += liveDocs(info); + } + if (total == 0) { + return super.findForcedMerges(infos, maxSegmentCount, segmentsToMerge, context); + } + // Largest first, so a segment bigger than a share of its own takes as many outputs as it needs + // and the small ones pack into what is left. + eligible.sort(Comparator.comparingLong(BalancedSegmentsMergePolicy::liveDocs).reversed()); + + final double share = (double) total / maxSegmentCount; + final MergeSpecification spec = new MergeSpecification(); + final List<SegmentCommitInfo> group = new ArrayList<>(); + long groupDocs = 0; + int assigned = 0; + for (int i = 0; i < eligible.size(); i++) { + group.add(eligible.get(i)); + groupDocs += liveDocs(eligible.get(i)); + final int left = eligible.size() - i - 1; + if (left > 0) { + final double withNext = groupDocs + liveDocs(eligible.get(i + 1)); + if (Math.abs(withNext - share) < Math.abs(groupDocs - share)) { + continue; // taking the next segment too gets this group closer to a share + } + } + // What this group has earned, less one output for every group still to be formed. Each of + // those needs at least one segment, so at most `left` of them remain. + final int outputs = + left == 0 + ? maxSegmentCount - assigned + : Math.min((int) Math.round(groupDocs / share), maxSegmentCount - assigned - 1); + if (outputs < 1) { + continue; // no output left to give it; it joins the next group + } + if (group.size() > 1 || outputs > 1) { + // A lone segment that is already about the right size needs no merge at all. Review Comment: This skips a single-input/single-output merge even when the segment has deletes. For example, `forceMerge(1)` on a single segment with 50% deletes would do nothing, whereas TieredMergePolicy rewrites it to reclaim the deletes. Should we only skip here when the segment is already merged, perhaps using `isMerged(infos, group.get(0), context)`? ```suggestion if (group.size() > 1 || outputs > 1 || isMerged(infos, group.getFirst(), context) == false) { // Merge when balancing requires multiple inputs or outputs, or when a lone segment is not // already fully merged. ``` ########## lucene/core/src/java/org/apache/lucene/index/IndexWriter.java: ########## @@ -5214,11 +5254,463 @@ private boolean assertSoftDeletesCount(CodecReader reader, int expectedCount) th return true; } + /** + * Turns one finished merge output into a segment on disk: packs its files into a compound file if + * the codec asks for one, writes its segment info, and warms it. + * + * <p>Shared by the single-output and the partitioned merge paths, which differ in how many + * outputs they produce and not in what packaging one of them means. The delete-on-failure and + * abort handling here is the delicate part, and is the reason this is one method rather than two + * similar ones. + * + * @return false if the merge was aborted while this ran, in which case the caller must abandon it + */ + private boolean packageMergedSegment( + MergePolicy.OneMerge merge, + MergePolicy mergePolicy, + SegmentCommitInfo info, + IOContext context) + throws IOException { + // Very important to do this before opening the reader + // because codec must know if prox was written for + // this segment: + boolean useCompoundFile; + synchronized (this) { // Guard segmentInfos + useCompoundFile = + info.info + .getCodec() + .compoundFormat() + .useCompoundFile(mergePolicy.size(info, this), mergePolicy); + } + + if (useCompoundFile) { + Collection<String> filesToRemove = info.files(); + // NOTE: Creation of the CFS file must be performed with the original + // directory rather than with the merging directory, so that it is not + // subject to merge throttling. + TrackingDirectoryWrapper trackingCFSDir = new TrackingDirectoryWrapper(directory); + try { + createCompoundFile(infoStream, trackingCFSDir, info.info, context, this::deleteNewFiles); + } catch (Throwable t) { + try { + synchronized (this) { + if (merge.isAborted()) { + // This can happen if rollback is called while we were building + // our CFS -- fall through to logic below to remove the non-CFS + // merged files: + if (infoStream.isEnabled("IW")) { + infoStream.message( + "IW", "hit merge abort exception creating compound file during merge: " + t); + } + return false; + } else { + handleMergeException(t, merge); + } + } + } finally { + if (infoStream.isEnabled("IW")) { + infoStream.message("IW", "hit exception creating compound file during merge: " + t); + } + // Safe: these files must exist + deleteNewFiles(info.files()); + } + } + + synchronized (this) { + + // delete new non cfs files directly: they were never + // registered with IFD + deleteNewFiles(filesToRemove); + + if (merge.isAborted()) { + if (infoStream.isEnabled("IW")) { + infoStream.message("IW", "abort merge after building CFS"); + } + // Safe: these files must exist + deleteNewFiles(info.files()); + return false; + } + } + + info.info.setUseCompoundFile(true); + } + + // Have codec write SegmentInfo. Must do this after + // creating CFS so that 1) .si isn't slurped into CFS, + // and 2) .si reflects useCompoundFile=true change + // above: + try { + config.getCodec().segmentInfoFormat().write(directory, info.info, context); + } catch (Throwable t) { + // Safe: these files must exist + deleteNewFiles(info.files()); + throw t; + } + + // TODO: ideally we would freeze info here!! + // because any changes after writing the .si will be + // lost... + + final IndexReaderWarmer mergedSegmentWarmer = config.getMergedSegmentWarmer(); + if (readerPool.isReaderPoolingEnabled() && mergedSegmentWarmer != null) { + final ReadersAndUpdates rld = getPooledInstance(info, true); + final SegmentReader sr = rld.getReader(IOContext.DEFAULT); + try { + mergedSegmentWarmer.warm(sr); + } finally { + synchronized (this) { + rld.release(sr); + release(rld); + } + } + } + return true; + } + + /** Opens a pooled reader for each of a merge's input segments, and holds their files open. */ + private void initMergeReaders(MergePolicy.OneMerge merge, IOContext context) throws IOException { + merge.initMergeReaders( + sci -> { + final ReadersAndUpdates rld = getPooledInstance(sci, true); + rld.setIsMerging(); + synchronized (this) { + return rld.getReaderForMerge( + context, mr -> deleter.incRef(mr.reader.getSegmentInfo().files())); + } + }); + } + + /** Fresh {@link SegmentCommitInfo} for one output of a merge. */ + private SegmentCommitInfo newMergeSegmentInfo(MergePolicy.OneMerge merge) { + boolean hasBlocks = false; + for (SegmentCommitInfo info : merge.segments) { + if (info.info.getHasBlocks()) { + hasBlocks = true; + break; + } + } + SegmentInfo si = + new SegmentInfo( + directoryOrig, + Version.LATEST, + null, + newSegmentName(), + -1, + false, + hasBlocks, + config.getCodec(), + Collections.emptyMap(), + StringHelper.randomId(), + Collections.emptyMap(), + config.getIndexSort()); + Map<String, String> details = new HashMap<>(); + details.put("mergeMaxNumSegments", "" + merge.maxNumSegments); + details.put("mergeFactor", Integer.toString(merge.segments.size())); + details.put("mergeOutputs", Integer.toString(merge.getOutputCount())); + setDiagnostics(si, SOURCE_MERGE, details); + return new SegmentCommitInfo(si, 0, 0, -1L, -1L, -1L, StringHelper.randomId()); + } + + /** + * Validate a doc-range partition spec. Only the SHAPE is checked -- one boundary array per input, + * all the same length, non-decreasing, starting at 0 and ending at that input's maxDoc. Given + * that shape, the two properties the merge depends on, disjointness and full coverage, hold by + * construction rather than by trusting the caller. + * + * <p>Plus the one thing that is not about the spec: an index sort. Correctness rests on the + * outputs being contiguous AND in order in the merged document space, which holds under a sort, + * where the merged order is key order and the outputs are key ranges. Without one, {@link + * DocIDMerger} concatenates input by input, so an output's documents land in one block per input + * rather than in a single run, and splitting a term's postings by document id would quietly hand + * documents to the wrong output. + */ Review Comment: Is an index sort actually required here? AFAICT, it isn't required by the current implementation, and the `getDocRangePartitions` javadocs say that it isn't required: "An index sort is not required..." -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
