Optimize Faiss Query With Filters: Reduce iteration and memory for id filter

[luyuncheng]

#1387 it discussed some optimization ideas about filters in ann search.

Problem Statement

we have met in some scenarios which match docids is high cardinality cause process OOM.

i think it is the original docid filters using an array object(jintArray filterIdsJ) to store all matched docids and transfer to jni, and then using an other bitmap/ to create IDSelectorBitmap/IDSelectorBatch do filter search.

these method cause many

• memory copy(like KNNWeight#createBitSet to create a bitmap and then copy it to array )

• iteration scan (like KNNWeight#bitSetToIntArray to iteration bitmap to a array; faiss_wrapper#bitSetToIntArray to iteration array to bitmap)

k-NN/src/main/java/org/opensearch/knn/index/query/KNNWeight.java

Lines 156 to 189 in 8c98265

	// If we already have a BitSet and no deletions, reuse the BitSet
	return ((BitSetIterator) filteredDocIdsIterator).getBitSet();
	}
	// Create a new BitSet from matching and live docs
	FilteredDocIdSetIterator filterIterator = new FilteredDocIdSetIterator(filteredDocIdsIterator) {
	@Override
	protected boolean match(int doc) {
	return liveDocs == null || liveDocs.get(doc);
	}
	};
	return BitSet.of(filterIterator, maxDoc);
	}
	private int[] getFilterIdsArray(final LeafReaderContext context) throws IOException {
	if (filterWeight == null) {
	return new int[0];
	}
	return bitSetToIntArray(getFilteredDocsBitSet(context, this.filterWeight));
	}
	private int[] getParentIdsArray(final LeafReaderContext context) throws IOException {
	if (knnQuery.getParentsFilter() == null) {
	return null;
	}
	return bitSetToIntArray(knnQuery.getParentsFilter().getBitSet(context));
	}
	private int[] bitSetToIntArray(final BitSet bitSet) {
	final int cardinality = bitSet.cardinality();
	final int[] intArray = new int[cardinality];
	final BitSetIterator bitSetIterator = new BitSetIterator(bitSet, cardinality);
	int index = 0;
	int docId = bitSetIterator.nextDoc();
	while (docId != DocIdSetIterator.NO_MORE_DOCS) {

• memory oom(like KNNWeight#getFilterIdsArray when bitmap is a high cardinality, it would OOM, because bitset iterator in JAVA, int array in JAVA, bitmap in JNI, array in JNI)

---

Proposal

in this pr i propose an idea to reduce these situations occurs.

• i use SparseFixedBitSet/FixedBitSet to do doExactSearch reduce one iteration/copy from bitset docid to array docid
• i use FixedBitSet long[] array pass direct to JNI, reduce memory copy and iteration translation. it can reduce iteration/copy from jni#filterIdsJ ->IDSelectorBitmap / IDSelectorBatch
• i move verify IDSelectorBitmap/IDSelectorBatch logic into JAVA.

[heemin32]

@luyuncheng Thanks for the contribution. Please rebase and resolve conflict.

@navneet1v Please take a look.

[navneet1v]

@luyuncheng I will look into the PR today. On a high level it seems like you have moved the code for using array or bitmaps based approach from C++ code to Java code to avoid OOM exceptions.

@heemin32 thanks for taking first pass.

[navneet1v]

we have met in some scenarios which match docids is high cardinality cause process OOM.

Can you add more details what are those cases? At what level of high cardinality we are facing the OOM issue?

[luyuncheng]

@luyuncheng I will look into the PR today. On a high level it seems like you have moved the code for using array or bitmaps based approach from C++ code to Java code to avoid OOM exceptions.

@navneet1v correct

Can you add more details what are those cases? At what level of high cardinality we are facing the OOM issue?

@navneet1v when in some scenario, we filter with some field like status, age, etc. the result docid array is very large. and shard merge into less segment, the docid number is 1<<30 and one docid is a long size(8 byte), so the memory in one thread are:

1. fixedbitset: 16MB，
2. the java transfer array is 8gb,
3. the jni array is 8gb (need copy from jvm to jni)
4. and also need to create another bitmap or array IDSelector. 16MB or 8gb

[navneet1v]

@luyuncheng I will look into the PR today. On a high level it seems like you have moved the code for using array or bitmaps based approach from C++ code to Java code to avoid OOM exceptions.

@navneet1v correct

Can you add more details what are those cases? At what level of high cardinality we are facing the OOM issue?

@navneet1v when in some scenario, we filter with some field like status, age, etc. the result docid array is very large. and shard merge into less segment, the docid number is 1<<30 and one docid is a long size(8 byte), so the memory in one thread are:

1. fixedbitset: 16MB，
2. the java transfer array is 8gb,
3. the jni array is 8gb (need copy from jvm to jni)
4. and also need to create another bitmap or array IDSelector. 16MB or 8gb

FYI lucene doc ids are ints and not longs. ref: https://lucene.apache.org/core/8_0_0/core/org/apache/lucene/search/DocIdSetIterator.html

Given than lucene per segment the max doc value can be INT_MAX - 1, the worst case a segment can have (2,147,483,647 - 1) documents which results to 8gb of array size. So are you saying that you are hitting cases when the segments have documents close to 2 Billion documents?

[luyuncheng]

Given than lucene per segment the max doc value can be INT_MAX - 1, the worst case a segment can have (2,147,483,647 - 1) documents which results to 8gb of array size. So are you saying that you are hitting cases when the segments have documents close to 2 Billion documents?

@navneet1v We hit some scenarios with 4gb segments (which is 1 << 30), but with concurrent search also with many indices this memory usage would be 4gb * thread number. AND 4gb is also in jvm and in jni memory

[navneet1v]

@luyuncheng can we get this last comment resolved. I am quite excited to see the impact of this change on the performance of filters in large cases.

[luyuncheng]

@luyuncheng can we get this last comment resolved.

DONE.

I am quite excited to see the impact of this change on the performance of filters in large cases.

@navneet1v Do you have any nightly benchmarks and datasets for the large filter case?

[navneet1v]

@luyuncheng can we get this last comment resolved.
DONE.

I am quite excited to see the impact of this change on the performance of filters in large cases.

@navneet1v Do you have any nightly benchmarks and datasets for the large filter case?

I have a dataset for testing with filters but it has only 1M documents in it. We can use that dataset with 1 shard and 1 segment to see what is impact before and after your change. That should provide us good insight on how much we are saving here.

[luyuncheng]

I have a dataset for testing with filters but it has only 1M documents in it. We can use that dataset with 1 shard and 1 segment to see what is impact before and after your change. That should provide us good insight on how much we are saving here.

LGTM, Thanks for the review !

[luyuncheng]

@navneet1v Let's try to close this PR quickly by providing prompt review to avoid another rebase.

@heemin32 @navneet1v can we make this PR finish? i think all comments are resolved

[navneet1v]

@navneet1v Let's try to close this PR quickly by providing prompt review to avoid another rebase.

@heemin32 @navneet1v can we make this PR finish? i think all comments are resolved

@luyuncheng this comment is marked as resolved but no changes are added, neither the resolution. Please fix this so that we can forward in PR

https://github.com/opensearch-project/k-NN/pull/1402/files#r1502838834

[luyuncheng]

@navneet1v Let's try to close this PR quickly by providing prompt review to avoid another rebase.

@heemin32 @navneet1v can we make this PR finish? i think all comments are resolved

@luyuncheng this comment is marked as resolved but no changes are added, neither the resolution. Please fix this so that we can forward in PR

https://github.com/opensearch-project/k-NN/pull/1402/files#r1502838834

will this really be cardinality * Long.BYTES * Long.BYTES

shouldn't this be cardinality * Integer.BYTES * Byte.SIZE

As Integer.BYTES : Provide the number of bytes in a int. And as lucene doc Ids are always int we should use Int here. Byte.SIZE : Represent the number of bits in a byte.

@navneet1v we use jlong array pass to jni#faiss_wrapper represent docid when selectorType is BATCH mode.

AND faiss#IDSelectorBatch using std::unordered_set<idx_t> set; filtered the docids which faiss::idx_t is int64_t in faiss#MetricType

So
Batch Memory = faiss::idx_t array in faiss = Cardinality * Long.BYTES

BitMap Memory = MaxId / 64 * Long.BYTES means using how many longs represent the words like FixedBitSet#bits2words and then multiple Long.Bytes for memory estimate

So the condition trans to Batch Memory < BitMap Memory ==> Cardinality * Long.BYTES / Long.BYTES * 64 < MaxId ==> Cardinality * Long.BYTES * Long.BYTES < MaxId (which Long.BYTES is 8 we can simplify 64)

@navneet1v i think there is no need to change because the explain are shows as comments deduction

[navneet1v]

@navneet1v Let's try to close this PR quickly by providing prompt review to avoid another rebase.

@heemin32 @navneet1v can we make this PR finish? i think all comments are resolved

@luyuncheng this comment is marked as resolved but no changes are added, neither the resolution. Please fix this so that we can forward in PR
https://github.com/opensearch-project/k-NN/pull/1402/files#r1502838834

will this really be cardinality * Long.BYTES * Long.BYTES
shouldn't this be cardinality * Integer.BYTES * Byte.SIZE
As Integer.BYTES : Provide the number of bytes in a int. And as lucene doc Ids are always int we should use Int here. Byte.SIZE : Represent the number of bits in a byte.

@navneet1v we use jlong array pass to jni#faiss_wrapper represent docid when selectorType is BATCH mode.
AND faiss#IDSelectorBatch using std::unordered_set<idx_t> set; filtered the docids which faiss::idx_t is int64_t in faiss#MetricType
So
Batch Memory = faiss::idx_t array in faiss = Cardinality * Long.BYTES
BitMap Memory = MaxId / 64 * Long.BYTES means using how many longs represent the words like FixedBitSet#bits2words and then multiple Long.Bytes for memory estimate
So the condition trans to Batch Memory < BitMap Memory ==> Cardinality * Long.BYTES / Long.BYTES * 64 < MaxId ==> Cardinality * Long.BYTES * Long.BYTES < MaxId (which Long.BYTES is 8 we can simplify 64)

@navneet1v i think there is no need to change because the explain are shows as comments deduction

@luyuncheng I think I got what you are saying. But we are using Long.Bytes to represent how many bits there in 1 byte. Can we atleast move towards this:

(cardinality * Long.BYTES * Byte.SIZE) <= filterIdsBitSet.length()
cardinality => number of docIds
Long.BYTES => number of bytes in long as Faiss uses Long to represent docIds.
Byte.SIZE => number of bits in a byte.

filterIdsBitSet.length() => number of bits in the filterIdsBitset.

This will ensure that logic is easy to understand for future readers.

[luyuncheng]

(cardinality * Long.BYTES * Byte.SIZE) <= filterIdsBitSet.length()
cardinality => number of docIds
Long.BYTES => number of bytes in long as Faiss uses Long to represent docIds.
Byte.SIZE => number of bits in a byte.

filterIdsBitSet.length() => number of bits in the filterIdsBitset.

This will ensure that logic is easy to understand for future readers.

@navneet1v LGTM, Update at f747ec7

[navneet1v]

@luyuncheng @heemin32 is on leave I have requested 1 more maintainer for approval.

[vamshin]

LGTM! Thanks