[WIP] Multi-Vector support for HNSW search

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene60/Lucene60FieldInfosFormat.java

@@ -29,6 +29,7 @@
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -218,6 +219,7 @@ private FieldInfo[] readFieldInfos(IndexInput input, int version) throws IOExcep
                 0,
                 VectorEncoding.FLOAT32,
                 VectorSimilarityFunction.EUCLIDEAN,
+                MultiVectorSimilarityFunction.Aggregation.NONE,
                 isSoftDeletesField,
                 false);
       } catch (IllegalStateException e) {

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90FieldInfosFormat.java

@@ -28,6 +28,7 @@
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -195,6 +196,7 @@ public FieldInfos read(
                     vectorDimension,
                     VectorEncoding.FLOAT32,
                     vectorDistFunc,
+                    MultiVectorSimilarityFunction.Aggregation.NONE,
                     isSoftDeletesField,
                     false);
             infos[i].checkConsistency();

lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldInfosFormat.java

@@ -27,6 +27,7 @@
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -199,6 +200,7 @@ public FieldInfos read(
                 vectorNumDimensions,
                 vectorEncoding,
                 vectorDistFunc,
+                MultiVectorSimilarityFunction.Aggregation.NONE,
                 isSoftDeletesField,
                 isParentField);
       }

lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/TestBlockWriter.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.ByteBuffersDataOutput;
@@ -120,6 +121,7 @@ private static FieldInfo getMockFieldInfo(String fieldName, int number) {
         0,
         VectorEncoding.FLOAT32,
         VectorSimilarityFunction.EUCLIDEAN,
+        MultiVectorSimilarityFunction.Aggregation.NONE,
         true,
         false);
   }

lucene/codecs/src/test/org/apache/lucene/codecs/uniformsplit/sharedterms/TestSTBlockReader.java

@@ -39,6 +39,7 @@
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.ImpactsEnum;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorEncoding;
@@ -207,6 +208,7 @@ private static FieldInfo mockFieldInfo(String fieldName, int number) {
         0,
         VectorEncoding.FLOAT32,
         VectorSimilarityFunction.EUCLIDEAN,
+        MultiVectorSimilarityFunction.Aggregation.NONE,
         false,
         false);
   }

lucene/core/src/java/module-info.java

@@ -77,7 +77,8 @@
   provides org.apache.lucene.codecs.KnnVectorsFormat with
       org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat,
       org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat,
-      org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat;
+      org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsFormat,
+      org.apache.lucene.codecs.lucene99.Lucene99HnswMultiVectorsFormat;
   provides org.apache.lucene.codecs.PostingsFormat with
       org.apache.lucene.codecs.lucene99.Lucene99PostingsFormat;
   provides org.apache.lucene.index.SortFieldProvider with

lucene/core/src/java/org/apache/lucene/codecs/hnsw/DefaultFlatMultiVectorScorer.java

@@ -0,0 +1,212 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.codecs.hnsw;
+
+import java.io.IOException;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.hnsw.RandomAccessVectorValues;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+
+/**
+ * Default multi-vector implementation of {@link FlatVectorsScorer}.
+ *
+ * @lucene.experimental
+ */
+// noCommit - pending tests
+public class DefaultFlatMultiVectorScorer extends DefaultFlatVectorScorer {
+
+  public static final DefaultFlatMultiVectorScorer INSTANCE = new DefaultFlatMultiVectorScorer();
+
+//  @Override
+//  public RandomVectorScorerSupplier getRandomVectorScorerSupplier(VectorSimilarityFunction similarityFunction, RandomAccessVectorValues vectorValues) throws IOException {
+//    throw new IllegalStateException("getRandomVectorScorerSupplier called on a Multi-Vector scorer. Expected call to getRandomMultiVectorScorerSupplier instead");
+//  }
+
+  @Override
+  public RandomVectorScorerSupplier getRandomMultiVectorScorerSupplier(
+      MultiVectorSimilarityFunction similarityFunction, RandomAccessVectorValues values)
+      throws IOException {
+    if (values instanceof RandomAccessVectorValues.Floats floatMultiVectorValues) {
+      return new FloatScoringSupplier(floatMultiVectorValues, similarityFunction);
+    } else if (values instanceof RandomAccessVectorValues.Bytes byteMultiVectorValues) {
+      return new ByteScoringSupplier(byteMultiVectorValues, similarityFunction);
+    }
+    throw new IllegalArgumentException(
+        "MultiVector values must be an instance of RandomAccessVectorValues.Floats or RandomAccessVectorValues.Bytes");
+  }
+
+  @Override
+  public RandomVectorScorer getRandomMultiVectorScorer(
+      MultiVectorSimilarityFunction similarityFunction,
+      RandomAccessVectorValues values,
+      float[] target)
+      throws IOException {
+    assert values instanceof RandomAccessVectorValues.Floats;
+    if (target.length % values.dimension() != 0) {
+      throw new IllegalArgumentException(
+          "query multiVector dimension differs from multiVector field dimension: "
+              + values.dimension());
+    }
+    return new FloatMultiVectorScorer(
+        (RandomAccessVectorValues.Floats) values, target, similarityFunction);
+  }
+
+  @Override
+  public RandomVectorScorer getRandomMultiVectorScorer(
+      MultiVectorSimilarityFunction similarityFunction,
+      RandomAccessVectorValues values,
+      byte[] target)
+      throws IOException {
+    assert values instanceof RandomAccessVectorValues.Bytes;
+    if (target.length % values.dimension() != 0) {
+      throw new IllegalArgumentException(
+          "query multiVector dimension differs from multiVector field dimension: "
+              + values.dimension());
+    }
+    return new ByteMultiVectorScorer(
+        (RandomAccessVectorValues.Bytes) values, target, similarityFunction);
+  }
+
+  @Override
+  public String toString() {
+    return "DefaultFlatMultiVectorScorer()";
+  }
+
+  /** RandomVectorScorerSupplier for bytes multiVectors */
+  private static final class ByteScoringSupplier implements RandomVectorScorerSupplier {
+    private final RandomAccessVectorValues.Bytes multiVectors;
+    private final RandomAccessVectorValues.Bytes multiVectors1;
+    private final RandomAccessVectorValues.Bytes multiVectors2;
+    private final MultiVectorSimilarityFunction similarityFunction;
+    private final int dimension;
+
+    private ByteScoringSupplier(
+        RandomAccessVectorValues.Bytes multiVectors,
+        MultiVectorSimilarityFunction similarityFunction)
+        throws IOException {
+      this.multiVectors = multiVectors;
+      multiVectors1 = multiVectors.copy();
+      multiVectors2 = multiVectors.copy();
+      this.similarityFunction = similarityFunction;
+      this.dimension = multiVectors.dimension();
+    }
+
+    @Override
+    public RandomVectorScorer scorer(int ord) {
+      return new RandomVectorScorer.AbstractRandomVectorScorer(multiVectors) {
+        @Override
+        public float score(int node) throws IOException {
+          return similarityFunction.compare(
+              multiVectors1.vectorValue(ord), multiVectors2.vectorValue(node), dimension);
+        }
+      };
+    }
+
+    @Override
+    public RandomVectorScorerSupplier copy() throws IOException {
+      return new ByteScoringSupplier(multiVectors, similarityFunction);
+    }
+  }
+
+  /** RandomVectorScorerSupplier for Float vector */
+  private static final class FloatScoringSupplier implements RandomVectorScorerSupplier {
+    private final RandomAccessVectorValues.Floats multiVectors;
+    private final RandomAccessVectorValues.Floats multiVectors1;
+    private final RandomAccessVectorValues.Floats multiVectors2;
+    private final MultiVectorSimilarityFunction similarityFunction;
+    private final int dimension;
+
+    private FloatScoringSupplier(
+        RandomAccessVectorValues.Floats multiVectors,
+        MultiVectorSimilarityFunction similarityFunction)
+        throws IOException {
+      this.multiVectors = multiVectors;
+      multiVectors1 = multiVectors.copy();
+      multiVectors2 = multiVectors.copy();
+      this.similarityFunction = similarityFunction;
+      this.dimension = multiVectors.dimension();
+    }
+
+    @Override
+    public RandomVectorScorer scorer(int ord) {
+      return new RandomVectorScorer.AbstractRandomVectorScorer(multiVectors) {
+        @Override
+        public float score(int node) throws IOException {
+          return similarityFunction.compare(
+              multiVectors1.vectorValue(ord), multiVectors2.vectorValue(node), dimension);
+        }
+      };
+    }
+
+    @Override
+    public RandomVectorScorerSupplier copy() throws IOException {
+      return new FloatScoringSupplier(multiVectors, similarityFunction);
+    }
+  }
+
+  /** A {@link RandomVectorScorer} for float multiVectors. */
+  private static class FloatMultiVectorScorer
+      extends RandomVectorScorer.AbstractRandomVectorScorer {
+    private final RandomAccessVectorValues.Floats values;
+    private final float[] query;
+    private final MultiVectorSimilarityFunction similarityFunction;
+    private final int dimension;
+
+    public FloatMultiVectorScorer(
+        RandomAccessVectorValues.Floats values,
+        float[] query,
+        MultiVectorSimilarityFunction similarityFunction) {
+      super(values);
+      this.values = values;
+      this.query = query;
+      this.similarityFunction = similarityFunction;
+      this.dimension = values.dimension();
+    }
+
+    @Override
+    public float score(int node) throws IOException {
+      return similarityFunction.compare(query, values.vectorValue(node), dimension);
+    }
+  }
+
+  /** A {@link RandomVectorScorer} for byte multiVectors. */
+  private static class ByteMultiVectorScorer extends RandomVectorScorer.AbstractRandomVectorScorer {
+    private final RandomAccessVectorValues.Bytes values;
+    private final byte[] query;
+    private final MultiVectorSimilarityFunction similarityFunction;
+    private final int dimension;
+
+    public ByteMultiVectorScorer(
+        RandomAccessVectorValues.Bytes values,
+        byte[] query,
+        MultiVectorSimilarityFunction similarityFunction) {
+      super(values);
+      this.values = values;
+      this.query = query;
+      this.similarityFunction = similarityFunction;
+      this.dimension = values.dimension();
+    }
+
+    @Override
+    public float score(int node) throws IOException {
+      return similarityFunction.compare(query, values.vectorValue(node), dimension);
+    }
+  }
+}

lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsScorer.java

@@ -18,6 +18,7 @@
 package org.apache.lucene.codecs.hnsw;
 
 import java.io.IOException;
+import org.apache.lucene.index.MultiVectorSimilarityFunction;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.util.hnsw.RandomAccessVectorValues;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
@@ -72,4 +73,54 @@ RandomVectorScorer getRandomVectorScorer(
       RandomAccessVectorValues vectorValues,
       byte[] target)
       throws IOException;
+
+  /**
+   * Returns a {@link RandomVectorScorerSupplier} that can be used to score multi-vector values
+   *
+   * @param similarityFunction the similarity function to use
+   * @param values the multi-vector values to score
+   * @return a {@link RandomVectorScorerSupplier} that can be used to score multi-vectors
+   * @throws IOException if an I/O error occurs
+   */
+  default RandomVectorScorerSupplier getRandomMultiVectorScorerSupplier(
+      MultiVectorSimilarityFunction similarityFunction, RandomAccessVectorValues values)
+      throws IOException {
+    throw new UnsupportedOperationException("Multi-Vector scoring not supported");
+  }
+
+  /**
+   * Returns a {@link RandomVectorScorer} for the given set of multi-vectors and target
+   * multi-vector.
+   *
+   * @param similarityFunction the similarity function to use
+   * @param values the multi-vector values to score
+   * @param target the target multi-vector with vector values packed in a single array
+   * @return a {@link RandomVectorScorer} for the given field and target multi-vector.
+   * @throws IOException if an I/O error occurs when reading from the index.
+   */
+  default RandomVectorScorer getRandomMultiVectorScorer(
+      MultiVectorSimilarityFunction similarityFunction,
+      RandomAccessVectorValues values,
+      float[] target)
+      throws IOException {
+    throw new UnsupportedOperationException("Multi-Vector scoring not supported");
+  }
+
+  /**
+   * Returns a {@link RandomVectorScorer} for the given set of multi-vectors and target
+   * multi-vector.
+   *
+   * @param similarityFunction the similarity function to use
+   * @param values the multi-vector values to score
+   * @param target the target multi-vector with vector values packed in a single array
+   * @return a {@link RandomVectorScorer} for the given field and target multi-vector.
+   * @throws IOException if an I/O error occurs when reading from the index.
+   */
+  default RandomVectorScorer getRandomMultiVectorScorer(
+      MultiVectorSimilarityFunction similarityFunction,
+      RandomAccessVectorValues values,
+      byte[] target)
+      throws IOException {
+    throw new UnsupportedOperationException("Multi-Vector scoring not supported");
+  }
 }