Optimize top-k counting for approximate queries

[alexklibisz]

Currently the biggest bottleneck at query time is the countHits method in MatchHashesAndScoreQuery. This counts the number of times each doc in the segment matches one of the query vector's hashes. https://github.com/alexklibisz/elastiknn/blob/master/elastiknn-lucene/src/main/java/org/apache/lucene/search/MatchHashesAndScoreQuery.java#L54-L73

AFAIK, Lucene is generally optimized for a small number of terms (e.g. the words in a search query). Elastiknn, on the other hand, can require retrieving doc IDs for tens or hundreds of terms (the hashes of a query vector).

It seems the main thing worth exploring is using a different PostingsFormat, or potentially implementing a custom one. Maybe there's a way to optimize the storage/access pattern for checking a larger number of terms? Maybe there's a way to simply wrap an existing postings format and offer some helpers that cache the expensive method calls?

Some specific numbers:

When running the MatchHashesAndScoreQueryPerformanceSuite, with 100k indexed vectors and 5k search vectors, the VisualVM sampler reports spending ~92% of its time in the countHits method:

When running the ContinuousBenchmark with the SIFT dataset (1M indexed vectors, 1k search vectors), the VisualVM sampler reports spending ~36% of the search thread time in countHits:

[mikemccand]

Hmm, why is Self time so high in your profiler output? What is countHits actually doing?

Is there any way to cluster multiple hashes into a single term during indexing? I.e. do certain subsets of hashes frequently co-occur at search time?

[alexklibisz]

Hi @mikemccand, thanks for the reply. As a side note, I've found many of your articles very helpful!

Hmm, why is Self time so high in your profiler output? What is countHits actually doing?

To clarify a bit, the first screenshot is of a benchmark that uses Lucene by itself, i.e. no Elasticsearch involved. The second is from inside the search thread in the Elasticsearch JVM. The self times in the first case is proportionally higher, but both are > 50% of the method runtime.

The countHits method is here: https://github.com/alexklibisz/elastiknn/blob/master/elastiknn-lucene/src/main/java/org/apache/lucene/search/MatchHashesAndScoreQuery.java#L54-L73

The rough outline of countHits is:

• Get an iterator for the terms at the field. There is a single field where the terms (hashes of a vector) are indexed.
• Instantiate an ArrayHitCounter. This is basically allocating a short[] with one entry for every doc in the segment. In this benchmark there are 1M docs in the segment. Each time a doc matches one of the hashes, increment the value at its docId in that array. This is in lieu of using a hashmap (docId -> count), which I found to be at least twice as slow as using an array, even when using primitive collections. More on that here: More efficient counter for query hits #156
• Iterate over the query vector hashes.
• For each hash, seek it in the terms iterator, retrieve its postings (docs that match that term).
• For each posting, call .increment(docId), which increments the array-based counter like described above.
• After those two loops, return the counter.

Perhaps it's taking a while to allocate the array-backed counter? I don't have a good sense for whether that should actually take such a significant amount of time. Even if it does, the only alternatives I can think of are even slower. I can try it out later with a dummy counter implementation that just throws away the counts to get a lower bound. Would appreciate any more input you have about this and #156 !

Is there any way to cluster multiple hashes into a single term during indexing? I.e. do certain subsets of hashes frequently co-occur at search time?

That would be a smart optimization, but I'm not sure there's a way to do that without having the entire dataset up-front, which doesn't fit the model of the plugin (index/update/delete vectors just like you would index/update/delete traditional docs).

[alexklibisz]

Some more notes-to-self for when I get back to this:

Here are the VisualVM hotspots from running the SIFT benchmark (1M stored vectors, 10k total queries) on a local ES JVM using the ArrayHitCounter:

Here are the VisualVM hotspots from the same configuration, but using a dummy counter that just returns the same count and maintains no internal state (i.e. no memory allocated/updated for counting):

I made sure to start the sampling right when it started running queries and stop right after it finished queries. Otherwise the percentages get skewed as the search thread continues running.

[mikemccand]

Hi @mikemccand, thanks for the reply. As a side note, I've found many of your articles very helpful!

Thanks, I am glad to hear that :)

The countHits method looks fine, though you should not iterate to docs.cost() but rather iterate until you see NO_MORE_DOCS returned. In general the cost method is not guaranteed to be a precise count of how many hits this PostingsEnum will return, though for the default Codec that appears to be the case for now.

You are reusing your TermsEnum and PostingsEnum which is great. You might want to pre-sort the Hash in unicode order ... it might give a tiny improvement since the TermsEnum can share more internal state on each seekExact.

Your countHits method is very similar to what Lucene's MultiTermQuery.rewrite does when there are many terms for the non-scoring case, except that method accumulates into a Lucene bitset rather than int[] since the Lucene Query is just concerned with knowing which documents matched, not how many terms matched each hit.

I can try it out later with a dummy counter implementation that just throws away the counts to get a lower bound.

Measuring just the iteration and no count incrementing is a good idea. Oh I see, you just concurrently did that, as I am here commenting, heh. Oddly, it looks like Self time increased with the dummy counter (28.7% to 31.4%)? And the overall time dropped a bit ... this is what I would expect: Lucene is working hard to iterate through all those matched docids.

That would be a smart optimization, but I'm not sure there's a way to do that without having the entire dataset up-front, which doesn't fit the model of the plugin (index/update/delete vectors just like you would index/update/delete traditional docs).

Yeah, it is not simple to implement, and would require that you have some way to tell which hashes frequently co-occur during indexing. It would also increase your index size, but then make queries faster. You might also index larger and larger co-occurrence sets, so that at search time you could "factor out" the largest set(s) that can apply to that query.

How large is each vector hash?

[alexklibisz]

Thanks again for digging into this a bit.

The countHits method looks fine, though you should not iterate to docs.cost()

Good to know, I'll fix that.

You might want to pre-sort the Hash in unicode order ... it might give a tiny improvement since the TermsEnum can share more internal state on each seekExact.

I might already be doing this in the query constructor, here
The comparator just uses Java's Arrays.compareUnsigned, which, IIRC, is what Lucene does for ByteRefs. Is that what you are suggesting or something different?

Your countHits method is very similar to what Lucene's MultiTermQuery.rewrite does when there are many terms for the non-scoring case, except that method accumulates into a Lucene bitset rather than int[] since the Lucene Query is just concerned with knowing which documents matched, not how many terms matched each hit.

I modeled the implementation after the TermsInSetQuery. Which, similarly, only tracks whether a doc contained a matching term, not the actual count. It seems like the main difference and probably the main bottleneck between this custom query and prior art in Lucene is that I actually need to know how many matches there are in each doc. If you have any suggestions/tricks for very fast counting on the JVM, I'm all ears. As mentioned before, hashmaps, even from primitive collection libraries, seem to be about 2x slower.

How large is each vector hash?

With this benchmark config, there are 100 hashes per vector, 20 bytes each. I'd say that's a pretty typical setup.

[alexklibisz]

Maybe there's a clever stopping criteria to visiting all of the terms? I started reading about MaxScore and WAND scoring. Maybe that's dead end here?

[alexklibisz]

Put together a working early-stopping solution today. Roughly the idea is:

• Compute the total number of term hits up front by iterating over the terms enum once.
• Iterate over the terms and docs for each term again.
• Use a counter to track the hits for each doc and a heap to track the top k hit counts.
• For every doc, decrement the total number of term hits by the number of hits in that doc.
• Stop iterating when the kth largest count is larger than the remaining number of term hits. Even if all of the remaining hits were for a single doc, it still couldn't exceed the kth largest.

Seems to work on the benchmark discussed previously. The times are roughly equivalent, but there's some obvious room for optimization. The method spends more time managing the heap than it does accessing lucene postings. I'm using a standard PriorityQueue<Integer>, and I bet if I wrote one using primitives it would be faster. The total runtime spends 10% of its time computing the kth largest hit from the array-backed counter. I just have to change some internal classes and methods to be able to use the kth largest from the heap instead of re-computing it.

[mikemccand]

You might want to pre-sort the Hash in unicode order ... it might give a tiny improvement since the TermsEnum can share more internal state on each seekExact.

I might already be doing this in the query constructor, here
The comparator just uses Java's Arrays.compareUnsigned, which, IIRC, is what Lucene does for ByteRefs. Is that what you are suggesting or something different?

compareUnsigned is perfect and matches Lucene's term sort order (Unicode sort order on utf-8 encoded byte[]).

Put together a working early-stopping solution today. Roughly the idea is:

That sounds promising! Do you take the docFreq (or maybe totalTermFreq) of terms into account? E.g., collecting all term + docFreq from the TermsEnum, then sort them in increasing docFreq order? That way your early termination would save the most work (by skipping the terms that had many docs)? Though, maybe your terms (vector hashes) are somewhat uniform and their docFreq do not vary much.

Be sure to test on real vectors. Testing on random data leads to random conclusions!

[alexklibisz]

That sounds promising! Do you take the docFreq (or maybe totalTermFreq) of terms into account? E.g., collecting all term + docFreq from the TermsEnum, then sort them in increasing docFreq order? That way your early termination would save the most work (by skipping the terms that had many docs)? Though, maybe your terms (vector hashes) are somewhat uniform and their docFreq do not vary much.

I thought about this but my understanding was that Lucene is optimized to seek over the terms in sorted order? Maybe the cost of unsorted seeks over terms is trivial compared to other things here?

[alexklibisz]

I'm afraid the early-stopping method as I described it isn't going to work. Specifically, it's pretty easy to find a case where a single vector matches for multiple consecutive hash terms and fills up the heap with its incrementing match counts. You end up determining that the cutoff for match counts is an old match count for that doc and you early-stop, but then there's only one doc that exceeds that cutoff. Replacing a doc's previous count in the heap might solve this?
Edit: I got it to work by clearing the heap after each term. If the smallest count in the heap for a specific term exceeds the remaining possible counts, then you early stop. But this ends up still being slower than just counting all the docs. Setting this aside for a few days.

[mikemccand]

That sounds promising! Do you take the docFreq (or maybe totalTermFreq) of terms into account? E.g., collecting all term + docFreq from the TermsEnum, then sort them in increasing docFreq order? That way your early termination would save the most work (by skipping the terms that had many docs)? Though, maybe your terms (vector hashes) are somewhat uniform and their docFreq do not vary much.

I thought about this but my understanding was that Lucene is optimized to seek over the terms in sorted order? Maybe the cost of unsorted seeks over terms is trivial compared to other things here?

You are right, Lucene would prefer that you seek the terms in Unicode order.

But it allows you to NOT do so, and the performance penalty might be lower than the gains you could see by early terminating sooner?

Just pure speculation ... I think you would need to test and see to know for sure.