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; package org.apache.lucene.codecs;
import java.io.IOException; import java.io.IOException;
import java.util.Set;
import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.ByteVectorValues;
import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentReadState;
@ -86,6 +87,11 @@ public abstract class KnnVectorsFormat implements NamedSPILoader.NamedSPI {
return Holder.getLoader().lookup(name); 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. */ /** Returns a {@link KnnVectorsWriter} to write the vectors to the index. */
public abstract KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException; 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.Locale;
import java.util.Map; import java.util.Map;
import java.util.Random; import java.util.Random;
import java.util.ServiceLoader;
import java.util.Set; import java.util.Set;
import java.util.TimeZone; import java.util.TimeZone;
import java.util.TreeSet; import java.util.TreeSet;
@ -3215,18 +3214,21 @@ public abstract class LuceneTestCase extends Assert {
return it; return it;
} }
protected KnnVectorsFormat randomVectorFormat(VectorEncoding vectorEncoding) { private static boolean supportsVectorEncoding(
ServiceLoader<KnnVectorsFormat> formats = java.util.ServiceLoader.load(KnnVectorsFormat.class); KnnVectorsFormat format, VectorEncoding vectorEncoding) {
List<KnnVectorsFormat> availableFormats = new ArrayList<>(); if (format instanceof HnswBitVectorsFormat) {
for (KnnVectorsFormat f : formats) { // special case, this only supports BYTE
if (f.getName().equals(HnswBitVectorsFormat.NAME)) { return vectorEncoding == VectorEncoding.BYTE;
if (vectorEncoding.equals(VectorEncoding.BYTE)) {
availableFormats.add(f);
}
} else {
availableFormats.add(f);
} }
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); return RandomPicks.randomFrom(random(), availableFormats);
} }
} }