Skip to content

Commit

Permalink
Remove incorrect/expensive use of ServiceLoader for choosing random f…
Browse files Browse the repository at this point in the history
…ormat (#13428)

# Conflicts:
#	lucene/test-framework/src/java/org/apache/lucene/tests/util/LuceneTestCase.java
  • Loading branch information
uschindler committed May 27, 2024
1 parent 4be6531 commit 273b146
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -86,6 +87,11 @@ public static KnnVectorsFormat forName(String 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. */
public abstract KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3281,4 +3281,22 @@ public static BytesRef newBytesRef(byte[] bytesIn, int offset, int length) {

return it;
}

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);
}
}

0 comments on commit 273b146

Please sign in to comment.