Implement ExpressionScoreMapper for FullTextFunction and BinaryLogic

carlosdelest · carlosdelest · commit e4eb86d776c9 · 2025-02-05T17:29:53.000+01:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java
@@ -10,25 +10,24 @@
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator;
 import org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator.ShardConfig;
+import org.elasticsearch.compute.lucene.LuceneQueryScoreEvaluator;
 import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.compute.operator.ScoreOperator;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TranslationAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
-import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
 import org.elasticsearch.xpack.esql.core.expression.function.Function;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
-import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
 import org.elasticsearch.xpack.esql.expression.predicate.logical.BinaryLogic;
 import org.elasticsearch.xpack.esql.expression.predicate.logical.Not;
-import org.elasticsearch.xpack.esql.expression.predicate.logical.Or;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
@@ -39,6 +38,7 @@
 import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.querydsl.query.TranslationAwareExpressionQuery;
+import org.elasticsearch.xpack.esql.score.ExpressionScoreMapper;
 
 import java.util.List;
 import java.util.Locale;
@@ -56,7 +56,12 @@
  * These functions needs to be pushed down to Lucene queries to be executed - there's no Evaluator for them, but depend on
  * {@link org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizer} to rewrite them into Lucene queries.
  */
-public abstract class FullTextFunction extends Function implements TranslationAware, PostAnalysisPlanVerificationAware, EvaluatorMapper {
+public abstract class FullTextFunction extends Function
+    implements
+        TranslationAware,
+        PostAnalysisPlanVerificationAware,
+        EvaluatorMapper,
+        ExpressionScoreMapper {
 
     private final Expression query;
     private final QueryBuilder queryBuilder;
@@ -204,79 +209,13 @@ private static void checkFullTextQueryFunctions(LogicalPlan plan, Failures failu
                 failures
             );
             checkFullTextFunctionsParents(condition, failures);
-
-            boolean usesScore = plan.output()
-                .stream()
-                .anyMatch(attr -> attr instanceof MetadataAttribute ma && ma.name().equals(MetadataAttribute.SCORE));
-            if (usesScore) {
-                checkFullTextSearchDisjunctions(condition, failures);
-            }
         } else {
             plan.forEachExpression(FullTextFunction.class, ftf -> {
                 failures.add(fail(ftf, "[{}] {} is only supported in WHERE commands", ftf.functionName(), ftf.functionType()));
             });
         }
     }
 
-    /**
-     * Checks whether a condition contains a disjunction with a full text search.
-     * If it does, check that every element of the disjunction is a full text search or combinations (AND, OR, NOT) of them.
-     * If not, add a failure to the failures collection.
-     *
-     * @param condition        condition to check for disjunctions of full text searches
-     * @param failures         failures collection to add to
-     */
-    private static void checkFullTextSearchDisjunctions(Expression condition, Failures failures) {
-        Holder<Boolean> isInvalid = new Holder<>(false);
-        condition.forEachDown(Or.class, or -> {
-            if (isInvalid.get()) {
-                // Exit early if we already have a failures
-                return;
-            }
-            if (checkDisjunctionPushable(or) == false) {
-                isInvalid.set(true);
-                failures.add(
-                    fail(
-                        or,
-                        "Invalid condition when using METADATA _score [{}]. Full text functions can be used in an OR condition, "
-                            + "but only if just full text functions are used in the OR condition",
-                        or.sourceText()
-                    )
-                );
-            }
-        });
-    }
-
-    /**
-     * Checks if a disjunction is pushable from the point of view of FullTextFunctions. Either it has no FullTextFunctions or
-     * all it contains are FullTextFunctions.
-     *
-     * @param or disjunction to check
-     * @return true if the disjunction is pushable, false otherwise
-     */
-    private static boolean checkDisjunctionPushable(Or or) {
-        boolean hasFullText = or.anyMatch(FullTextFunction.class::isInstance);
-        return hasFullText == false || onlyFullTextFunctionsInExpression(or);
-    }
-
-    /**
-     * Checks whether an expression contains just full text functions or negations (NOT) and combinations (AND, OR) of full text functions
-     *
-     * @param expression expression to check
-     * @return true if all children are full text functions or negations of full text functions, false otherwise
-     */
-    private static boolean onlyFullTextFunctionsInExpression(Expression expression) {
-        if (expression instanceof FullTextFunction) {
-            return true;
-        } else if (expression instanceof Not) {
-            return onlyFullTextFunctionsInExpression(expression.children().get(0));
-        } else if (expression instanceof BinaryLogic binaryLogic) {
-            return onlyFullTextFunctionsInExpression(binaryLogic.left()) && onlyFullTextFunctionsInExpression(binaryLogic.right());
-        }
-
-        return false;
-    }
-
     /**
      * Checks all commands that exist before a specific type satisfy conditions.
      *
@@ -365,4 +304,15 @@ public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvalua
         }
         return new LuceneQueryExpressionEvaluator.Factory(shardConfigs);
     }
+
+    @Override
+    public ScoreOperator.ExpressionScorer.Factory toScorer(ToScorer toScorer) {
+        List<EsPhysicalOperationProviders.ShardContext> shardContexts = toScorer.shardContexts();
+        ShardConfig[] shardConfigs = new ShardConfig[shardContexts.size()];
+        int i = 0;
+        for (EsPhysicalOperationProviders.ShardContext shardContext : shardContexts) {
+            shardConfigs[i++] = new ShardConfig(shardContext.toQuery(queryBuilder()), shardContext.searcher());
+        }
+        return new LuceneQueryScoreEvaluator.Factory(shardConfigs);
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/logical/BinaryLogic.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/logical/BinaryLogic.java
@@ -8,6 +8,11 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.DoubleVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.ScoreOperator;
 import org.elasticsearch.xpack.esql.capabilities.TranslationAware;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
@@ -22,14 +27,18 @@
 import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
+import org.elasticsearch.xpack.esql.score.ExpressionScoreMapper;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isBoolean;
 
-public abstract class BinaryLogic extends BinaryOperator<Boolean, Boolean, Boolean, BinaryLogicOperation> implements TranslationAware {
+public abstract class BinaryLogic extends BinaryOperator<Boolean, Boolean, Boolean, BinaryLogicOperation>
+    implements
+        TranslationAware,
+        ExpressionScoreMapper {
 
     protected BinaryLogic(Source source, Expression left, Expression right, BinaryLogicOperation operation) {
         super(source, left, right, operation);
@@ -108,4 +117,30 @@ public static Query boolQuery(Source source, Query left, Query right, boolean is
         }
         return new BoolQuery(source, isAnd, queries);
     }
+
+    @Override
+    public ScoreOperator.ExpressionScorer.Factory toScorer(ToScorer toScorer) {
+        return context -> new BinaryLogicScorer(context, toScorer.toScorer(left()).get(context), toScorer.toScorer(right()).get(context));
+    }
+
+    private record BinaryLogicScorer(DriverContext driverContext, ScoreOperator.ExpressionScorer left, ScoreOperator.ExpressionScorer right)
+        implements
+            ScoreOperator.ExpressionScorer {
+        @Override
+        public DoubleBlock score(Page page) {
+            DoubleVector.Builder builder = driverContext.blockFactory().newDoubleVectorFixedBuilder(page.getPositionCount());
+            try (DoubleVector leftVector = left.score(page).asVector(); DoubleVector rightVector = right.score(page).asVector()) {
+                for (int i = 0; i < page.getPositionCount(); i++) {
+                    builder.appendDouble(leftVector.getDouble(i) + rightVector.getDouble(i));
+                }
+            }
+            return builder.build().asBlock();
+        }
+
+        @Override
+        public void close() {
+            left.close();
+            right.close();
+        }
+    }
 }
