zihanx commented on code in PR #16496:
URL: https://github.com/apache/lucene/pull/16496#discussion_r3778757452
##########
lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java:
##########
@@ -95,32 +94,81 @@ public static int subIndex(int n, List<LeafReaderContext>
leaves) {
}
/**
- * 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()} — tracking the
+ * ordinals is cheap enough (see the {@code PartitionByLeafBenchmark}) that
a separate no-ordinals
+ * method is not worth maintaining.
+ *
+ * <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);
+ }
+
+ // Pack (docId, ordinal) into a single long -- docId in the high 32 bits,
the original input
+ // position in the low 32 bits -- then sort as primitives. Sorting
ascending orders by docId,
+ // with the ordinal as a tiebreak (docIds are unique here, so the tiebreak
never applies). This
+ // relies on both values being non-negative (Lucene doc IDs and array
indices), so the sign bit
+ // is always clear and signed long order matches ascending docId order.
Packing lets us use the
+ // tuned primitive Arrays.sort(long[]) with no per-comparison callbacks
and a single contiguous
+ // array, which is faster than a comparator/IntroSorter over parallel
int[]s.
+ final long[] packed = new long[globalDocIds.length];
+ for (int i = 0; i < packed.length; i++) {
+ packed[i] = ((long) globalDocIds[i] << 32) | (i & 0xFFFFFFFFL);
}
- int[] sortedDocIds = new int[hits.length];
- for (int i = 0; i < hits.length; i++) {
- sortedDocIds[i] = hits[i].doc;
+ Arrays.sort(packed);
+ final int[] sortedDocIds = new int[packed.length];
+ final int[] sortedOrdinals = new int[packed.length];
+ for (int i = 0; i < packed.length; i++) {
Review Comment:
You are right, we should directly decode in the loop. There's no need for
the extra allocation, will change it in the next commit. Thanks!
--
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]
