cpoerschke commented on code in PR #13539:
URL: https://github.com/apache/lucene/pull/13539#discussion_r1664572651
##########
lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsWriter.java:
##########
@@ -143,61 +145,81 @@ public int nextDoc() throws IOException {
public static final class MergedVectorValues {
private MergedVectorValues() {}
- /** Returns a merged view over all the segment's {@link
FloatVectorValues}. */
- public static FloatVectorValues mergeFloatVectorValues(
- FieldInfo fieldInfo, MergeState mergeState) throws IOException {
+ private static void validateFieldEncoding(FieldInfo fieldInfo,
VectorEncoding expected) {
assert fieldInfo != null && fieldInfo.hasVectorValues();
- if (fieldInfo.getVectorEncoding() != VectorEncoding.FLOAT32) {
+ VectorEncoding fieldEncoding = fieldInfo.getVectorEncoding();
+ if (fieldEncoding != expected) {
throw new UnsupportedOperationException(
- "Cannot merge vectors encoded as [" +
fieldInfo.getVectorEncoding() + "] as FLOAT32");
+ "Cannot merge vectors encoded as [" + fieldEncoding + "] as " +
expected);
}
- List<VectorValuesSub> subs = new ArrayList<>();
- for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
- KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i];
+ }
+
+ private static <V, S> List<S> mergeVectorValues(
+ KnnVectorsReader[] knnVectorsReaders,
+ MergeState.DocMap[] docMaps,
+ IOFunction<KnnVectorsReader, V> valuesSupplier,
+ BiFunction<MergeState.DocMap, V, S> newSub)
+ throws IOException {
+ List<S> subs = new ArrayList<>();
+ for (int i = 0; i < knnVectorsReaders.length; i++) {
+ KnnVectorsReader knnVectorsReader = knnVectorsReaders[i];
if (knnVectorsReader != null) {
- FloatVectorValues values =
knnVectorsReader.getFloatVectorValues(fieldInfo.name);
+ V values = valuesSupplier.apply(knnVectorsReader);
if (values != null) {
- subs.add(new VectorValuesSub(mergeState.docMaps[i], values));
+ subs.add(newSub.apply(docMaps[i], values));
}
}
}
- return new MergedFloat32VectorValues(subs, mergeState);
+ return subs;
+ }
+
+ /** Returns a merged view over all the segment's {@link
FloatVectorValues}. */
+ public static FloatVectorValues mergeFloatVectorValues(
+ FieldInfo fieldInfo, MergeState mergeState) throws IOException {
+ validateFieldEncoding(fieldInfo, VectorEncoding.FLOAT32);
+ return new MergedFloat32VectorValues(
+ mergeVectorValues(
+ mergeState.knnVectorsReaders,
+ mergeState.docMaps,
+ knnVectorsReader -> {
+ return knnVectorsReader.getFloatVectorValues(fieldInfo.name);
+ },
+ (docMap, values) -> {
+ return new FloatVectorValuesSub(docMap, values);
+ }),
+ mergeState);
}
/** Returns a merged view over all the segment's {@link ByteVectorValues}.
*/
public static ByteVectorValues mergeByteVectorValues(FieldInfo fieldInfo,
MergeState mergeState)
throws IOException {
- assert fieldInfo != null && fieldInfo.hasVectorValues();
- if (fieldInfo.getVectorEncoding() != VectorEncoding.BYTE) {
- throw new UnsupportedOperationException(
- "Cannot merge vectors encoded as [" +
fieldInfo.getVectorEncoding() + "] as BYTE");
- }
- List<ByteVectorValuesSub> subs = new ArrayList<>();
- for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
- KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i];
- if (knnVectorsReader != null) {
- ByteVectorValues values =
knnVectorsReader.getByteVectorValues(fieldInfo.name);
- if (values != null) {
- subs.add(new ByteVectorValuesSub(mergeState.docMaps[i], values));
- }
- }
- }
- return new MergedByteVectorValues(subs, mergeState);
+ validateFieldEncoding(fieldInfo, VectorEncoding.BYTE);
+ return new MergedByteVectorValues(
+ mergeVectorValues(
+ mergeState.knnVectorsReaders,
+ mergeState.docMaps,
+ knnVectorsReader -> {
+ return knnVectorsReader.getByteVectorValues(fieldInfo.name);
+ },
+ (docMap, values) -> {
+ return new ByteVectorValuesSub(docMap, values);
+ }),
+ mergeState);
}
static class MergedFloat32VectorValues extends FloatVectorValues {
Review Comment:
And maybe this could drop the `32` for consistency with
`MergedByteVectorValues` naming?
```suggestion
static class MergedFloatVectorValues extends FloatVectorValues {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org
For additional commands, e-mail: issues-h...@lucene.apache.org
