supabase / vecs

Postgres/pgvector Python Client
https://supabase.github.io/vecs/latest
Apache License 2.0
219 stars 33 forks source link

User-definable ivfflat.probes #19

Closed olirice closed 1 year ago

olirice commented 1 year ago

When an index is present on a collection, the ivfflat.probes parameter defines controls the balance of speed of results and precision of results.

Currently, ivfflat.probes is fixed at 10 so the user does not have control over the precision of results

https://github.com/supabase/vecs/blob/ffeab7fd19df505d28ce7935df8a66ad5a6928e8/src/vecs/collection.py#L356

There should be a key-word only argument in the query method to allow that to be user adjustable

https://github.com/supabase/vecs/blob/ffeab7fd19df505d28ce7935df8a66ad5a6928e8/src/vecs/collection.py#L291-L299

note, I'm not sure if settings can use bind params like

text("set ivfflat.probes = :probes").bindparams(probes=probes)

so it might need to use an f-string where the user's parameter has been confirmed to be an int

if not isinstance(probes, int):
  raise ...
text(f"set ivfflat.probes = {probes}"
jbritain commented 1 year ago

I have added the parameter (not PRed yet), and it passes tests, however these tests weren't actually built to include the parameter, and the default value is still 10. Should the tests be modified somehow to check if different values work, or is this a case where if it doesn't work it'd be user error?

olirice commented 1 year ago

I have added the parameter (not PRed yet), and it passes tests

:raised_hands:

these tests weren't actually built to include the parameter....Should the tests be modified somehow to check if different values work

yes, please add tests for different levels of probes. Its a little complex to test the actual expected behavior that results from increasing probes but we should at least check that manually passing in a different values does not result in an error

the default value is still 10

from the pgvector docs https://github.com/pgvector/pgvector#indexing it recommends sqrt(lists) but based on issue that feels a little low. Lets keep it at 10 for now and I'll open a separate issue to figure out the best default.

a case where if it doesn't work it's be user error

it would be good to guard against (raise an exception for) probes < 1, but there's no need to enforce a maximum

jbritain commented 1 year ago

What parameters would be affected by the number of probes, and how best could I test for this? The number affects precision, but as far as I can tell we can't predict this precision. Should I just test that the correct number of results is returned like with the baseline query test, or is there something else which is both affected by the number of probes and easy to assert?

olirice commented 1 year ago

parameters would be affected by the number of probes

the only externally visible changes are

but, like you point out, those are harder to test, so for now its fine to assert that you get the correct number of results back when changing the number of probes

assert len(docs.query(..., limit =10, probes=10)) == 10
assert len(docs.query(..., limit =10, probes=5)) == 10
assert len(docs.query(..., limit =10, probes=1)) == 10
assert len(docs.query(..., limit =10, probes=999)) == 10

# error case: probes must be greater than 0
with pytest.raises(vecs.exc.ArgError):
    docs.query(..., probes=0)

# error case: probes must be greater than 0
with pytest.raises(vecs.exc.ArgError):
    docs.query(..., probes=-1)

# error case: probes must be an integer
with pytest.raises(vecs.exc.ArgError):
    docs.query(..., probes=7.5)