Implement usesScoring and solve scoring in binary logic

Author: carlosdelest

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneQueryExpressionEvaluator.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.compute.data.BooleanVector;
 import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DocVector;
+import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
@@ -83,7 +84,7 @@ public Block eval(Page page) {
     }
 
     @Override
-    public Block score(Page page, BlockFactory blockFactory) {
+    public DoubleBlock score(Page page, BlockFactory blockFactory) {
         if (scoreVector == null) {
             Releasables.closeExpectNoException(eval(page));
         }
@@ -152,8 +153,10 @@ private Tuple<BooleanVector, DoubleVector> evalSlow(DocVector docs) throws IOExc
         int prevShard = -1;
         int prevSegment = -1;
         SegmentState segmentState = null;
-        try (BooleanVector.Builder builder = blockFactory.newBooleanVectorFixedBuilder(docs.getPositionCount());
-             DoubleVector.Builder scoreBuilder = blockFactory.newDoubleVectorFixedBuilder(docs.getPositionCount())) {
+        try (
+            BooleanVector.Builder builder = blockFactory.newBooleanVectorFixedBuilder(docs.getPositionCount());
+            DoubleVector.Builder scoreBuilder = blockFactory.newDoubleVectorFixedBuilder(docs.getPositionCount())
+        ) {
             for (int i = 0; i < docs.getPositionCount(); i++) {
                 int shard = docs.shards().getInt(docs.shards().getInt(map[i]));
                 int segment = docs.segments().getInt(map[i]);
@@ -170,7 +173,7 @@ private Tuple<BooleanVector, DoubleVector> evalSlow(DocVector docs) throws IOExc
                     segmentState.scoreSingleDocWithScorer(builder, scoreBuilder, docs.docs().getInt(map[i]));
                 }
             }
-            try (BooleanVector outOfOrder = builder.build();DoubleVector outOfOrderScores = scoreBuilder.build()) {
+            try (BooleanVector outOfOrder = builder.build(); DoubleVector outOfOrderScores = scoreBuilder.build()) {
                 return Tuple.tuple(
                     outOfOrder.filter(docs.shardSegmentDocMapBackwards()),
                     outOfOrderScores.filter(docs.shardSegmentDocMapBackwards())
@@ -261,8 +264,10 @@ private SegmentState(Weight weight, LeafReaderContext ctx) {
         Tuple<BooleanVector, DoubleVector> scoreDense(int min, int max) throws IOException {
             int length = max - min + 1;
             if (noMatch) {
-                return Tuple.tuple(blockFactory.newConstantBooleanVector(false, length),
-                    blockFactory.newConstantDoubleVector(NO_MATCH_SCORE, length));
+                return Tuple.tuple(
+                    blockFactory.newConstantBooleanVector(false, length),
+                    blockFactory.newConstantDoubleVector(NO_MATCH_SCORE, length)
+                );
             }
             if (bulkScorer == null ||  // The bulkScorer wasn't initialized
                 Thread.currentThread() != bulkScorerThread // The bulkScorer was initialized on a different thread
@@ -271,8 +276,10 @@ Tuple<BooleanVector, DoubleVector> scoreDense(int min, int max) throws IOExcepti
                 bulkScorer = weight.bulkScorer(ctx);
                 if (bulkScorer == null) {
                     noMatch = true;
-                    return Tuple.tuple(blockFactory.newConstantBooleanVector(false, length),
-                        blockFactory.newConstantDoubleVector(NO_MATCH_SCORE, length));
+                    return Tuple.tuple(
+                        blockFactory.newConstantBooleanVector(false, length),
+                        blockFactory.newConstantDoubleVector(NO_MATCH_SCORE, length)
+                    );
                 }
             }
 
@@ -300,8 +307,10 @@ Tuple<BooleanVector, DoubleVector> scoreSparse(IntVector docs) throws IOExceptio
                     blockFactory.newConstantDoubleVector(NO_MATCH_SCORE, docs.getPositionCount())
                 );
             }
-            try (BooleanVector.Builder builder = blockFactory.newBooleanVectorFixedBuilder(docs.getPositionCount());
-                DoubleVector.Builder scoreBuilder = blockFactory.newDoubleVectorFixedBuilder(docs.getPositionCount())) {
+            try (
+                BooleanVector.Builder builder = blockFactory.newBooleanVectorFixedBuilder(docs.getPositionCount());
+                DoubleVector.Builder scoreBuilder = blockFactory.newDoubleVectorFixedBuilder(docs.getPositionCount())
+            ) {
                 for (int i = 0; i < docs.getPositionCount(); i++) {
                     scoreSingleDocWithScorer(builder, scoreBuilder, docs.getInt(i));
                 }

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

@@ -9,7 +9,7 @@
 
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.ConstantNullVector;
+import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -87,7 +87,7 @@ default boolean eagerEvalSafeInLazy() {
          * Retrieves the score for the expression
          * Only expressions that can be scored, or that can combine scores, should override this method
          */
-        default Block score(Page page, BlockFactory blockFactory) {
+        default DoubleBlock score(Page page, BlockFactory blockFactory) {
             return blockFactory.newConstantDoubleBlockWith(0.0, page.getPositionCount());
         }
     }

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

@@ -8,7 +8,10 @@
 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;
@@ -17,13 +20,17 @@
 
 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 +39,33 @@ 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 (BooleanBlock filterResultBlock = (BooleanBlock) evaluator.eval(page)) {
+            if (filterResultBlock.areAllValuesNull()) {
                 // All results are null which is like false. No values selected.
                 page.releaseBlocks();
                 return null;
             }
+
             // 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 (filterResultBlock.getBoolean(filterResultBlock.getFirstValueIndex(p))) {
                     positions[rowCount++] = p;
                 }
             }
@@ -64,7 +74,14 @@ protected Page process(Page page) {
                 page.releaseBlocks();
                 return null;
             }
-            if (rowCount == page.getPositionCount()) {
+            final DoubleBlock scoreBlock;
+            if (usesScoring) {
+                scoreBlock = evaluator.score(page, blockFactory);
+            } else {
+                scoreBlock = null;
+            }
+
+            if (rowCount == page.getPositionCount() && (usesScoring == false || scoreBlock.asVector().isConstant())) {
                 return page;
             }
             positions = Arrays.copyOf(positions, rowCount);
@@ -73,10 +90,15 @@ 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, page.getBlock(i), scoreBlock, positions);
+                    } else {
+                        filteredBlocks[i] = page.getBlock(i).filter(positions);
+                    }
                 }
                 success = true;
             } finally {
+                Releasables.closeExpectNoException(scoreBlock);
                 page.releaseBlocks();
                 if (success == false) {
                     Releasables.closeExpectNoException(filteredBlocks);
@@ -86,6 +108,15 @@ protected Page process(Page page) {
         }
     }
 
+    private Block createScoresBlock(int rowCount, DoubleBlock originalScoreBlock, DoubleBlock newScoreBlock, 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(originalScoreBlock.getDouble(positions[j]) + newScoreBlock.getDouble(positions[j]));
+        }
+        return updatedScoresBuilder.build().asBlock().filter(positions);
+    }
+
     @Override
     public String toString() {
         return "FilterOperator[" + "evaluator=" + evaluator + ']';

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryExpressionEvaluatorTests.java

@@ -184,7 +184,9 @@ private List<Page> runQuery(Set<String> values, Query query, boolean shuffleDocs
             );
             LuceneQueryExpressionEvaluator luceneQueryEvaluator = new LuceneQueryExpressionEvaluator(
                 blockFactory,
-                new LuceneQueryExpressionEvaluator.ShardConfig[] { shard }, false);
+                new LuceneQueryExpressionEvaluator.ShardConfig[] { shard },
+                false
+            );
 
             List<Operator> operators = new ArrayList<>();
             if (shuffleDocs) {

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/FilterOperatorTests.java

@@ -66,7 +66,7 @@ public EvalOperator.ExpressionEvaluator get(DriverContext context) {
             public String toString() {
                 return "SameLastDigit[lhs=0, rhs=1]";
             }
-        });
+        }, false);
     }
 
     @Override
@@ -113,7 +113,7 @@ public void testReadFromBlock() {
             new SequenceBooleanBlockSourceOperator(context.blockFactory(), List.of(true, false, true, false))
         );
         List<Page> results = drive(
-            new FilterOperator.FilterOperatorFactory(dvrCtx -> new EvalOperatorTests.LoadFromPage(0)).get(context),
+            new FilterOperator.FilterOperatorFactory(dvrCtx -> new EvalOperatorTests.LoadFromPage(0), false).get(context),
             input.iterator(),
             context
         );