Remove FilterScoringOperator

Author: carlosdelest

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterOperator.java

@@ -8,22 +8,31 @@
 package org.elasticsearch.compute.operator;
 
 import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.core.Releasables;
 
 import java.util.Arrays;
 
+import static org.elasticsearch.compute.lucene.LuceneQueryExpressionEvaluator.SCORE_FOR_FALSE;
+
 public class FilterOperator extends AbstractPageMappingOperator {
 
+    public static final int SCORE_BLOCK_INDEX = 1;
+
     private final EvalOperator.ExpressionEvaluator evaluator;
+    private final boolean usesScoring;
+    private final BlockFactory blockFactory;
 
-    public record FilterOperatorFactory(ExpressionEvaluator.Factory evaluatorSupplier) implements OperatorFactory {
+    public record FilterOperatorFactory(ExpressionEvaluator.Factory evaluatorSupplier, boolean usesScoring) implements OperatorFactory {
 
         @Override
         public Operator get(DriverContext driverContext) {
-            return new FilterOperator(evaluatorSupplier.get(driverContext));
+            return new FilterOperator(evaluatorSupplier.get(driverContext), usesScoring, driverContext.blockFactory());
         }
 
         @Override
@@ -32,30 +41,46 @@ public String describe() {
         }
     }
 
-    public FilterOperator(EvalOperator.ExpressionEvaluator evaluator) {
+    public FilterOperator(ExpressionEvaluator evaluator, boolean usesScoring, BlockFactory blockFactory) {
         this.evaluator = evaluator;
+        this.usesScoring = usesScoring;
+        this.blockFactory = blockFactory;
     }
 
     @Override
     protected Page process(Page page) {
         int rowCount = 0;
         int[] positions = new int[page.getPositionCount()];
 
-        try (BooleanBlock test = (BooleanBlock) evaluator.eval(page)) {
-            if (test.areAllValuesNull()) {
+        try (Block filterResultBlock = evaluator.eval(page)) {
+            if (filterResultBlock.areAllValuesNull()) {
                 // All results are null which is like false. No values selected.
                 page.releaseBlocks();
                 return null;
             }
+
+            // Explicit types to avoid casting on every element
+            DoubleBlock scoreBlock = null;
+            BooleanBlock testBlock = null;
+            if (usesScoring) {
+                assert filterResultBlock instanceof DoubleBlock : "Evaluated block should be a DoubleBlock when using scoring";
+                scoreBlock = (DoubleBlock) filterResultBlock;
+            } else {
+                assert filterResultBlock instanceof BooleanBlock : "Evaluated block should be a BooleanBlock when not using scoring";
+                testBlock = (BooleanBlock) filterResultBlock;
+            }
+
             // TODO we can detect constant true or false from the type
             // TODO or we could make a new method in bool-valued evaluators that returns a list of numbers
             for (int p = 0; p < page.getPositionCount(); p++) {
-                if (test.isNull(p) || test.getValueCount(p) != 1) {
+                if (filterResultBlock.isNull(p) || filterResultBlock.getValueCount(p) != 1) {
                     // Null is like false
                     // And, for now, multivalued results are like false too
                     continue;
                 }
-                if (test.getBoolean(test.getFirstValueIndex(p))) {
+                if (usesScoring && scoreBlock.getDouble(scoreBlock.getFirstValueIndex(p)) != SCORE_FOR_FALSE) {
+                    positions[rowCount++] = p;
+                } else if (usesScoring == false && testBlock.getBoolean(testBlock.getFirstValueIndex(p))) {
                     positions[rowCount++] = p;
                 }
             }
@@ -73,7 +98,11 @@ protected Page process(Page page) {
             boolean success = false;
             try {
                 for (int i = 0; i < page.getBlockCount(); i++) {
-                    filteredBlocks[i] = page.getBlock(i).filter(positions);
+                    if (usesScoring && i == SCORE_BLOCK_INDEX) {
+                        filteredBlocks[i] = createScoresBlock(rowCount, (DoubleBlock) filterResultBlock, positions);
+                    } else {
+                        filteredBlocks[i] = page.getBlock(i).filter(positions);
+                    }
                 }
                 success = true;
             } finally {
@@ -86,6 +115,15 @@ protected Page process(Page page) {
         }
     }
 
+    private Block createScoresBlock(int rowCount, DoubleBlock scoreBlock, int[] positions) {
+        // Create a new scores block with the retrieved scores, that will replace the existing one on the result page
+        DoubleVector.Builder updatedScoresBuilder = blockFactory.newDoubleVectorBuilder(rowCount);
+        for (int j = 0; j < rowCount; j++) {
+            updatedScoresBuilder.appendDouble(scoreBlock.getDouble(positions[j]));
+        }
+        return updatedScoresBuilder.build().asBlock();
+    }
+
     @Override
     public String toString() {
         return "FilterOperator[" + "evaluator=" + evaluator + ']';

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterScoringOperator.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.EvalOperator.EvalOperatorFactory;
 import org.elasticsearch.compute.operator.FilterOperator.FilterOperatorFactory;
-import org.elasticsearch.compute.operator.FilterScoringOperator.FilterScoringOperatorFactory;
 import org.elasticsearch.compute.operator.LocalSourceOperator;
 import org.elasticsearch.compute.operator.LocalSourceOperator.LocalSourceFactory;
 import org.elasticsearch.compute.operator.MvExpandOperator;
@@ -671,13 +670,7 @@ private PhysicalOperation planFilter(FilterExec filter, LocalExecutionPlannerCon
             usesScore
         );
 
-        OperatorFactory operatorFactory;
-        if (PlannerUtils.usesScoring(filter)) {
-            operatorFactory = new FilterScoringOperatorFactory(evaluatorFactory);
-        } else {
-            operatorFactory = new FilterOperatorFactory(evaluatorFactory);
-        }
-        return source.with(operatorFactory, source.layout);
+        return source.with(new FilterOperatorFactory(evaluatorFactory, PlannerUtils.usesScoring(filter)), source.layout);
     }
 
     private PhysicalOperation planLimit(LimitExec limit, LocalExecutionPlannerContext context) {