Fp16 re ranking and prefetch support

lucene/core/src/java/org/apache/lucene/codecs/lucene95/OffHeapFloat16VectorValues.java

@@ -91,6 +91,24 @@ public VectorEncoding getEncoding() {
     return VectorEncoding.FLOAT16;
   }
 
+  @Override
+  public void prefetch(final int[] ordsToPrefetch, int numOrds) throws IOException {
+    if (ordsToPrefetch == null) {
+      return;
+    }
+
+    int finalNumOrds = Math.min(numOrds, ordsToPrefetch.length);
+    if (finalNumOrds <= 1) {
+      return;
+    }
+
+    // calculate offset and prefetch immediately
+    for (int i = 0; i < finalNumOrds; i++) {
+      long offset = (long) ordsToPrefetch[i] * byteSize;
+      slice.prefetch(offset, byteSize);
+    }
+  }
+
   public static OffHeapFloat16VectorValues load(
       VectorSimilarityFunction vectorSimilarityFunction,
       FlatVectorsScorer flatVectorsScorer,

lucene/core/src/java/org/apache/lucene/search/AbstractFullPrecisionVectorSimilarityValuesSource.java

@@ -0,0 +1,132 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+/**
+ * Base class for {@link DoubleValuesSource} implementations that compute vector similarity between
+ * a query vector and the raw full precision vectors of a KNN vector field. Subclasses supply the
+ * vector-type-specific access to the field values, leaving the common scoring flow here.
+ */
+abstract class AbstractFullPrecisionVectorSimilarityValuesSource extends DoubleValuesSource {
+
+  protected final String fieldName;
+  protected final VectorSimilarityFunction vectorSimilarityFunction;
+
+  protected AbstractFullPrecisionVectorSimilarityValuesSource(
+      String fieldName, VectorSimilarityFunction vectorSimilarityFunction) {
+    this.fieldName = fieldName;
+    this.vectorSimilarityFunction = vectorSimilarityFunction;
+  }
+
+  /** Returns the full precision similarity scores for documents in the given leaf. */
+  public DoubleValues getSimilarityScores(LeafReaderContext ctx) throws IOException {
+    return getValues(ctx, null);
+  }
+
+  @Override
+  public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
+    final KnnVectorValues vectorValues = getVectorValues(ctx);
+    if (vectorValues == null) {
+      checkField(ctx);
+      return DoubleValues.EMPTY;
+    }
+    final FieldInfo fi = ctx.reader().getFieldInfos().fieldInfo(fieldName);
+    if (fi.getVectorDimension() != queryDimension()) {
+      throw new IllegalArgumentException(
+          "Query vector dimension does not match field dimension: "
+              + queryDimension()
+              + " != "
+              + fi.getVectorDimension());
+    }
+
+    if (vectorSimilarityFunction == null) {
+      VectorScorer scorer = fullPrecisionRescorer(vectorValues);
+      if (scorer == null) {
+        return DoubleValues.EMPTY;
+      }
+      DocIdSetIterator iterator = scorer.iterator();
+      return new DoubleValues() {
+        @Override
+        public double doubleValue() throws IOException {
+          return scorer.score();
+        }
+
+        @Override
+        public boolean advanceExact(int doc) throws IOException {
+          return doc >= iterator.docID()
+              && (iterator.docID() == doc || iterator.advance(doc) == doc);
+        }
+      };
+    }
+    final KnnVectorValues.DocIndexIterator iterator = vectorValues.iterator();
+    return new DoubleValues() {
+      @Override
+      public double doubleValue() throws IOException {
+        return compareToQuery(vectorValues, iterator.index());
+      }
+
+      @Override
+      public boolean advanceExact(int doc) throws IOException {
+        return doc >= iterator.docID() && (iterator.docID() == doc || iterator.advance(doc) == doc);
+      }
+    };
+  }
+
+  /**
+   * Returns the full precision vector values for the field, or {@code null} if the field is absent.
+   */
+  protected abstract KnnVectorValues getVectorValues(LeafReaderContext ctx) throws IOException;
+
+  /** Raises the appropriate error when the field has no vector values of the expected type. */
+  protected abstract void checkField(LeafReaderContext ctx);
+
+  /** Returns the dimension of the query vector. */
+  protected abstract int queryDimension();
+
+  /**
+   * Returns a {@link VectorScorer} that scores the query vector against the full precision vectors,
+   * or {@code null} if no scorer is available.
+   */
+  protected abstract VectorScorer fullPrecisionRescorer(KnnVectorValues vectorValues)
+      throws IOException;
+
+  /** Compares the query vector against the field vector at the given ordinal. */
+  protected abstract double compareToQuery(KnnVectorValues vectorValues, int ord)
+      throws IOException;
+
+  @Override
+  public boolean needsScores() {
+    return false;
+  }
+
+  @Override
+  public DoubleValuesSource rewrite(IndexSearcher reader) throws IOException {
+    return this;
+  }
+
+  @Override
+  public boolean isCacheable(LeafReaderContext ctx) {
+    return true;
+  }
+}

lucene/core/src/java/org/apache/lucene/search/DoubleValuesSource.java

@@ -288,6 +288,21 @@ public static DoubleValues similarityToQueryVector(
     return new FloatVectorSimilarityValuesSource(queryVector, vectorField).getValues(ctx, null);
   }
 
+  /**
+   * Returns a DoubleValues instance for computing the vector similarity score per document against
+   * the float16 query vector
+   *
+   * @param ctx the context for which to return the DoubleValues
+   * @param queryVector float16 query vector
+   * @param vectorField knn float16 field name
+   * @return DoubleValues instance
+   * @throws IOException if an {@link IOException} occurs
+   */
+  public static DoubleValues similarityToQueryVector(
+      LeafReaderContext ctx, short[] queryVector, String vectorField) throws IOException {
+    return new Float16VectorSimilarityValuesSource(queryVector, vectorField).getValues(ctx, null);
+  }
+
   /**
    * Creates a DoubleValuesSource that wraps a generic NumericDocValues field. Assumes no monotonic
    * direction by default (Monotonicity.NONE).

lucene/core/src/java/org/apache/lucene/search/Float16VectorSimilarityValuesSource.java

@@ -0,0 +1,71 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.LeafReaderContext;
+
+/**
+ * A {@link DoubleValuesSource} which computes the vector similarity scores between the query vector
+ * and the {@link org.apache.lucene.document.KnnFloat16VectorField} for documents.
+ */
+class Float16VectorSimilarityValuesSource extends VectorSimilarityValuesSource {
+
+  private final short[] queryVector;
+
+  public Float16VectorSimilarityValuesSource(short[] vector, String fieldName) {
+    super(fieldName);
+    this.queryVector = vector;
+  }
+
+  @Override
+  public VectorScorer getScorer(LeafReaderContext ctx) throws IOException {
+    final Float16VectorValues vectorValues = ctx.reader().getFloat16VectorValues(fieldName);
+    if (vectorValues == null) {
+      Float16VectorValues.checkField(ctx.reader(), fieldName);
+      return null;
+    }
+    return vectorValues.scorer(queryVector);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(fieldName, Arrays.hashCode(queryVector));
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) return true;
+    if (obj == null || getClass() != obj.getClass()) return false;
+    Float16VectorSimilarityValuesSource other = (Float16VectorSimilarityValuesSource) obj;
+    return Objects.equals(fieldName, other.fieldName)
+        && Arrays.equals(queryVector, other.queryVector);
+  }
+
+  @Override
+  public String toString() {
+    return "Float16VectorSimilarityValuesSource(fieldName="
+        + fieldName
+        + " queryVector="
+        + Arrays.toString(queryVector)
+        + ")";
+  }
+}

lucene/core/src/java/org/apache/lucene/search/FullPrecisionFloat16VectorSimilarityValuesSource.java

@@ -0,0 +1,116 @@
+/*
+ * 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.search;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+import org.apache.lucene.index.Float16VectorValues;
+import org.apache.lucene.index.KnnVectorValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.VectorSimilarityFunction;
+
+/**
+ * A {@link DoubleValuesSource} that computes vector similarity between a query vector and the raw
+ * half-precision (FLOAT16) vectors indexed in the provided {@link
+ * org.apache.lucene.document.KnnFloat16VectorField} in documents.
+ */
+public class FullPrecisionFloat16VectorSimilarityValuesSource
+    extends AbstractFullPrecisionVectorSimilarityValuesSource {
+
+  private final short[] queryVector;
+
+  /**
+   * Creates a {@link DoubleValuesSource} that returns the vector similarity score between the
+   * provided query vector and the field for documents.
+   *
+   * @param vector the query vector
+   * @param fieldName the field name of the {@link org.apache.lucene.document.KnnFloat16VectorField}
+   * @param vectorSimilarityFunction the vector similarity function to use
+   */
+  public FullPrecisionFloat16VectorSimilarityValuesSource(
+      short[] vector, String fieldName, VectorSimilarityFunction vectorSimilarityFunction) {
+    super(fieldName, vectorSimilarityFunction);
+    this.queryVector = vector;
+  }
+
+  /**
+   * Creates a {@link DoubleValuesSource} that returns the vector similarity score between the
+   * provided query vector and the field for documents, using the similarity function configured for
+   * the field.
+   *
+   * @param vector the query vector
+   * @param fieldName the field name of the {@link org.apache.lucene.document.KnnFloat16VectorField}
+   */
+  public FullPrecisionFloat16VectorSimilarityValuesSource(short[] vector, String fieldName) {
+    this(vector, fieldName, null);
+  }
+
+  @Override
+  protected KnnVectorValues getVectorValues(LeafReaderContext ctx) throws IOException {
+    return ctx.reader().getFloat16VectorValues(fieldName);
+  }
+
+  @Override
+  protected void checkField(LeafReaderContext ctx) {
+    Float16VectorValues.checkField(ctx.reader(), fieldName);
+  }
+
+  @Override
+  protected int queryDimension() {
+    return queryVector.length;
+  }
+
+  @Override
+  protected VectorScorer fullPrecisionRescorer(KnnVectorValues vectorValues) throws IOException {
+    return ((Float16VectorValues) vectorValues).rescorer(queryVector);
+  }
+
+  @Override
+  protected double compareToQuery(KnnVectorValues vectorValues, int ord) throws IOException {
+    return vectorSimilarityFunction.compare(
+        queryVector, ((Float16VectorValues) vectorValues).vectorValue(ord));
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(fieldName, Arrays.hashCode(queryVector), vectorSimilarityFunction);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) return true;
+    if (obj == null || getClass() != obj.getClass()) return false;
+    FullPrecisionFloat16VectorSimilarityValuesSource other =
+        (FullPrecisionFloat16VectorSimilarityValuesSource) obj;
+    return Objects.equals(fieldName, other.fieldName)
+        && Objects.equals(vectorSimilarityFunction, other.vectorSimilarityFunction)
+        && Arrays.equals(queryVector, other.queryVector);
+  }
+
+  @Override
+  public String toString() {
+    return "FullPrecisionFloat16VectorSimilarityValuesSource(fieldName="
+        + fieldName
+        + " vectorSimilarityFunction="
+        + vectorSimilarityFunction
+        + " queryVector="
+        + Arrays.toString(queryVector)
+        + ")";
+  }
+}