Python support

[flying-sheep]

Hi, “Python support” is listed under “Missing features”.

Would you accept a PR that adds it?

cc @ilan-gold

[ilan-gold]

Just for context: our use-case would involve using dask as a compute engine for some pre-neighbors tasks, and then dumping the results block-by-block into an on-disk index. We could then create a new dask array that reads from that on-disk index the n-nearest neighbors once it is built.

Just want to make sure I understand what this library does in the context of our need. If this understanding would be correct, then python bindings for us would be a must.

[irevoire]

Uuuh, not sure I understood your use case.

This library stores vectors (also called embeddings).
Then, you have to build the database. Once this is done, you can provide your own vector and retrieve the nearest vectors in the database.

Is this what you need?

[ilan-gold]

I smashed the two steps in the middle together (storing vectors, building the database), but yes that is exactly what we need.

We would take a very large PCA, store its results on disk using dask (lazily and row-chunk by row-chunk), build the database (outside of dask), and then create a new matrix where each observation has its n-nearest neighbors in the columns marked.

Concretely, my vision would be to call add_item in every chunk, row-by-row, in parallel, from python (roughly speaking, probably makes sense to pass in the chunk to rust and then call add_item within rust from a for-loop/iterator). Then once that is completed, we would build outside of dask, called from python. Finally we would use the Reader to build up the nearest-neighbors matrix (lazily, again using dask).

[flying-sheep]

Hmm, to implement this properly, we’d accept distance as a runtime parameter, but Distance doesn’t seem to be dyn compatible.

How would you recommend creating a database with a distance type only determined at runtime?

[irevoire]

Concretely, my vision would be to call add_item in every chunk, row-by-row, in parallel, from python (roughly speaking, probably makes sense to pass in the chunk to rust and then call add_item within rust from a for-loop/iterator).

For this part, take care, you cannot write multiple thing in parallel in an LMDB database. You'll get a runtime error if you try.
So either you do everything on one thread, or you get a "writer" thread that holds the database, and the other are just pushing stuff in a channel or something.

Then once that is completed, we would build outside of dask, called from python. Finally we would use the Reader to build up the nearest-neighbors matrix (lazily, again using dask).

That should work 👍

Hmm, to implement this properly, we’d accept distance as a runtime parameter, but Distance doesn’t seem to be dyn compatible.

Yeah, I must admit we didn't take the best decision in terms of usage 😭
Currently, the only way to do it, as far as I know, is to wrap the Database type of head directly with the runtime param that represents the distance in a struct, and then every time you must do an operation, you start by matching the runtime param and then create the associated Writer and do your operation.

That's what I did here in Meilisearch except I only handle binary quantized distances or not.

---

I never tried it, but IMO you should also be able to wrap all this logic in a new "RuntimeDistance" trait that stores the dynamic distance in one of its fields and do the match in all of its methods.
But I'm not sure the Distance trait can hold data, so maybe it's not feasible that way, I don't remember.

[flying-sheep]

IMO you should also be able to wrap all this logic in a new "RuntimeDistance" trait that stores the dynamic distance in one of its fields and do the match in all of its methods.

I thought about that, but would that work with the associated types necessary to implement Distance?

Also it would make more sense performance-wise to do the match once and then hold a reference to the implementation and have it do stuff. But as said: not dyn-compatible …

[irevoire]

I thought about that, but would that work with the associated types necessary to implement Distance?

I don't think so, to be honest, but it could be worth trying.

Also it would make more sense performance-wise to do the match once and then hold a reference to the implementation and have it do stuff. But as said: not dyn-compatible …

Yep, not ideal, sorry, and this idea of statically dispatching everything is used everywhere so much in arroy that I don't even see how we could change it today without rewriting half of the crate, to be honest 😭

[ilan-gold]

For this part, take care, you cannot write multiple thing in parallel in an LMDB database. You'll get a runtime error if you try.
So either you do everything on one thread, or you get a "writer" thread that holds the database, and the other are just pushing stuff in a channel or something.

That's a shame, but also part of the reason I posted detail here because I wasn't sure. Is there an option to merge databases (before actually building the index) or something similar?¹ Or is the process of building the database itself somehow tied to having access to all of the data? The option of having a reserved way of writing I'm sure will be plenty fast, just curious about other options.

Footnotes

1. Looking at FAISS' documentation gives me a headache (a me problem) but they appear to have this concept of merging: https://github.com/facebookresearch/faiss/tree/main/benchs/distributed_ondisk#splitting-across-inverted-lists although it appears to be for merging already built indices. Just thinking out loud ↩

[ilan-gold]

Quick update: I spoke with Phil separately, and if we can't do the above on-disk merge idea, we can handle ensuring that only one process does the writing (which is a genuine concern with dask where processes are used).

So I am confident that a python API written here would be workable for us