Make ScoreOperator and LuceneQueryEvaluator more robust

Author: carlosdelest

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

@@ -61,10 +61,16 @@ protected LuceneQueryEvaluator(BlockFactory blockFactory, ShardConfig[] shards)
     }
 
     public Block executeQuery(Page page) {
-        // Lucene based operators retrieve DocVectors as first block
-        Block block = page.getBlock(0);
-        assert block instanceof DocBlock : "LuceneQueryExpressionEvaluator expects DocBlock as input";
-        DocVector docs = (DocVector) block.asVector();
+        // Search for DocVector block
+        Block docBlock = null;
+        for (int i = 0; i < page.getBlockCount(); i++) {
+            if (page.getBlock(i) instanceof DocBlock) {
+                docBlock = page.getBlock(i);
+                break;
+            }
+        }
+        assert docBlock != null : "LuceneQueryExpressionEvaluator expects a DocBlock";
+        DocVector docs = (DocVector) docBlock.asVector();
         try {
             if (docs.singleSegmentNonDecreasing()) {
                 return evalSingleSegmentNonDecreasing(docs).asBlock();

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

@@ -9,7 +9,6 @@
 
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.DocVector;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.Page;
@@ -46,9 +45,9 @@ public ScoreOperator(BlockFactory blockFactory, ExpressionScorer scorer, int sco
 
     @Override
     protected Page process(Page page) {
-        assert page.getBlockCount() >= 2 : "Expected at least 2 blocks, got " + page.getBlockCount();
-        assert page.getBlock(0).asVector() instanceof DocVector : "Expected a DocVector, got " + page.getBlock(0).asVector();
-        assert page.getBlock(1).asVector() instanceof DoubleVector : "Expected a DoubleVector, got " + page.getBlock(1).asVector();
+        assert page.getBlockCount() > scoreBlockPosition : "Expected to get a score block in position " + scoreBlockPosition;
+        assert page.getBlock(scoreBlockPosition).asVector() instanceof DoubleVector
+            : "Expected a DoubleVector as a score block, got " + page.getBlock(scoreBlockPosition).asVector();
 
         Block[] blocks = new Block[page.getBlockCount()];
         for (int i = 0; i < page.getBlockCount(); i++) {