Refactoring

Author: carlosdelest

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/CosineSimilarity.java

@@ -9,19 +9,16 @@
 
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesTo;
 import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesToLifecycle;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.Param;
-import org.elasticsearch.xpack.esql.expression.function.scalar.math.Pow;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 
 import java.io.IOException;
-import java.util.List;
 
 import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
 
@@ -48,29 +45,25 @@ public CosineSimilarity(
             description = "second dense_vector to calculate cosine similarity"
         ) Expression right
     ) {
-        super(source, List.of(left, right), left, right);
+        super(source, left, right);
     }
 
     private CosineSimilarity(StreamInput in) throws IOException {
         this(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(Expression.class), in.readNamedWriteable(Expression.class));
     }
 
     @Override
-    protected NodeInfo<? extends Expression> info() {
-        return NodeInfo.create(this, Pow::new, left(), right());
+    protected SimilarityEvaluatorFunction getSimilarityFunction() {
+        return COSINE::compare;
     }
 
     @Override
-    public String getWriteableName() {
-        return ENTRY.name;
+    protected NodeInfo<? extends Expression> info() {
+        return NodeInfo.create(this, CosineSimilarity::new, left(), right());
     }
 
     @Override
-    public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
-        return new SimilarityEvaluatorFactory(
-            toEvaluator.apply(left()),
-            toEvaluator.apply(right()),
-            (leftScratch, rightScratch) -> COSINE.compare(leftScratch, rightScratch)
-        );
+    public String getWriteableName() {
+        return ENTRY.name;
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/VectorSimilarityFunction.java

@@ -17,7 +17,6 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
-import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
 
 import java.io.IOException;
@@ -28,21 +27,16 @@
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
 
-public abstract class VectorSimilarityFunction extends EsqlScalarFunction implements VectorFunction {
-    protected Expression left;
-    protected Expression right;
-
-    public VectorSimilarityFunction(
-        Source source,
-        List<Expression> fields,
-        @Param(name = "left", type = { "dense_vector" }, description = "first dense_vector to calculate cosine similarity") Expression left,
-        @Param(
-            name = "right",
-            type = { "dense_vector" },
-            description = "second dense_vector to calculate cosine similarity"
-        ) Expression right
-    ) {
-        super(source, fields);
+/**
+ * Base class for vector similarity functions, which compute a similarity score between two dense vectors
+ */
+abstract class VectorSimilarityFunction extends EsqlScalarFunction implements VectorFunction {
+
+    private final Expression left;
+    private final Expression right;
+
+    protected VectorSimilarityFunction(Source source, Expression left, Expression right) {
+        super(source, List.of(left, right));
         this.left = left;
         this.right = right;
     }
@@ -85,29 +79,33 @@ public Expression right() {
         return right;
     }
 
+    /**
+     * Functional interface for evaluating the similarity between two float arrays
+     */
     @FunctionalInterface
     public interface SimilarityEvaluatorFunction {
         float calculateSimilarity(float[] leftScratch, float[] rightScratch);
     }
 
-    protected class SimilarityEvaluatorFactory implements EvalOperator.ExpressionEvaluator.Factory {
+    @Override
+    public final EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
+        return new SimilarityEvaluatorFactory(toEvaluator.apply(left()), toEvaluator.apply(right()), getSimilarityFunction());
+    }
 
-        private final EvalOperator.ExpressionEvaluator.Factory left;
-        private final EvalOperator.ExpressionEvaluator.Factory right;
-        private final SimilarityEvaluatorFunction similarityFunction;
+    /**
+     * Returns the similarity function to be used for evaluating the similarity between two vectors.
+     */
+    protected abstract SimilarityEvaluatorFunction getSimilarityFunction();
 
-        SimilarityEvaluatorFactory(
-            EvalOperator.ExpressionEvaluator.Factory left,
-            EvalOperator.ExpressionEvaluator.Factory right,
-            SimilarityEvaluatorFunction similarityFunction
-        ) {
-            this.left = left;
-            this.right = right;
-            this.similarityFunction = similarityFunction;
-        }
+    private record SimilarityEvaluatorFactory(
+        EvalOperator.ExpressionEvaluator.Factory left,
+        EvalOperator.ExpressionEvaluator.Factory right,
+        SimilarityEvaluatorFunction similarityFunction
+    ) implements EvalOperator.ExpressionEvaluator.Factory {
 
         @Override
         public EvalOperator.ExpressionEvaluator get(DriverContext context) {
+            // TODO check whether to use this custom evaluator or reuse / define an existing one
             return new EvalOperator.ExpressionEvaluator() {
                 @Override
                 public Block eval(Page page) {