GITHUB-12252: Add function queries for computing similarity scores between knn vectors

lucene/queries/src/java/org/apache/lucene/queries/function/FunctionValues.java

@@ -70,6 +70,14 @@ public boolean boolVal(int doc) throws IOException {
     return intVal(doc) != 0;
   }
 
+  public float[] floatVectorVal(int doc) throws IOException {
+    throw new UnsupportedOperationException();
+  }
+
+  public byte[] byteVectorVal(int doc) throws IOException {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * returns the bytes representation of the string val - TODO: should this return the indexed raw
    * bytes not?

lucene/queries/src/java/org/apache/lucene/queries/function/valuesource/ByteVectorFieldSource.java

@@ -0,0 +1,83 @@
+/*
+ * 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.queries.function.valuesource;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Map;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.queries.function.FunctionValues;
+import org.apache.lucene.queries.function.ValueSource;
+import org.apache.lucene.search.DocIdSetIterator;
+
+public class ByteVectorFieldSource extends ValueSource {
+  private final String fieldName;
+
+  public ByteVectorFieldSource(String fieldName) {
+    this.fieldName = fieldName;
+  }
+
+  @Override
+  public FunctionValues getValues(Map<Object, Object> context, LeafReaderContext readerContext)
+      throws IOException {
+
+    final ByteVectorValues vectorValues = readerContext.reader().getByteVectorValues(fieldName);
+    return new VectorFieldFunction(this) {
+      byte[] defaultVector = null;
+
+      @Override
+      public byte[] byteVectorVal(int doc) throws IOException {
+        if (exists(doc)) {
+          return vectorValues.vectorValue();
+        } else {
+          return defaultVector();
+        }
+      }
+
+      @Override
+      protected DocIdSetIterator getVectorIterator() {
+        return vectorValues;
+      }
+
+      private byte[] defaultVector() {
+        if (defaultVector == null) {
+          defaultVector = new byte[vectorValues.dimension()];
+          Arrays.fill(defaultVector, (byte) 0);
+        }
+        return defaultVector;
+      }
+    };
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (o.getClass() != ByteVectorFieldSource.class) return false;
+    ByteVectorFieldSource other = (ByteVectorFieldSource) o;
+    return fieldName.equals(other.fieldName);
+  }
+
+  @Override
+  public int hashCode() {
+    return getClass().hashCode() * 31 + fieldName.getClass().hashCode();
+  }
+
+  @Override
+  public String description() {
+    return "denseByteVectorField(" + fieldName + ")";
+  }
+}

lucene/queries/src/java/org/apache/lucene/queries/function/valuesource/ByteVectorSimilarityFunction.java

@@ -0,0 +1,35 @@
+/*
+ * 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.queries.function.valuesource;
+
+import java.io.IOException;
+
+import org.apache.lucene.queries.function.FunctionValues;
+import org.apache.lucene.queries.function.ValueSource;
+
+public class ByteVectorSimilarityFunction extends VectorSimilarityFunction {
+  public ByteVectorSimilarityFunction(
+          org.apache.lucene.index.VectorSimilarityFunction similarityFunction, ValueSource vector1, ValueSource vector2) {
+    super(similarityFunction, vector1, vector2);
+  }
+
+  @Override
+  protected float func(int doc, FunctionValues f1, FunctionValues f2) throws IOException {
+    assert f1.byteVectorVal(doc).length == f2.byteVectorVal(doc).length: "Vectors must have the same length";
+    return similarityFunction.compare(f1.byteVectorVal(doc), f2.byteVectorVal(doc));
+  }
+}

lucene/queries/src/java/org/apache/lucene/queries/function/valuesource/ByteVectorValueSource.java

@@ -0,0 +1,74 @@
+/*
+ * 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.queries.function.valuesource;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.queries.function.FunctionValues;
+import org.apache.lucene.queries.function.ValueSource;
+/** Function that returns a constant byte vector value for every document. */
+public class ByteVectorValueSource extends ValueSource {
+  byte[] vector;
+
+  public ByteVectorValueSource(List<Number> constVector) {
+    this.vector = new byte[constVector.size()];
+    for (int i = 0; i < constVector.size(); i++) {
+      vector[i] = constVector.get(i).byteValue();
+    }
+  }
+
+  @Override
+  public FunctionValues getValues(Map<Object, Object> context, LeafReaderContext readerContext)
+      throws IOException {
+    return new FunctionValues() {
+      @Override
+      public byte[] byteVectorVal(int doc) {
+        return vector;
+      }
+
+      @Override
+      public String strVal(int doc) {
+        return Arrays.toString(vector);
+      }
+
+      @Override
+      public String toString(int doc) throws IOException {
+        return description() + '=' + strVal(doc);
+      }
+    };
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (!(o instanceof ByteVectorValueSource)) return false;
+    ByteVectorValueSource other = (ByteVectorValueSource) o;
+    return Arrays.equals(vector, other.vector);
+  }
+
+  @Override
+  public int hashCode() {
+    return getClass().hashCode() * 31 + Arrays.hashCode(vector);
+  }
+
+  @Override
+  public String description() {
+    return "denseVectorConst(" + Arrays.toString(vector) + ')';
+  }
+}

lucene/queries/src/java/org/apache/lucene/queries/function/valuesource/FloatVectorFieldSource.java

@@ -0,0 +1,83 @@
+/*
+ * 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.queries.function.valuesource;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Map;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.queries.function.FunctionValues;
+import org.apache.lucene.queries.function.ValueSource;
+import org.apache.lucene.search.DocIdSetIterator;
+
+public class FloatVectorFieldSource extends ValueSource {
+  private final String fieldName;
+
+  public FloatVectorFieldSource(String fieldName) {
+    this.fieldName = fieldName;
+  }
+
+  @Override
+  public FunctionValues getValues(Map<Object, Object> context, LeafReaderContext readerContext)
+      throws IOException {
+
+    final FloatVectorValues vectorValues = readerContext.reader().getFloatVectorValues(fieldName);
+    return new VectorFieldFunction(this) {
+      float[] defaultVector = null;
+
+      @Override
+      public float[] floatVectorVal(int doc) throws IOException {
+        if (exists(doc)) {
+          return vectorValues.vectorValue();
+        } else {
+          return defaultVector();
+        }
+      }
+
+      @Override
+      protected DocIdSetIterator getVectorIterator() {
+        return vectorValues;
+      }
+
+      private float[] defaultVector() {
+        if (defaultVector == null) {
+          defaultVector = new float[vectorValues.dimension()];
+          Arrays.fill(defaultVector, 0.f);
+        }
+        return defaultVector;
+      }
+    };
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (o.getClass() != FloatVectorFieldSource.class) return false;
+    FloatVectorFieldSource other = (FloatVectorFieldSource) o;
+    return fieldName.equals(other.fieldName);
+  }
+
+  @Override
+  public int hashCode() {
+    return getClass().hashCode() * 31 + fieldName.getClass().hashCode();
+  }
+
+  @Override
+  public String description() {
+    return "denseFloatVectorField(" + fieldName + ")";
+  }
+}

lucene/queries/src/java/org/apache/lucene/queries/function/valuesource/FloatVectorSimilarityFunction.java

@@ -0,0 +1,35 @@
+/*
+ * 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.queries.function.valuesource;
+
+import java.io.IOException;
+
+import org.apache.lucene.queries.function.FunctionValues;
+import org.apache.lucene.queries.function.ValueSource;
+
+public class FloatVectorSimilarityFunction extends VectorSimilarityFunction {
+  public FloatVectorSimilarityFunction(
+          org.apache.lucene.index.VectorSimilarityFunction similarityFunction, ValueSource vector1, ValueSource vector2) {
+    super(similarityFunction, vector1, vector2);
+  }
+
+  @Override
+  protected float func(int doc, FunctionValues f1, FunctionValues f2) throws IOException {
+    assert f1.floatVectorVal(doc).length == f2.floatVectorVal(doc).length: "Vectors must have the same length";
+    return similarityFunction.compare(f1.floatVectorVal(doc), f2.floatVectorVal(doc));
+  }
+}