ReaderUtil.partitionByLeaf API design: ordinal tracking for scatter/gather pattern

[zihanx]

Description:

Following up on the recently merged ReaderUtil.partitionByLeaf API (#15507), we're planning a second task to build a generic utility class for efficiently retrieving doc-values for an int[] of global doc IDs.

The doc-values retriever would:

• Take an IndexReader, int[] globalDocIds, and DoubleValuesSource[] (or abstract LeafVisitor factory)

• Take an Executor/TaskExecutor for concurrent per-leaf retrieval

• Return values in 1:1 correspondence with the input doc ID order

This requires a scatter/gather pattern:

• Scatter: Partition global doc IDs by leaf (using partitionByLeaf)

• Map: Retrieve doc-values per leaf (concurrent)

• Gather: Reassemble results back to original input order

The current partitionByLeaf(ScoreDoc[], leaves) API doesn't track original positions, so we can't reassemble results in input order. We need ordinal tracking.

Design options:

1. Two separate APIs (current prototype):

// Existing - no ordinals (for callers who don't need gather and only take in ScoreDoc[])
public static int[][] partitionByLeaf(ScoreDoc[] hits, List<LeafReaderContext> leaves)

// New - with ordinals for scatter/gather
public record PartitionedHits(int[][] docIdsByLeaf, int[] ordinals) {}
public static PartitionedHits partitionByLeafWithOrdinals(int[] globalDocIds, List<LeafReaderContext> leaves)

// Private - helper method for duplicate partition logic 
private static int[][] partitionSortedDocIds(int[] sortedDocIds, List<LeafReaderContext> leaves) 

Pros: Clear intent, no wasted compute, type-safe

Cons: Two entry points

2. Single API with boolean flag:

public record PartitionedHits(int[][] docIdsByLeaf, int[] ordinals) {}
public static PartitionedHits partitionByLeaf(..., boolean includeOrdinals)

Pros: Single entry point

Cons: ordinals is null when not requested, still need to handle different input types (ScoreDoc[] vs int[])

3. Always return ordinals:

public static PartitionedHits partitionByLeaf(int[] globalDocIds, List<LeafReaderContext> leaves)

Pros: Simplest API

Cons: Wastes compute for callers who don't need ordinals

Questions for the community:

1. Should we add an int[] input method? The merged PR only has ScoreDoc[] input
   (last PR comment), but the scatter/gather use case needs raw doc IDs.

2. Which API design is preferred for ordinal tracking? Or any better ideas?

[gsmiller]

I have a few thoughts but I'm more curious what others think. As a disclaimer I work with Zihan on Amazon's product search, but I'm not particularly close to this project (so I'll probably ask silly questions) :)

1. Between the first two options, I think two separate methods is the way to go (option 1). Adding a boolean to control behavior (option 2) has a bit of a bad smell IMO.
2. Is there a reason you want the newly-proposed partitionByLeafWithOrdinals method to take int[] instead of ScoreDoc[] as input (in option 1)? There's a bit of asymmetry that might be worth fleshing out. It feels like these two methods are conflating whether-or-not the caller has ScoreDoc[] or int[] as input and whether-or-not they need "ordinal tracking." This sort of carries into option 3 too. If you went with that approach, would you use ScoreDoc[] as input (as the current API does), or would you need to change it to int[]?
3. Would option 3 be as simple as leveraging IntroSorter to sort the docIDs and ordinals as parallel arrays [1]? I wonder how much overhead it would actually introduce to do this over the current Arrays#sort approach. I kind of think we should try this first and see if it actually results in a meaningful performance difference. So I think I'm in favor of starting with option 3. If there really is a meaningful performance impact, I'd be in favor of option 1 (but sorting out the discussion of ScoreDoc[] vs. int[] as input).

Also, I'm assuming the ordinals array you have in mind runs parallel to the sorted output but flattened across leaves? That's what I'd expect, but just confirming.

[1] Something like...

    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 = ordinals[i];
        ordinals[i] = ordinals[j];
        ordinals[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);

[gsmiller]

Hmm... actually option 3 might be worse than I thought. I put together a micro-benchmark (#15933) and keeping track of the parallel ordinal array adds quite a bit of overhead. Here's the raw output of the benchmark:

Benchmark                                       (numDocIds)  (numLeaves)   Mode  Cnt     Score    Error   Units
PartitionByLeafBenchmark.arraysSortOnly                 100            5  thrpt   15  1213.898 ± 23.702  ops/ms
PartitionByLeafBenchmark.arraysSortOnly                 100           50  thrpt   15  1146.033 ± 30.084  ops/ms
PartitionByLeafBenchmark.arraysSortOnly                 100          200  thrpt   15   849.219 ± 10.757  ops/ms
PartitionByLeafBenchmark.arraysSortOnly                1000            5  thrpt   15    84.469 ±  0.637  ops/ms
PartitionByLeafBenchmark.arraysSortOnly                1000           50  thrpt   15    65.685 ±  3.071  ops/ms
PartitionByLeafBenchmark.arraysSortOnly                1000          200  thrpt   15    88.932 ±  5.194  ops/ms
PartitionByLeafBenchmark.arraysSortOnly               10000            5  thrpt   15     5.031 ±  0.394  ops/ms
PartitionByLeafBenchmark.arraysSortOnly               10000           50  thrpt   15     5.562 ±  0.467  ops/ms
PartitionByLeafBenchmark.arraysSortOnly               10000          200  thrpt   15     5.691 ±  0.560  ops/ms
PartitionByLeafBenchmark.arraysSortOnly              100000            5  thrpt   15     0.265 ±  0.008  ops/ms
PartitionByLeafBenchmark.arraysSortOnly              100000           50  thrpt   15     0.265 ±  0.007  ops/ms
PartitionByLeafBenchmark.arraysSortOnly              100000          200  thrpt   15     0.265 ±  0.007  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals          100            5  thrpt   15   911.904 ± 31.288  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals          100           50  thrpt   15   866.220 ± 22.404  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals          100          200  thrpt   15   690.300 ± 17.029  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals         1000            5  thrpt   15    63.449 ±  3.054  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals         1000           50  thrpt   15    64.223 ±  2.953  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals         1000          200  thrpt   15    69.654 ±  0.704  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals        10000            5  thrpt   15     3.443 ±  0.234  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals        10000           50  thrpt   15     3.471 ±  0.283  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals        10000          200  thrpt   15     3.200 ±  0.322  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals       100000            5  thrpt   15     0.174 ±  0.006  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals       100000           50  thrpt   15     0.174 ±  0.005  ops/ms
PartitionByLeafBenchmark.introSortWithOrdinals       100000          200  thrpt   15     0.166 ±  0.011  ops/ms

And a summary of the results:

numDocIds	numLeaves	Arrays.sort (ops/ms)	IntroSorter+ordinals (ops/ms)	Overhead
100	5	1214	912	~25%
100	50	1146	866	~24%
100	200	849	690	~19%
1,000	5	84	63	~25%
1,000	50	66	64	~3%
1,000	200	89	70	~22%
10,000	5	5.0	3.4	~32%
10,000	50	5.6	3.5	~38%
10,000	200	5.7	3.2	~44%
100,000	5	0.265	0.174	~34%
100,000	50	0.265	0.174	~34%
100,000	200	0.265	0.166	~37%

That said... I wonder how much we'd care about this in practice? This operation would probably be done once in the handling of a given search, and I suspect the difference in performance here would really show up in a meaningful way relative to the other work being done. But maybe I'm wrong? I dunno. WDYT?

[zihanx]

Thanks @gsmiller for the feedback and the benchmark! Really appreciate your help!
To answer your questions:

1. ScoreDoc[] v.s. int[]: For the doc-retriever use case, the API should be a generic class that takes int[] globalDocIds as input not ScoreDoc[]. @mikemccand originally described the retriever class as:

A generic (reusable to open source) utility class to retrieve doc-values (DoubleValuesSource) for an int[] of global docids to return. API could be simple-ish, taking an open IndexReader, the int[] globalHits, and maybe DoubleValuesSource[]... Or, so other things (stored fields maybe) are supported, the API could be a bit more abstract, and caller passes in some sort of LeafVisitor factory lambda-a-jigger thingy

@mikemccand can you please help confirm this is the intended design? If so, even with Option 3, we'd still need two APIs (one for ScoreDoc[], one for int[]), which is essentially Option 1 (or we could say even worse because option 1 would not introduce overhead).

2. Overhead: Agree, the overhead might be noise compare to other work being done.

I think right now the problem we need to solve is what input type does the ordinal-tracking API need?
Our current state only has one api: partitionByLeaf(ScoreDoc[], leaves) → int[][], we still need to decide ScoreDoc[]/int[] input.
If we need int[] as the input for retriever class, I'm leaning towards option 1, which seems cleanest: two APIs with different input types and different return types based on use case.

[gsmiller]

With-respect-to ScoreDoc[] vs. int[], I agree with what you're saying about Option 1 vs. Option 3. Ultimately, I'm slightly concerned we're going to end up with four methods proposed by the end of this that model the different combinations of ScoreDoc[] and int[] as input across ordinals or no-ordinals :). Maybe it was a mistake to move from int[] to ScoreDoc[] in the previous API design. If int[] really is the more generic version of this, maybe the right thing is to ask callers to deal with translating their ScoreDoc[] to int[]. I think the earlier premise was that the common use-case was for callers to have ScoreDoc[], so it made sense to not expect them to translate. But stepping back a little, maybe we should consider this decision separately?

Just to make things complicated though (sorry!), did you consider an alternative approach that would just return ordinals as int[][]? So instead of sorting the actual ScoreDoc[] or int[] input, we'd provide ordinal references back into the original data. Maybe an example would be easier:

ScoreDoc[] hits = {new ScoreDoc(42, ..), new ScoreDoc(7, ..), new ScoreDoc(100, ..), new ScoreDoc(23, ..)};

int[][] ordinalsByLeaf = ReaderUtil.partitionByLeaf(hits, leaves);
// Suppose leaf 0 covers docs 0-49, leaf 1 covers docs 50-99, leaf 2 covers 100+
// ordinalsByLeaf[0] = {1, 3, 0}  → hits[1]=doc7, hits[3]=doc23, hits[0]=doc42 (sorted within leaf)
// ordinalsByLeaf[1] = {}
// ordinalsByLeaf[2] = {2}        → hits[2]=doc100

So this handles both the current use-case and the scatter-gather use-case I think in a single method? (Different topic on whether-or-not you need two versions for accepting ScoreDoc[] vs. accepting int[]). I think this will still underperform similarly as Option 3 based on the benchmarking I did earlier though. You still couldn't use Arrays#sort directly; you'd need need something like IntroSorter, and sorting would have to do two array dereferences to access each element for comparison (i.e., hits[ordinalsByLeaf[i]]). But maybe it avoids some data copying? It might be worth adding to the jmh benchmark I created just to see...