Ways to implement prefetching in BKD trees

[abiesps]

I see that Lucene has added prefetching #13179 in number of data structures to decouple search concurrency with IO concurrency.

In the linked issue, there is an unchecked item on prefetching support to points. I want to know if lucene is thinking of adding prefetching support to BKD trees in near future.

I was also exploring the ways to achieve this. At a high level, I was thinking that for range queries, point-lookup etc we can traverse the BKD tree in same fashion as we do. But instead of visiting leaf nodes (kdd files). We can call prefetch on matching leaf nodes file pointers. Iterator can also save the matching leaf nodes during traversal.
Once the traversal of BKD tree is complete, we can actually make a call to visit doc IDs or doc Values from the matching leaf nodes. Hopefully they may be prefetched by this time.

Opensearch uses similar variation of BKD tree traversal for range queries with an exception of early termination i.e if the matching doc id count becomes greater than equal to the size parameter requested in the query we stop the traversal.
For such use cases to support early termination, I was thinking of adding a vInt to inner nodes (kdi files) that denotes the total number of documents below that node.

Let me know if this make sense or what does community feel about this.

[abiesps]

I dont think we need to store matching doc count in inner nodes, probably can get it from size(). Exploring this more for my poc.

[abiesps]

Another thing that I realised, especially relevant for early termination is :
leafNode may not always exactly match thats why we have two functions in BKDReader visitDocIDs (for the case of cell inside query) and visitDocValues (for the case of cell crossing the query). visitDocValues is the case where leaf node may not match in case of more than one dimension in BKD tree. And we know for sure by actually reading the leaf Node (I/O). Whereas in case of visitDocIDs, all document in the leaf always matches the range in the query.

If we don't want early termination - then prefetching leaf nodes from both kind of matches is fine, otherwise we can keep on incrementing match count only in case of "r == PointValues.Relation.CELL_INSIDE_QUERY". With this we may do some wasteful traversal of BKD tree but saving on doing IO in query path (through prefetching leaf nodes).

And size() returns matching points not necessarily doc count.

With above consideration, let me explore more on this and raise a draft PR in couple of days or so.

[abiesps]

Ok, I got a working solution. Code is really in very early stage (I need to clean up a lot) - I need more time to raise a PR.

Approach

At a high level, this is what I am doing:

1. Start traversing the BKD tree as is,

2. If cell is inside the query
   a) Call goes to visitDocIDs method variation where if isLeaf is true,
   b) Do not read the leaf node yet. Instead, get the leafOrdinal, if this leaf ordinal is in continuation to last matching leaf ordinal (which I am storing in visitor) i.e leafOrdinal == visitor.lastMatchingLeafOrdinal() + 1. Do not call prefetch as this should be hopefully taken care by kernel readaheads. Otherwise if its not a continuous match i.e leafOrdinal != visitor.lastMatchingLeafOrdinal() + 1, then call prefetch on this leaf node file pointer's first page. Also store this leaf node fp in visitor for visting matching doc IDs later. For early termination, also pass number of matching points to visitor from "int count = isLastLeaf() ? lastLeafNodePointCount : config.maxPointsInLeafNode();"
   c) If its not a leafNode continue with recursion as is.

3. For remaining traversal code remains unchanged.

4. Once the traversal is complete, I am calling visitDocIDs on the leaf file pointers I stored for visting later on.

Benchmark Results

I was able to do a cold index test, this is the test that I wrote (https://github.com/abiesps/lucene-learnings/blob/main/src/main/java/com/sps/lucene/learnings/LuceneBKDTraversalPrefetchBenchmark.java), and I am seeing following results

Iteration	p50 - with prefetching (nanos)	p50 - without prefetching (nanos)	p90 - with prefetching (nanos)	p90 - without prefetching (nanos)	p99 - with prefetching (nanos)	p99 - without prefetching (nanos)
1	1428747	1873421	3204614	4011134	6223980	7440390
2	1423662	1868136	3257577	3999692	6045371	7452711
3	1431170	1881138	3286315	4026663	6614156	7655345
4	1463966	1882202	3351880	4031743	6434328	7826272
5	1448586	1860455	3253163	3993159	5995588	7826855
6	1451429	1882752	3235150	4013533	5681142	7661777
7	1447876	1840063	3301089	4025598	6419166	7884995
8	1427931	1835098	3188111	4007786	6166556	7742863
9	1473577	1868620	3358062	4106067	7045409	8093894
10	1444055	1863376	3342137	3995365	6286142	7813705
Avg across all iteration	1444099.9	1865526.1	3277809.8	4021074	6291183.8	7739880.7

I ran 1000 range queries, for each iteration with explicitly clearing the page cache between runs for with and without prefetching.

[abiesps]

As a next step, I want to do a test on Hot path (when all data is present in the page cache) and see the results with this approach (If results look ok/good) I'll raise a PR.