gsmiller commented on code in PR #16496:
URL: https://github.com/apache/lucene/pull/16496#discussion_r3770752122
##########
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:
Cool, just pushed those changes (should only have touched the benchmark
code, not any of your actual code). Also posting the results I got from running
the benchmark here. It's really interesting (to me at least) to see the
benchmark validate the merit of the packed-long approach vs. intro-sort. It's
also interesting to see how much it closes the perf gap to not doing any
ordinal tracking. I'd be curious what you think of keeping both APIs vs.
consolidating into one (essentially not worrying about the performance penalty
of ordinal tracking for cases where ordinals aren't needed in favor of one
API). Also please feel free to run the benchmarks and see if you're getting
similar results, forming different conclusions, etc., etc. :) Thanks!
## Raw results
| numDocIds | numLeaves | partitionByLeaf (plain) |
partitionByLeafWithOrdinals (packed) | …IntroSort |
|---:|---:|---:|---:|---:|
| 100 | 10 | 1486.7 ± 243.8 | 1188.3 ± 81.3 | 1134.1 ± 63.7 |
| 100 | 200 | 594.1 ± 49.9 | 570.3 ± 8.4 | 433.5 ± 8.1 |
| 1000 | 10 | 89.52 ± 41.50 | 92.03 ± 10.96 | 69.28 ± 2.13 |
| 1000 | 200 | 73.71 ± 7.86 | 65.14 ± 9.20 | 51.59 ± 2.21 |
| 10000 | 10 | 5.095 ± 1.321 | 5.885 ± 0.986 | 2.750 ± 0.321 |
| 10000 | 200 | 5.545 ± 2.494 | 5.916 ± 0.499 | 3.194 ± 0.587 |
| 100000 | 10 | 0.242 ± 0.015 | 0.243 ± 0.010 | 0.168 ± 0.006 |
| 100000 | 200 | 0.251 ± 0.007 | 0.245 ± 0.004 | 0.167 ± 0.006 |
## Packed-long vs. IntroSort
| numDocIds | numLeaves | Packed speedup |
|---:|---:|---:|
| 100 | 10 | 1.05× |
| 100 | 200 | 1.32× |
| 1000 | 10 | 1.33× |
| 1000 | 200 | 1.26× |
| 10000 | 10 | 2.14× |
| 10000 | 200 | 1.85× |
| 100000 | 10 | 1.45× |
| 100000 | 200 | 1.47× |
## Ordinal overhead vs the plain no-ordinals API
(Overhead = (plain − variant) ÷ plain)
| numDocIds | numLeaves | Packed overhead | IntroSort overhead |
|---:|---:|---:|---:|
| 100 | 10 | +20.1% | +23.7% |
| 100 | 200 | +4.0% | +27.0% |
| 1000 | 10 | -2.8% | +22.6% |
| 1000 | 200 | +11.6% | +30.0% |
| 10000 | 10 | -15.5% | +46.0% |
| 10000 | 200 | -6.7% | +42.4% |
| 100000 | 10 | -0.4% | +30.6% |
| 100000 | 200 | +2.4% | +33.5% |
## Raw JMH output
```
Benchmark (numDocIds)
(numLeaves) Mode Cnt Score Error Units
PartitionByLeafBenchmark.partitionByLeaf 100
10 thrpt 8 1486.740 ± 243.770 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 100
200 thrpt 8 594.129 ± 49.858 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 1000
10 thrpt 8 89.516 ± 41.496 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 1000
200 thrpt 8 73.707 ± 7.861 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 10000
10 thrpt 8 5.095 ± 1.321 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 10000
200 thrpt 8 5.545 ± 2.494 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 100000
10 thrpt 8 0.242 ± 0.015 ops/ms
PartitionByLeafBenchmark.partitionByLeaf 100000
200 thrpt 8 0.251 ± 0.007 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 100
10 thrpt 8 1188.346 ± 81.269 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 100
200 thrpt 8 570.284 ± 8.370 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 1000
10 thrpt 8 92.032 ± 10.960 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 1000
200 thrpt 8 65.144 ± 9.195 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 10000
10 thrpt 8 5.885 ± 0.986 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 10000
200 thrpt 8 5.916 ± 0.499 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 100000
10 thrpt 8 0.243 ± 0.010 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinals 100000
200 thrpt 8 0.245 ± 0.004 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 100
10 thrpt 8 1134.130 ± 63.681 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 100
200 thrpt 8 433.474 ± 8.146 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 1000
10 thrpt 8 69.280 ± 2.132 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 1000
200 thrpt 8 51.595 ± 2.211 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 10000
10 thrpt 8 2.750 ± 0.321 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 10000
200 thrpt 8 3.194 ± 0.587 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 100000
10 thrpt 8 0.168 ± 0.006 ops/ms
PartitionByLeafBenchmark.partitionByLeafWithOrdinalsIntroSort 100000
200 thrpt 8 0.167 ± 0.006 ops/ms
```
--
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]
