Searching an index over a subset of coordinates

[riv20]

I'm working on a problem that involves carrying out many k-NN searches over different coordinate subsets of an index set X.

Concretely, imagine an array of vectors X; I would like to perform a k-NN search on X[:,I] for many (i.e. hundreds of) different choices of I.

Currently I am simply loading X[:,I] to an index each time. On GPU, the overhead incurred by copying many index sets to GPU memory more than cancels out the speedup that comes with using GPU.

Obviously repeatedly copying X[:,I] to memory is wasteful, and in my case the entire X is not too big to fit entirely in memory.

Ideally, I would like to be able to get rid of this overhead by writing something like this:

index.add(X)
index.search(xq, k, coords = I) # find the k-NNs of xq in X[:,I]

Is something like this possible with FAISS?

Many thanks.