Index sorting to reduce number of shards/segments accessed

[alexklibisz]

Copying (and slightly modifying) this idea from a comment for better visibility:

In September 2021 I attended Adrien Grand's talk "Speed up your Lucene queries by avoiding searching" at the 2021 Virtual ApacheCon.

Index Sorting was one of the methods that Adrien covered. In short, Index Sorting allows the user to customize how a Lucene index sorts documents stored in its segments. For example, an application might sort documents by date if they are always returned in chronological order. Elasticsearch also has an API for index sorting.

How does Index Sorting apply to vector search? My current hypothesis is that Index Sorting can be used to assign vectors to shards and segments more intelligently than a random assignment. At index time, we use index sorting to influence the placement of vectors into shards in a way that preserves spatial proximity. In other words, each shard is a cluster of spatially proximate vectors. At query time, the vector is still hashed via LSH, but now all or most of the hashes are stored in very few shards. So the query needs to access fewer shards in order to access and re-rank the candidates. Fewer shards means less IO, means faster query.

The initial experiment for Index Sorting involved sorting the vectors by two values:

1. Distance to from the vector to the origin (vector of all zeros).
2. Cosine of the angle between the vector and the origin vector.

Sorting the index by these two values led to a ~60% query-time speedup on one of the ten million vector datasets used for the 2021 "big ANN benchmarks" challenge. More testing and optimization is needed, but it looks like a fruitful direction.

The code was on the elastiknn-278-lucene-benchmarks branch, in the files LuceneAlgorithm.scala and LuceneStore.scala. I'll attach a zip of that directory just in case that branch goes missing at some point. elastiknn-elastiknn-278-lucene-benchmarks.tar.gz.

I also had a Gist with directions on running that mess of a branch:
d3e47e30d3c7468f3ee56844e5ee6f7a-ee25c17a896845c0c4fc63f64b44fc896fdd668c.zip

[alexklibisz]

Closing this as it's been moved to an issue: #628