zihanx commented on code in PR #16496:
URL: https://github.com/apache/lucene/pull/16496#discussion_r3769236437
##########
lucene/core/src/java/org/apache/lucene/index/ReaderUtil.java:
##########
@@ -95,25 +95,122 @@ 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.
+ * Partitions global doc IDs by leaf. Doc IDs may be supplied in any order;
the returned per-leaf
+ * arrays are sorted in ascending docId order.
*
- * @param hits the ScoreDoc array (typically from TopDocs.scoreDocs)
+ * <p>This is an optimized subset of {@link
#partitionByLeafWithOrdinals(int[], List)} for callers
+ * that only need the per-leaf grouping and do not need to map results back
to the original input
+ * order. It sorts with {@link Arrays#sort(int[])} and skips the extra
bookkeeping required to
+ * track input ordinals. Callers that need to reassemble per-leaf results
into input order
+ * (scatter/gather) should use {@link #partitionByLeafWithOrdinals(int[],
List)} instead.
+ *
+ * <p>The input array is not mutated.
+ *
+ * @param globalDocIds global doc IDs in any 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 array indexed by leaf ord, containing the (sorted) global doc IDs
for that leaf (empty
+ * if no hits land in that leaf)
*/
- public static int[][] partitionByLeaf(ScoreDoc[] hits,
List<LeafReaderContext> leaves) {
+ public static int[][] partitionByLeaf(int[] globalDocIds,
List<LeafReaderContext> leaves) {
int numLeaves = leaves.size();
- int[][] result = new int[numLeaves][];
- if (hits.length == 0) {
+ if (globalDocIds.length == 0) {
+ int[][] result = new int[numLeaves][];
Arrays.fill(result, EMPTY_INT_ARRAY);
return result;
}
- int[] sortedDocIds = new int[hits.length];
- for (int i = 0; i < hits.length; i++) {
- sortedDocIds[i] = hits[i].doc;
- }
+ int[] sortedDocIds = globalDocIds.clone();
Arrays.sort(sortedDocIds);
+ return partitionSortedDocIds(sortedDocIds, 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]}.
+ *
+ * <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>This is the fuller-featured counterpart to {@link
#partitionByLeaf(int[], List)}: it returns
+ * the same per-leaf grouping and additionally records, for every
partitioned doc ID, its index in
+ * the original {@code globalDocIds} array. Tracking these ordinals carries
a small amount of
+ * extra work relative to {@link #partitionByLeaf(int[], List)}, so callers
that do not need to
+ * map results back to input order should prefer that method.
+ *
+ * <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 per-leaf sorted doc IDs alongside per-leaf ordinals into the
input array
+ */
+ public static PartitionedHits partitionByLeafWithOrdinals(
+ int[] globalDocIds, List<LeafReaderContext> leaves) {
+ int numLeaves = leaves.size();
+ 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);
+ }
+
+ // Sort doc IDs and ordinals as parallel arrays, so we keep the original
positions while
+ // moving doc IDs into ascending order. IntroSorter avoids the
boxing/lambda overhead a
+ // comparator-based Arrays.sort would incur on parallel int[]s.
+ final int[] sortedDocIds = globalDocIds.clone();
+ final int[] sortedOrdinals = new int[globalDocIds.length];
+ for (int i = 0; i < sortedOrdinals.length; i++) {
+ sortedOrdinals[i] = i;
+ }
+ new IntroSorter() {
Review Comment:
Sure! Feel free to merge it.
--
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]
