You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by gs...@apache.org on 2026/08/14 02:35:14 UTC
(lucene) branch branch_10x updated: Redesign ReaderUtil.partitionByLeaf to take int[] and track input ordinals for scatter/gather (#16496)
This is an automated email from the ASF dual-hosted git repository.
gsmiller pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new a4e4aec411e Redesign ReaderUtil.partitionByLeaf to take int[] and track input ordinals for scatter/gather (#16496)
a4e4aec411e is described below
commit a4e4aec411e700112619b1797c51c7d67973bb99
Author: Zihan Xu <zi...@andrew.cmu.edu>
AuthorDate: Thu Aug 13 18:47:09 2026 -0700
Redesign ReaderUtil.partitionByLeaf to take int[] and track input ordinals for scatter/gather (#16496)
---
lucene/CHANGES.txt | 4 +
.../benchmark/jmh/PartitionByLeafBenchmark.java | 206 ++++++++++++++++++++-
.../java/org/apache/lucene/index/ReaderUtil.java | 90 ++++++---
.../org/apache/lucene/index/TestReaderUtil.java | 157 +++++++++-------
4 files changed, 360 insertions(+), 97 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 38a145ee575..c12e4273341 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -16,6 +16,10 @@ API Changes
* GITHUB#16363: Add API to use a custom `DictionarySuggester` instead of hardcoded `GeneratingSuggester`
+* GITHUB#16496: Redesign ReaderUtil.partitionByLeaf to take an int[] of global doc IDs instead of a
+ ScoreDoc[] and return a PartitionedHits record of per-leaf doc IDs plus per-leaf original input
+ positions (ordinals) for scatter/gather. (Zihan Xu)
+
* GITHUB#16385: Remove unused public methods in AllGroupHeadsCollector.
Update grouping package info to use CollectorManager instead of Collector. (Binlong Gao)
diff --git a/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/PartitionByLeafBenchmark.java b/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/PartitionByLeafBenchmark.java
index 9f9ef2f3fa7..a70ff210d05 100644
--- a/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/PartitionByLeafBenchmark.java
+++ b/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/PartitionByLeafBenchmark.java
@@ -20,6 +20,7 @@ import java.util.Arrays;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.IntroSorter;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
@@ -35,13 +36,35 @@ import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;
/**
- * Benchmark comparing partition strategies for ReaderUtil#partitionByLeaf. Both benchmarks operate
- * on pre-sorted doc IDs to isolate the partition step from sorting overhead.
+ * Benchmark for partition strategies behind ReaderUtil#partitionByLeaf.
+ *
+ * <p>Two families of benchmarks live here:
*
* <ul>
- * <li>linearPartition: linear-scan partition (previous implementation)
- * <li>binarySearchPartition: binary-search partition using leaf boundaries (current
- * implementation)
+ * <li><b>Partition step only</b> (operate on pre-sorted doc IDs, isolating the partition from
+ * sorting overhead):
+ * <ul>
+ * <li>linearPartition: linear-scan partition
+ * <li>binarySearchPartition: binary-search partition using leaf boundaries
+ * </ul>
+ * <li><b>Full partition task</b> (start from unsorted input, so they include the sort): these
+ * compare the two candidate public-API shapes for a caller that only wants the per-leaf doc
+ * IDs and does <em>not</em> need input ordinals:
+ * <ul>
+ * <li>partitionByLeaf: clone + sort {@code int[]} + binary-search partition. No ordinal
+ * bookkeeping.
+ * <li>partitionByLeafWithOrdinals: pack {@code (docId, ordinal)} into a {@code long[]} +
+ * sort + binary-search partition, producing both per-leaf doc IDs and per-leaf
+ * ordinals. This mimics a caller that is "forced" through the ordinal-tracking API: it
+ * pays to build the ordinals array too, then uses only the per-leaf doc IDs. The delta
+ * against {@code partitionByLeaf} measures the overhead of standardizing on a single
+ * ordinal-tracking API.
+ * <li>partitionByLeafWithOrdinalsIntroSort: the same ordinal-tracking task, but sorting the
+ * doc IDs and ordinals as parallel {@code int[]}s with an {@link IntroSorter} instead
+ * of packing into a {@code long[]}. The delta against {@code
+ * partitionByLeafWithOrdinals} isolates the packed-long sort's win over the
+ * parallel-array sort.
+ * </ul>
* </ul>
*/
@BenchmarkMode(Mode.Throughput)
@@ -67,6 +90,12 @@ public class PartitionByLeafBenchmark {
/** Pre-sorted doc IDs to partition. */
private int[] sortedDocIds;
+ /**
+ * Unsorted doc IDs (the shuffled pool prefix), used by the full-task benchmarks so they include
+ * the sort. Never mutated by the benchmarks: each iteration clones/packs it before sorting.
+ */
+ private int[] unsortedDocIds;
+
/** Leaf boundaries: leafDocBase[i] is the docBase for leaf i. */
private int[] leafDocBase;
@@ -96,7 +125,10 @@ public class PartitionByLeafBenchmark {
pool[i] = pool[j];
pool[j] = tmp;
}
- sortedDocIds = ArrayUtil.copyOfSubArray(pool, 0, numDocIds);
+ // Unsorted (shuffled) prefix drives the full-task benchmarks; the sorted copy drives the
+ // partition-only benchmarks.
+ unsortedDocIds = ArrayUtil.copyOfSubArray(pool, 0, numDocIds);
+ sortedDocIds = unsortedDocIds.clone();
Arrays.sort(sortedDocIds);
}
@@ -111,9 +143,81 @@ public class PartitionByLeafBenchmark {
}
/**
- * Partition sorted doc IDs across leaves using a linear scan. This mirrors the previous
- * implementation in ReaderUtil#partitionByLeaf.
+ * Full task for a caller that only wants per-leaf doc IDs: clone + sort the {@code int[]}, then
+ * binary-search partition.
+ */
+ @Benchmark
+ public void partitionByLeaf(Blackhole bh) {
+ int[] sorted = unsortedDocIds.clone();
+ Arrays.sort(sorted);
+ bh.consume(partitionSortedBinarySearch(sorted));
+ }
+
+ /**
+ * Full task through the packed-long ordinal-tracking strategy for a caller that ultimately
+ * ignores the ordinals: pack {@code (docId, ordinal)} into a {@code long[]}, sort, then
+ * binary-search partition, producing both per-leaf doc IDs and per-leaf ordinals. This models a
+ * caller that doesn't need ordinals but is forced through a single ordinal-tracking API: the
+ * ordinals are built regardless of what the caller looks at. We consume the whole result (doc IDs
+ * and ordinals) so that ordinal work can't be eliminated as dead code and the measurement stays
+ * faithful to the work the ordinal-tracking strategy performs.
+ */
+ @Benchmark
+ public void partitionByLeafWithOrdinals(Blackhole bh) {
+ long[] packed = new long[unsortedDocIds.length];
+ for (int i = 0; i < packed.length; i++) {
+ packed[i] = ((long) unsortedDocIds[i] << 32) | i;
+ }
+ Arrays.sort(packed);
+ bh.consume(partitionSortedPackedWithOrdinals(packed));
+ }
+
+ /**
+ * The same ordinal-tracking full task as {@link #partitionByLeafWithOrdinals}, but sorting the
+ * doc IDs and their ordinals as two parallel {@code int[]}s with an {@link IntroSorter} rather
+ * than packing into a {@code long[]}. Kept alongside the packed-long variant so the two can be
+ * compared directly. As with the packed variant, we consume both output arrays so the ordinal
+ * work can't be eliminated as dead code.
*/
+ @Benchmark
+ public void partitionByLeafWithOrdinalsIntroSort(Blackhole bh) {
+ final int[] sortedDocIds = unsortedDocIds.clone();
+ final int[] sortedOrdinals = new int[unsortedDocIds.length];
+ for (int i = 0; i < sortedOrdinals.length; i++) {
+ sortedOrdinals[i] = i;
+ }
+ new IntroSorter() {
+ int pivot;
+
+ @Override
+ protected int compare(int i, int j) {
+ return Integer.compare(sortedDocIds[i], sortedDocIds[j]);
+ }
+
+ @Override
+ protected void swap(int i, int j) {
+ int tmp = sortedDocIds[i];
+ sortedDocIds[i] = sortedDocIds[j];
+ sortedDocIds[j] = tmp;
+ tmp = sortedOrdinals[i];
+ sortedOrdinals[i] = sortedOrdinals[j];
+ sortedOrdinals[j] = tmp;
+ }
+
+ @Override
+ protected void setPivot(int i) {
+ pivot = sortedDocIds[i];
+ }
+
+ @Override
+ protected int comparePivot(int j) {
+ return Integer.compare(pivot, sortedDocIds[j]);
+ }
+ }.sort(0, sortedDocIds.length);
+ bh.consume(partitionSortedParallelWithOrdinals(sortedDocIds, sortedOrdinals));
+ }
+
+ /** Partition sorted doc IDs across leaves using a linear scan. */
private int[][] partitionSortedLinear(int[] sortedDocIds) {
int[][] result = new int[numLeaves][];
if (sortedDocIds.length == 0) {
@@ -145,6 +249,92 @@ public class PartitionByLeafBenchmark {
return result;
}
+ /**
+ * Binary-search partition over a sorted array of packed {@code (docId << 32) | ordinal} entries,
+ * unpacking into per-leaf doc IDs <em>and</em> per-leaf ordinals. Returns both arrays as {@code
+ * {docIds, ordinals}} so neither the ordinal allocation nor its unpack loop can be eliminated as
+ * dead code — that produced-but-unused work is precisely the overhead this benchmark measures.
+ */
+ private int[][][] partitionSortedPackedWithOrdinals(long[] sortedPacked) {
+ int[][] docIdsByLeaf = new int[numLeaves][];
+ int[][] ordinalsByLeaf = new int[numLeaves][];
+ if (sortedPacked.length == 0) {
+ Arrays.fill(docIdsByLeaf, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, EMPTY_INT_ARRAY);
+ return new int[][][] {docIdsByLeaf, ordinalsByLeaf};
+ }
+ int from = 0;
+ int leafIdx = 0;
+ for (; leafIdx < numLeaves && from < sortedPacked.length; leafIdx++) {
+ int leafEnd = leafDocBase[leafIdx] + docsPerLeaf;
+ long leafEndPacked = ((long) leafEnd) << 32;
+ if (sortedPacked[from] >= leafEndPacked) {
+ docIdsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
+ ordinalsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
+ continue;
+ }
+ int to = Arrays.binarySearch(sortedPacked, from, sortedPacked.length, leafEndPacked);
+ if (to < 0) {
+ to = -to - 1;
+ }
+ int count = to - from;
+ int[] leafDocs = new int[count];
+ int[] leafOrds = new int[count];
+ for (int i = 0; i < count; i++) {
+ long packed = sortedPacked[from + i];
+ leafDocs[i] = (int) (packed >>> 32);
+ leafOrds[i] = (int) packed;
+ }
+ docIdsByLeaf[leafIdx] = leafDocs;
+ ordinalsByLeaf[leafIdx] = leafOrds;
+ from = to;
+ }
+ Arrays.fill(docIdsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ return new int[][][] {docIdsByLeaf, ordinalsByLeaf};
+ }
+
+ /**
+ * Binary-search partition over parallel sorted doc IDs and their ordinals, slicing each per-leaf
+ * run out with {@link System#arraycopy}. Returns both arrays as {@code {docIds, ordinals}} so
+ * neither the ordinal allocation nor its copy can be eliminated as dead code — that
+ * produced-but-unused work is precisely the overhead this benchmark measures.
+ */
+ private int[][][] partitionSortedParallelWithOrdinals(int[] sortedDocIds, int[] sortedOrdinals) {
+ int[][] docIdsByLeaf = new int[numLeaves][];
+ int[][] ordinalsByLeaf = new int[numLeaves][];
+ if (sortedDocIds.length == 0) {
+ Arrays.fill(docIdsByLeaf, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, EMPTY_INT_ARRAY);
+ return new int[][][] {docIdsByLeaf, ordinalsByLeaf};
+ }
+ int from = 0;
+ int leafIdx = 0;
+ for (; leafIdx < numLeaves && from < sortedDocIds.length; leafIdx++) {
+ int leafEnd = leafDocBase[leafIdx] + docsPerLeaf;
+ if (sortedDocIds[from] >= leafEnd) {
+ docIdsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
+ ordinalsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
+ continue;
+ }
+ int to = Arrays.binarySearch(sortedDocIds, from, sortedDocIds.length, leafEnd);
+ if (to < 0) {
+ to = -to - 1;
+ }
+ int count = to - from;
+ int[] leafDocs = new int[count];
+ int[] leafOrds = new int[count];
+ System.arraycopy(sortedDocIds, from, leafDocs, 0, count);
+ System.arraycopy(sortedOrdinals, from, leafOrds, 0, count);
+ docIdsByLeaf[leafIdx] = leafDocs;
+ ordinalsByLeaf[leafIdx] = leafOrds;
+ from = to;
+ }
+ Arrays.fill(docIdsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ return new int[][][] {docIdsByLeaf, ordinalsByLeaf};
+ }
+
/**
* Partition sorted doc IDs across leaves using binary search on leaf boundaries. For each leaf,
* binary search for its end boundary in the sorted doc IDs to find the slice belonging to that
diff --git a/lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java b/lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java
index 046528a602c..ecc0c6dc271 100644
--- a/lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java
+++ b/lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java
@@ -18,7 +18,6 @@ package org.apache.lucene.index;
import java.util.Arrays;
import java.util.List;
-import org.apache.lucene.search.ScoreDoc;
/**
* Common util methods for dealing with {@link IndexReader}s and {@link IndexReaderContext}s.
@@ -95,45 +94,90 @@ public final class ReaderUtil {
}
/**
- * Partitions global doc IDs from ScoreDoc array by leaf. Extracts doc IDs, sorts them, and
- * partitions across leaves.
+ * Result of partitioning doc IDs by leaf, including the original input ordinals for
+ * scatter/gather. {@code docIdsByLeaf[k]} holds the sorted global doc IDs that fall in leaf
+ * {@code k}, and {@code ordinalsByLeaf[k][i]} is the index in the original {@code globalDocIds}
+ * input array of the doc ID at {@code docIdsByLeaf[k][i]}.
*
- * @param hits the ScoreDoc array (typically from TopDocs.scoreDocs)
+ * <p>Both arrays have the same shape: {@code ordinalsByLeaf[k].length == docIdsByLeaf[k].length}
+ * for every leaf {@code k}.
+ *
+ * @param docIdsByLeaf per-leaf sorted global doc IDs; {@code docIdsByLeaf[k]} holds the doc IDs
+ * that fall in leaf {@code k} (empty if none)
+ * @param ordinalsByLeaf per-leaf original input positions; {@code ordinalsByLeaf[k][i]} is the
+ * index in the original input array of the doc ID at {@code docIdsByLeaf[k][i]}
+ */
+ public record PartitionedHits(int[][] docIdsByLeaf, int[][] ordinalsByLeaf) {}
+
+ /**
+ * Partitions global doc IDs by leaf, tracking each doc ID's original position in the input array
+ * so callers can reassemble per-leaf results back to input order (scatter/gather).
+ *
+ * <p>Doc IDs may be supplied in any order (e.g., ranking order); within each leaf the returned
+ * doc IDs are sorted in ascending order. For every partitioned doc ID the result also records its
+ * index in the original {@code globalDocIds} array, so callers can map per-leaf results back to
+ * input order. Callers that do not need the ordinals can simply ignore {@link
+ * PartitionedHits#ordinalsByLeaf()} and use {@link PartitionedHits#docIdsByLeaf()}.
+ *
+ * <p>The input array is not mutated.
+ *
+ * @param globalDocIds global doc IDs in any order (e.g., ranking order)
* @param leaves the index reader's leaves
- * @return array indexed by leaf ord, containing global doc IDs for that leaf (empty if no hits)
+ * @return per-leaf sorted doc IDs alongside per-leaf ordinals into the input array
*/
- public static int[][] partitionByLeaf(ScoreDoc[] hits, List<LeafReaderContext> leaves) {
+ public static PartitionedHits partitionByLeaf(
+ int[] globalDocIds, List<LeafReaderContext> leaves) {
int numLeaves = leaves.size();
- int[][] result = new int[numLeaves][];
- if (hits.length == 0) {
- Arrays.fill(result, EMPTY_INT_ARRAY);
- return result;
+ if (globalDocIds.length == 0) {
+ int[][] docIdsByLeaf = new int[numLeaves][];
+ int[][] ordinalsByLeaf = new int[numLeaves][];
+ Arrays.fill(docIdsByLeaf, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, EMPTY_INT_ARRAY);
+ return new PartitionedHits(docIdsByLeaf, ordinalsByLeaf);
}
- int[] sortedDocIds = new int[hits.length];
- for (int i = 0; i < hits.length; i++) {
- sortedDocIds[i] = hits[i].doc;
+
+ // Pack each (docId, ordinal) into a long: docId in the high 32 bits, ordinal in the low 32,
+ // then sort with the primitive Arrays.sort(long[]). Both values are non-negative, so ascending
+ // long order matches ascending docId order and the ordinal rides along.
+ final long[] packed = new long[globalDocIds.length];
+ for (int i = 0; i < packed.length; i++) {
+ packed[i] = ((long) globalDocIds[i] << 32) | (i & 0xFFFFFFFFL);
}
- Arrays.sort(sortedDocIds);
+ Arrays.sort(packed);
+
+ // Partition into per-leaf slices via binary search on each leaf's end boundary
+ int[][] docIdsByLeaf = new int[numLeaves][];
+ int[][] ordinalsByLeaf = new int[numLeaves][];
int from = 0;
int leafIdx = 0;
- for (; leafIdx < numLeaves && from < sortedDocIds.length; leafIdx++) {
+ for (; leafIdx < numLeaves && from < packed.length; leafIdx++) {
LeafReaderContext leaf = leaves.get(leafIdx);
- int leafEnd = leaf.docBase + leaf.reader().maxDoc();
- if (sortedDocIds[from] >= leafEnd) {
- result[leafIdx] = EMPTY_INT_ARRAY;
+ long leafEndPacked = ((long) (leaf.docBase + leaf.reader().maxDoc())) << 32;
+ if (packed[from] >= leafEndPacked) {
+ docIdsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
+ ordinalsByLeaf[leafIdx] = EMPTY_INT_ARRAY;
continue;
}
- int to = Arrays.binarySearch(sortedDocIds, from, sortedDocIds.length, leafEnd);
+ int to = Arrays.binarySearch(packed, from, packed.length, leafEndPacked);
if (to < 0) {
to = -to - 1;
}
int count = to - from;
assert count > 0;
- result[leafIdx] = new int[count];
- System.arraycopy(sortedDocIds, from, result[leafIdx], 0, count);
+ int[] leafDocs = new int[count];
+ int[] leafOrds = new int[count];
+ for (int i = 0; i < count; i++) {
+ long p = packed[from + i];
+ leafDocs[i] = (int) (p >>> 32);
+ leafOrds[i] = (int) p;
+ }
+ docIdsByLeaf[leafIdx] = leafDocs;
+ ordinalsByLeaf[leafIdx] = leafOrds;
from = to;
}
- Arrays.fill(result, leafIdx, numLeaves, EMPTY_INT_ARRAY);
- return result;
+
+ Arrays.fill(docIdsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ Arrays.fill(ordinalsByLeaf, leafIdx, numLeaves, EMPTY_INT_ARRAY);
+ return new PartitionedHits(docIdsByLeaf, ordinalsByLeaf);
}
}
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestReaderUtil.java b/lucene/core/src/test/org/apache/lucene/index/TestReaderUtil.java
index e1ccb1191e1..ba6e64205b8 100644
--- a/lucene/core/src/test/org/apache/lucene/index/TestReaderUtil.java
+++ b/lucene/core/src/test/org/apache/lucene/index/TestReaderUtil.java
@@ -22,50 +22,34 @@ import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.apache.lucene.document.Document;
-import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.util.LuceneTestCase;
public class TestReaderUtil extends LuceneTestCase {
- public void testPartitionByLeafEmptyDocIds() throws IOException {
+ public void testEmpty() throws IOException {
try (Directory dir = newDirectory();
IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig())) {
writer.addDocument(new Document());
writer.addDocument(new Document());
try (DirectoryReader reader = DirectoryReader.open(writer)) {
List<LeafReaderContext> leaves = reader.leaves();
- int[][] result = ReaderUtil.partitionByLeaf(new ScoreDoc[0], leaves);
- assertEquals(leaves.size(), result.length);
- for (int[] leaf : result) {
- assertEquals(0, leaf.length);
+ ReaderUtil.PartitionedHits partitioned = ReaderUtil.partitionByLeaf(new int[0], leaves);
+ assertEquals(leaves.size(), partitioned.docIdsByLeaf().length);
+ assertEquals(leaves.size(), partitioned.ordinalsByLeaf().length);
+ for (int i = 0; i < leaves.size(); i++) {
+ assertEquals(0, partitioned.docIdsByLeaf()[i].length);
+ assertEquals(0, partitioned.ordinalsByLeaf()[i].length);
}
}
}
}
- public void testPartitionByLeafSingleSegment() throws IOException {
- try (Directory dir = newDirectory();
- IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig())) {
- for (int i = 0; i < 10; i++) {
- writer.addDocument(new Document());
- }
- try (DirectoryReader reader = DirectoryReader.open(writer)) {
- List<LeafReaderContext> leaves = reader.leaves();
- assertEquals(1, leaves.size());
-
- ScoreDoc[] hits = {
- new ScoreDoc(0, 1f), new ScoreDoc(3, 1f), new ScoreDoc(5, 1f), new ScoreDoc(9, 1f)
- };
- int[][] result = ReaderUtil.partitionByLeaf(hits, leaves);
-
- assertEquals(1, result.length);
- assertArrayEquals(new int[] {0, 3, 5, 9}, result[0]);
- }
- }
- }
-
- public void testPartitionByLeafMultipleSegments() throws IOException {
+ /**
+ * Worked two-segment example with deliberately unsorted input: verifies per-leaf doc IDs are
+ * sorted, ordinals point back into the original input array, and the round-trip invariant holds.
+ */
+ public void testTracksOriginalPositions() throws IOException {
try (Directory dir = newDirectory();
IndexWriter writer =
new IndexWriter(dir, new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) {
@@ -73,8 +57,6 @@ public class TestReaderUtil extends LuceneTestCase {
writer.addDocument(new Document());
}
writer.commit();
-
- // Create second segment
for (int i = 0; i < 10; i++) {
writer.addDocument(new Document());
}
@@ -84,27 +66,41 @@ public class TestReaderUtil extends LuceneTestCase {
List<LeafReaderContext> leaves = reader.leaves();
assertEquals(2, leaves.size());
- // Hits in both segments
- ScoreDoc[] hits = {
- new ScoreDoc(2, 1f), new ScoreDoc(9, 1f), new ScoreDoc(10, 1f), new ScoreDoc(18, 1f)
- };
- int[][] result = ReaderUtil.partitionByLeaf(hits, leaves);
-
- assertEquals(2, result.length);
- // First segment: docs 0-9
- assertArrayEquals(new int[] {2, 9}, result[0]);
- // Second segment: docs 10-19
- assertArrayEquals(new int[] {10, 18}, result[1]);
+ // Input order is intentionally not docId-sorted; we want to verify ordinals point back
+ // into this array (and that the per-leaf doc IDs come out sorted).
+ int[] docIds = {18, 2, 10, 9};
+ ReaderUtil.PartitionedHits partitioned = ReaderUtil.partitionByLeaf(docIds, leaves);
+
+ // Leaf 0 holds docs 0-9: {2 (ord 1), 9 (ord 3)}
+ assertArrayEquals(new int[] {2, 9}, partitioned.docIdsByLeaf()[0]);
+ assertArrayEquals(new int[] {1, 3}, partitioned.ordinalsByLeaf()[0]);
+
+ // Leaf 1 holds docs 10-19: {10 (ord 2), 18 (ord 0)}
+ assertArrayEquals(new int[] {10, 18}, partitioned.docIdsByLeaf()[1]);
+ assertArrayEquals(new int[] {2, 0}, partitioned.ordinalsByLeaf()[1]);
+
+ // Round-trip: docIdsByLeaf[k][i] == globalDocIds[ordinalsByLeaf[k][i]]
+ for (int k = 0; k < leaves.size(); k++) {
+ int[] leafDocs = partitioned.docIdsByLeaf()[k];
+ int[] leafOrds = partitioned.ordinalsByLeaf()[k];
+ assertEquals(leafDocs.length, leafOrds.length);
+ for (int i = 0; i < leafDocs.length; i++) {
+ assertEquals(docIds[leafOrds[i]], leafDocs[i]);
+ }
+ }
}
}
}
- public void testPartitionByLeafSkipsSegmentsWithNoHits() throws IOException {
+ /**
+ * Leaves with no hits (including an empty middle leaf and trailing empty leaves) come back empty.
+ */
+ public void testSkipsSegmentsWithNoHits() throws IOException {
try (Directory dir = newDirectory();
IndexWriter writer =
new IndexWriter(dir, new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) {
- // Create 3 segments
- for (int seg = 0; seg < 3; seg++) {
+ // Create 4 segments (docs 0-9, 10-19, 20-29, 30-39).
+ for (int seg = 0; seg < 4; seg++) {
for (int i = 0; i < 10; i++) {
writer.addDocument(new Document());
}
@@ -113,21 +109,46 @@ public class TestReaderUtil extends LuceneTestCase {
try (DirectoryReader reader = DirectoryReader.open(writer)) {
List<LeafReaderContext> leaves = reader.leaves();
- assertEquals(3, leaves.size());
+ assertEquals(4, leaves.size());
+
+ // Hits only in the first and third segments: empty middle leaf (1) and trailing empty
+ // leaf (3). Input unsorted to also exercise sorting/ordinals.
+ int[] docIds = {25, 3};
+ ReaderUtil.PartitionedHits partitioned = ReaderUtil.partitionByLeaf(docIds, leaves);
+
+ assertEquals(4, partitioned.docIdsByLeaf().length);
+ assertEquals(4, partitioned.ordinalsByLeaf().length);
+
+ assertArrayEquals(new int[] {3}, partitioned.docIdsByLeaf()[0]);
+ assertArrayEquals(new int[] {1}, partitioned.ordinalsByLeaf()[0]); // 3 was input index 1
+ assertEquals(0, partitioned.docIdsByLeaf()[1].length); // empty middle leaf
+ assertEquals(0, partitioned.ordinalsByLeaf()[1].length);
+ assertArrayEquals(new int[] {25}, partitioned.docIdsByLeaf()[2]);
+ assertArrayEquals(new int[] {0}, partitioned.ordinalsByLeaf()[2]); // 25 was input index 0
+ assertEquals(0, partitioned.docIdsByLeaf()[3].length); // trailing empty leaf
+ assertEquals(0, partitioned.ordinalsByLeaf()[3].length);
+ }
+ }
+ }
- // Hits only in first and third segment (skip middle)
- ScoreDoc[] hits = {new ScoreDoc(3, 1f), new ScoreDoc(25, 1f)};
- int[][] result = ReaderUtil.partitionByLeaf(hits, leaves);
+ public void testDoesNotMutateInput() throws IOException {
+ try (Directory dir = newDirectory();
+ IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig())) {
+ for (int i = 0; i < 10; i++) {
+ writer.addDocument(new Document());
+ }
+ try (DirectoryReader reader = DirectoryReader.open(writer)) {
+ List<LeafReaderContext> leaves = reader.leaves();
- assertEquals(3, result.length);
- assertArrayEquals(new int[] {3}, result[0]);
- assertEquals(0, result[1].length); // middle segment has no hits
- assertArrayEquals(new int[] {25}, result[2]);
+ int[] docIds = {9, 3, 0, 5};
+ int[] copy = docIds.clone();
+ ReaderUtil.partitionByLeaf(docIds, leaves);
+ assertArrayEquals(copy, docIds);
}
}
}
- public void testPartitionByLeafRandomized() throws IOException {
+ public void testRandomized() throws IOException {
for (int iter = 0; iter < 100; iter++) {
int numSegments = random().nextInt(10) + 1;
int totalDocs = 0;
@@ -152,39 +173,43 @@ public class TestReaderUtil extends LuceneTestCase {
List<LeafReaderContext> leaves = reader.leaves();
assertEquals(numSegments, leaves.size());
- // Generate random hits (0 to totalDocs inclusive - covers empty and all-match)
+ // Random hits (0 to totalDocs inclusive - covers empty and all-match).
int numHits = random().nextInt(totalDocs + 1);
Set<Integer> hitSet = new HashSet<>();
while (hitSet.size() < numHits) {
hitSet.add(random().nextInt(totalDocs));
}
int[] docIds = hitSet.stream().mapToInt(Integer::intValue).toArray();
- ScoreDoc[] hits = new ScoreDoc[docIds.length];
- for (int i = 0; i < docIds.length; i++) {
- hits[i] = new ScoreDoc(docIds[i], 1f);
- }
- int[][] result = ReaderUtil.partitionByLeaf(hits, leaves);
+ ReaderUtil.PartitionedHits partitioned = ReaderUtil.partitionByLeaf(docIds, leaves);
+ int[][] docIdsByLeaf = partitioned.docIdsByLeaf();
+ int[][] ordinalsByLeaf = partitioned.ordinalsByLeaf();
- // Verify: result length matches leaves
- assertEquals(numSegments, result.length);
+ assertEquals(numSegments, docIdsByLeaf.length);
+ assertEquals(numSegments, ordinalsByLeaf.length);
- // Verify: total hits preserved
- int totalResultDocs = Arrays.stream(result).mapToInt(a -> a.length).sum();
+ // Total hits preserved.
+ int totalResultDocs = Arrays.stream(docIdsByLeaf).mapToInt(a -> a.length).sum();
assertEquals(docIds.length, totalResultDocs);
- // Verify: each doc in correct leaf and sorted
- for (int leafIdx = 0; leafIdx < result.length; leafIdx++) {
- int[] leafDocs = result[leafIdx];
+ for (int leafIdx = 0; leafIdx < docIdsByLeaf.length; leafIdx++) {
+ int[] leafDocs = docIdsByLeaf[leafIdx];
+ int[] leafOrds = ordinalsByLeaf[leafIdx];
+ assertEquals(leafDocs.length, leafOrds.length);
LeafReaderContext leaf = leaves.get(leafIdx);
int docBase = leaf.docBase;
int maxDoc = leaf.reader().maxDoc();
+
for (int i = 0; i < leafDocs.length; i++) {
int docId = leafDocs[i];
+ // docId belongs to this leaf
assertTrue(docId >= docBase && docId < docBase + maxDoc);
+ // sorted ascending within leaf
if (i > 0) {
assertTrue(leafDocs[i] > leafDocs[i - 1]);
}
+ // ordinal points back to the original input
+ assertEquals(docIds[leafOrds[i]], docId);
}
}
}