[FP16 Default] Enable optimistic search in memory optimized search

[0ctopus13prime]

Description

Recently Lucene has added optimistic search in place of early termination enabled parallel search - apache/lucene#14226
This is not yet available in memory optimized search, which needs to be pulled in.
From the coarse performance benchmark with Cohere-10M with FP32, optimistic search is showing 10% better QPS than early termination based parallel search.

[0ctopus13prime]

SIMD enabled memory optimized search with optimistic search seems showing slightly better performance and recall than early termination based method in multi segments scenario. Recall went up from 93% to 94% and mean QPS went up from 347 to 367 (5% improvement)
Note that Faiss C++ is taking 161.88 ops/s for the same workloads.

• Data : cohere-10m
• Segments : 126
• Instance : r7g.4xlarge, CPU cores=16, memory=128GB
• Data type : FP16

SIMD + Memory optimized search + Early termination

|                                                 Min Throughput |   prod-queries |      306.12 |  ops/s |
|                                                Mean Throughput |   prod-queries |      341.07 |  ops/s |
|                                              Median Throughput |   prod-queries |      343.28 |  ops/s |
|                                                 Max Throughput |   prod-queries |      345.76 |  ops/s |
|                                        50th percentile latency |   prod-queries |     54.2734 |     ms |
|                                        90th percentile latency |   prod-queries |     77.5883 |     ms |
|                                        99th percentile latency |   prod-queries |     101.928 |     ms |
|                                      99.9th percentile latency |   prod-queries |     126.095 |     ms |
|                                     99.99th percentile latency |   prod-queries |     147.692 |     ms |
|                                       100th percentile latency |   prod-queries |     153.879 |     ms |
|                                   50th percentile service time |   prod-queries |     54.2734 |     ms |
|                                   90th percentile service time |   prod-queries |     77.5883 |     ms |
|                                   99th percentile service time |   prod-queries |     101.928 |     ms |
|                                 99.9th percentile service time |   prod-queries |     126.095 |     ms |
|                                99.99th percentile service time |   prod-queries |     147.692 |     ms |
|                                  100th percentile service time |   prod-queries |     153.879 |     ms |
|                                                     error rate |   prod-queries |           0 |      % |
|                                                  Mean recall@k |   prod-queries |        0.93 |        |
|                                                  Mean recall@1 |   prod-queries |        0.97 |        |

SIMD + Memory optimized search + Optimistic search

|                                                 Min Throughput |   prod-queries |      335.85 |  ops/s |
|                                                Mean Throughput |   prod-queries |      367.12 |  ops/s |
|                                              Median Throughput |   prod-queries |         371 |  ops/s |
|                                                 Max Throughput |   prod-queries |      377.84 |  ops/s |
|                                        50th percentile latency |   prod-queries |      50.641 |     ms |
|                                        90th percentile latency |   prod-queries |     57.4001 |     ms |
|                                        99th percentile latency |   prod-queries |     90.2634 |     ms |
|                                      99.9th percentile latency |   prod-queries |     104.845 |     ms |
|                                     99.99th percentile latency |   prod-queries |     124.002 |     ms |
|                                       100th percentile latency |   prod-queries |      148.17 |     ms |
|                                   50th percentile service time |   prod-queries |      50.641 |     ms |
|                                   90th percentile service time |   prod-queries |     57.4001 |     ms |
|                                   99th percentile service time |   prod-queries |     90.2634 |     ms |
|                                 99.9th percentile service time |   prod-queries |     104.845 |     ms |
|                                99.99th percentile service time |   prod-queries |     124.002 |     ms |
|                                  100th percentile service time |   prod-queries |      148.17 |     ms |
|                                                     error rate |   prod-queries |           0 |      % |
|                                                  Mean recall@k |   prod-queries |        0.94 |        |
|                                                  Mean recall@1 |   prod-queries |        0.97 |        |

[0ctopus13prime]

The tricky part of the implementation is getting an iterator from VectorReader.

The optimistic search consists of two phases. In the first phase, a k-NN search is run on each segment with an adjusted k value, determined by this formula
. The idea is to assign a larger k value to larger segments.

In the second phase, only the segments whose minimum score is greater than or equal to the kth score from the merged results are searched again. This time, the vector search uses the document ids collected in the first phase as seeds.

The problem is that document IDs don’t necessarily match internal vector IDs. In sparse or nested cases, they can differ. So we can’t directly use the collected document IDs from the first phase as seeds for the second. We need an iterator to map document IDs to their corresponding vector IDs.

The issue is that NativeEngines990KnnVectorsReader returns FloatVectorValues from Lucene’s flat vector reader, as shown below:

@Override
public FloatVectorValues getFloatVectorValues(final String field) throws IOException {
    // Return iterator from Lucene's flat vector! Note that `FloatVectorValues` is a sort of iterator.
    return flatVectorsReader.getFloatVectorValues(field);
}

However, what we actually need is an iterator that provides both document IDs and vector IDs from the Faiss index. Instead, we’re currently getting an iterator over Lucene’s flat vectors.

I think we can fix this, but even after doing so, we’ll still need to implement iterator getters for all Faiss flat vector types (e.g., FP16, FP32, and binary). The amount of work is turning out to be bigger than I expected. 😅