Remove incorrect/expensive use of ServiceLoader for choosing random format (#13428)

This commit is contained in:
Uwe Schindler 2024-05-27 23:59:28 +02:00 committed by GitHub
parent c06fff6a52
commit b1189adeb3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 12 deletions

View File

@ -18,6 +18,7 @@
package org.apache.lucene.codecs;
import java.io.IOException;
import java.util.Set;
import org.apache.lucene.index.ByteVectorValues;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.SegmentReadState;
@ -86,6 +87,11 @@ public abstract class KnnVectorsFormat implements NamedSPILoader.NamedSPI {
return Holder.getLoader().lookup(name);
}
/** returns a list of all available format names */
public static Set<String> availableKnnVectorsFormats() {
return Holder.getLoader().availableServices();
}
/** Returns a {@link KnnVectorsWriter} to write the vectors to the index. */
public abstract KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException;

View File

@ -89,7 +89,6 @@ import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Random;
import java.util.ServiceLoader;
import java.util.Set;
import java.util.TimeZone;
import java.util.TreeSet;
@ -3215,18 +3214,21 @@ public abstract class LuceneTestCase extends Assert {
return it;
}
protected KnnVectorsFormat randomVectorFormat(VectorEncoding vectorEncoding) {
ServiceLoader<KnnVectorsFormat> formats = java.util.ServiceLoader.load(KnnVectorsFormat.class);
List<KnnVectorsFormat> availableFormats = new ArrayList<>();
for (KnnVectorsFormat f : formats) {
if (f.getName().equals(HnswBitVectorsFormat.NAME)) {
if (vectorEncoding.equals(VectorEncoding.BYTE)) {
availableFormats.add(f);
}
} else {
availableFormats.add(f);
}
private static boolean supportsVectorEncoding(
KnnVectorsFormat format, VectorEncoding vectorEncoding) {
if (format instanceof HnswBitVectorsFormat) {
// special case, this only supports BYTE
return vectorEncoding == VectorEncoding.BYTE;
}
return true;
}
protected static KnnVectorsFormat randomVectorFormat(VectorEncoding vectorEncoding) {
List<KnnVectorsFormat> availableFormats =
KnnVectorsFormat.availableKnnVectorsFormats().stream()
.map(KnnVectorsFormat::forName)
.filter(format -> supportsVectorEncoding(format, vectorEncoding))
.toList();
return RandomPicks.randomFrom(random(), availableFormats);
}
}