ProductQuantization.compute fails with an opaque wrapped error when the corpus is smaller than the cluster count

[tae898]

Hello from the ArcadeDB side, which vendors JVector (currently 4.0.0-rc.9) for its vector indexes.

What happens

Calling ProductQuantization.compute(...) with fewer points than clusters fails with an IllegalArgumentException that is raised inside a parallel task, so the caller sees it wrapped and pointed at an internal class rather than at the parameter that is wrong:

java.lang.IllegalArgumentException: java.lang.IllegalArgumentException: Number of clusters 256 cannot exceed number of points 3
    at java.base/java.util.concurrent.ForkJoinTask.getException(Unknown Source)
    at java.base/java.util.concurrent.ForkJoinTask.reportException(Unknown Source)
    at java.base/java.util.concurrent.ForkJoinTask.join(Unknown Source)
    at io.github.jbellis.jvector.quantization.ProductQuantization.createCodebooks(ProductQuantization.java:491)
    at io.github.jbellis.jvector.quantization.ProductQuantization.compute(ProductQuantization.java:134)
    at io.github.jbellis.jvector.quantization.ProductQuantization.compute(ProductQuantization.java:86)
Caused by: java.lang.IllegalArgumentException: Number of clusters 256 cannot exceed number of points 3
    at io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.chooseInitialCentroids(KMeansPlusPlusClusterer.java:176)

The message itself is accurate and helpful. The problem is where it comes from: the double wrapping plus a stack that bottoms out in chooseInitialCentroids gives no indication that the fix is either "index more vectors first" or "lower clusters", and the IllegalArgumentException about clustering is easy to mistake for a data problem rather than a configuration one.

Why it is easy to hit

Any integration that exposes product quantization as an index option will hit it on small datasets: tests, examples, demos, and the first seconds of a real ingest all sit below 256 vectors. In our case a background rebuild fired on a 3-vector index, the build died, and the index was left with no graph and silently returned zero results afterwards, which took a while to trace back to this line.

Reproduction

Any corpus smaller than clusters:

ProductQuantization.compute(ravv /* 3 vectors */, subspaces, 256 /* clusters */, false);

Suggestion

Validate in ProductQuantization.compute before the parallel section, so the error is raised on the caller's thread and names the knob, for example:

cannot train product quantization with 256 clusters on 3 vectors: reduce clusters to at most 3 or index more vectors first

Optionally, clamping clusters to min(clusters, pointCount) with a warning would let small indexes work instead of failing, though I appreciate that silently changing a quantization parameter may be the wrong default for a library. Either behaviour is easier to act on than the current wrapped throw.

Happy to send a PR for the up-front validation if that is a shape you would accept.

[ashkrisk]

Thanks for pointing this out. I've raised a PR which improves the error message and moves it outside the parallel block, see #706.

I think it's better for JVector to continue to throw an error if the number of training vectors is less than the number of clusters, which forces the application to be more specific about what it wants. AFAIK Cassandra and Opensearch deal with this problem by delaying the use of PQ until the number of available vectors crosses a specific threshold.

Clamping is arguably a simpler solution, and it's easy enough to implement client-side.

ProductQuantization.compute(ravv, subspaces, Math.min(256, ravv.size()), false);

[r-devulap]

Clamping is arguably a simpler solution, and it's easy enough to implement client-side.

Agreed.