ESQL: Fix AsyncOperator status values and add emitted rows

docs/changelog/132738.yaml

@@ -0,0 +1,5 @@
+pr: 132738
+summary: Fix `AsyncOperator` status values and add emitted rows
+area: ES|QL
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -367,6 +367,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion RESOLVE_INDEX_MODE_ADDED = def(9_141_0_00);
     public static final TransportVersion DATA_STREAM_WRITE_INDEX_ONLY_SETTINGS = def(9_142_0_00);
     public static final TransportVersion SCRIPT_RESCORER = def(9_143_0_00);
+    public static final TransportVersion ESQL_LOOKUP_OPERATOR_EMITTED_ROWS = def(9_144_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -238,7 +238,7 @@ public IsBlockedResult isBlocked() {
 
     @Override
     public final Operator.Status status() {
-        return status(Math.max(0L, checkpoint.getMaxSeqNo()), Math.max(0L, checkpoint.getProcessedCheckpoint()), processNanos.sum());
+        return status(checkpoint.getMaxSeqNo() + 1, checkpoint.getProcessedCheckpoint() + 1, processNanos.sum());
     }
 
     protected Operator.Status status(long receivedPages, long completedPages, long processNanos) {
@@ -289,7 +289,7 @@ public long completedPages() {
             return completedPages;
         }
 
-        public long procesNanos() {
+        public long processNanos() {
             return processNanos;
         }
 
@@ -310,8 +310,8 @@ protected final XContentBuilder innerToXContent(XContentBuilder builder) throws
             if (builder.humanReadable()) {
                 builder.field("process_time", TimeValue.timeValueNanos(processNanos));
             }
-            builder.field("received_pages", receivedPages);
-            builder.field("completed_pages", completedPages);
+            builder.field("pages_received", receivedPages);
+            builder.field("pages_completed", completedPages);
             return builder;
         }
 

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

@@ -39,17 +39,17 @@ protected AsyncOperator.Status mutateInstance(AsyncOperator.Status in) throws IO
             case 0 -> new AsyncOperator.Status(
                 randomValueOtherThan(in.receivedPages(), ESTestCase::randomNonNegativeLong),
                 in.completedPages(),
-                in.procesNanos()
+                in.processNanos()
             );
             case 1 -> new AsyncOperator.Status(
                 in.receivedPages(),
                 randomValueOtherThan(in.completedPages(), ESTestCase::randomNonNegativeLong),
-                in.procesNanos()
+                in.processNanos()
             );
             case 2 -> new AsyncOperator.Status(
                 in.receivedPages(),
                 in.completedPages(),
-                randomValueOtherThan(in.procesNanos(), ESTestCase::randomNonNegativeLong)
+                randomValueOtherThan(in.processNanos(), ESTestCase::randomNonNegativeLong)
             );
             default -> throw new AssertionError("unknown ");
         };
@@ -62,8 +62,8 @@ public void testToXContent() {
             {
               "process_nanos" : 10,
               "process_time" : "10nanos",
-              "received_pages" : 100,
-              "completed_pages" : 50
+              "pages_received" : 100,
+              "pages_completed" : 50
             }"""));
     }
 }

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

@@ -87,7 +87,7 @@ protected Matcher<String> expectedToStringOfSimple() {
     }
 
     @Override
-    protected void assertEmptyStatus(Map<String, Object> map) {
+    protected void assertStatus(Map<String, Object> map, List<Page> input, List<Page> output) {
         assertThat(map, nullValue());
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitStatusTests.java → x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/LimitOperatorStatusTests.java

@@ -16,7 +16,7 @@
 
 import static org.hamcrest.Matchers.equalTo;
 
-public class LimitStatusTests extends AbstractWireSerializingTestCase<LimitOperator.Status> {
+public class LimitOperatorStatusTests extends AbstractWireSerializingTestCase<LimitOperator.Status> {
     public void testToXContent() {
         assertThat(Strings.toString(new LimitOperator.Status(10, 1, 1, 111, 222)), equalTo("""
             {"limit":10,"limit_remaining":1,"pages_processed":1,"rows_received":111,"rows_emitted":222}"""));

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

@@ -18,10 +18,13 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.LongStream;
 
 import static org.elasticsearch.compute.test.RandomBlock.randomElementType;
+import static org.elasticsearch.test.MapMatcher.assertMap;
+import static org.elasticsearch.test.MapMatcher.matchesMap;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.sameInstance;
 
@@ -192,4 +195,17 @@ Block randomBlock(BlockFactory blockFactory, int size) {
         }
         return RandomBlock.randomBlock(blockFactory, randomElementType(), size, false, 1, 1, 0, 0).block();
     }
+
+    @Override
+    protected final void assertStatus(Map<String, Object> map, List<Page> input, List<Page> output) {
+        var emittedRows = output.stream().mapToInt(Page::getPositionCount).sum();
+
+        var mapMatcher = matchesMap().entry("rows_received", emittedRows)
+            .entry("pages_processed", output.size())
+            .entry("rows_emitted", emittedRows)
+            .entry("limit", 100)
+            .entry("limit_remaining", 100 - emittedRows);
+
+        assertMap(map, mapMatcher);
+    }
 }

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.operator;
 
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.test.AnyOperatorTestCase;
 import org.hamcrest.Matcher;
 
@@ -52,7 +53,7 @@ public void testBigDescription() {
     }
 
     @Override
-    protected void assertEmptyStatus(Map<String, Object> map) {
+    protected void assertStatus(Map<String, Object> map, List<Page> input, List<Page> output) {
         assertThat(map, nullValue());
     }
 }

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.operator.SourceOperator;
@@ -23,11 +24,14 @@
 import org.hamcrest.Matcher;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.test.MapMatcher.assertMap;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.matchesPattern;
+import static org.hamcrest.Matchers.notNullValue;
 
 /**
  * Superclass for testing any {@link Operator}, including {@link SourceOperator}s.
@@ -106,39 +110,58 @@ public final void testSimpleToString() {
     }
 
     /**
-     * Ensures that the Operator.Status of this operator has the standard fields.
+     * Ensures that the Operator.Status of this operator has the standard fields, set to 0.
      */
-    public final void testOperatorStatus() throws IOException {
+    public void testEmptyOperatorStatus() {
         DriverContext driverContext = driverContext();
         try (var operator = simple().get(driverContext)) {
-            Operator.Status status = operator.status();
-            if (status == null) {
-                assertEmptyStatus(null);
-                return;
-            }
+            assertOperatorStatus(operator, List.of(), List.of());
+        }
+    }
 
-            try (var xContentBuilder = XContentBuilder.builder(XContentType.JSON.xContent())) {
-                status.toXContent(xContentBuilder, ToXContent.EMPTY_PARAMS);
+    /**
+     * Extracts and asserts the operator status.
+     */
+    protected final void assertOperatorStatus(Operator operator, List<Page> input, List<Page> output) {
+        Operator.Status status = operator.status();
 
-                var bytesReference = BytesReference.bytes(xContentBuilder);
-                var map = XContentHelper.convertToMap(bytesReference, false, xContentBuilder.contentType()).v2();
+        if (status == null) {
+            assertStatus(null, input, output);
+            return;
+        }
 
-                assertEmptyStatus(map);
-            }
+        var xContent = XContentType.JSON.xContent();
+        try (var xContentBuilder = XContentBuilder.builder(xContent)) {
+            status.toXContent(xContentBuilder, ToXContent.EMPTY_PARAMS);
+
+            var bytesReference = BytesReference.bytes(xContentBuilder);
+            var map = XContentHelper.convertToMap(bytesReference, false, xContentBuilder.contentType()).v2();
+
+            assertStatus(map, input, output);
+        } catch (IOException e) {
+            fail(e, "Failed to convert operator status to XContent");
         }
     }
 
     /**
      * Assert that the status is sane.
+     * <p>
+     *     This method should be overridden with custom logics and for better assertions and for operators without status.
+     * </p>
      */
-    protected void assertEmptyStatus(@Nullable Map<String, Object> map) {
+    protected void assertStatus(@Nullable Map<String, Object> map, List<Page> input, List<Page> output) {
+        assertThat(map, notNullValue());
+
+        var totalInputRows = input.stream().mapToInt(Page::getPositionCount).sum();
+        var totalOutputRows = output.stream().mapToInt(Page::getPositionCount).sum();
+
         MapMatcher matcher = matchesMap().extraOk();
         if (map.containsKey("pages_processed")) {
-            matcher = matcher.entry("pages_processed", 0);
+            matcher = matcher.entry("pages_processed", greaterThanOrEqualTo(0));
         } else {
-            matcher = matcher.entry("pages_received", 0).entry("pages_emitted", 0);
+            matcher = matcher.entry("pages_received", input.size()).entry("pages_emitted", output.size());
         }
-        matcher = matcher.entry("rows_received", 0).entry("rows_emitted", 0);
+        matcher = matcher.entry("rows_received", totalInputRows).entry("rows_emitted", totalOutputRows);
         assertMap(map, matcher);
     }
 

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AsyncOperatorTestCase.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.test;
+
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.test.MapMatcher;
+import org.hamcrest.Matcher;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.elasticsearch.test.MapMatcher.assertMap;
+import static org.elasticsearch.test.MapMatcher.matchesMap;
+import static org.hamcrest.Matchers.either;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+
+public abstract class AsyncOperatorTestCase extends OperatorTestCase {
+    @Override
+    @SuppressWarnings("unchecked")
+    protected final void assertStatus(Map<String, Object> map, List<Page> input, List<Page> output) {
+        var mapMatcher = matchesMap().entry("pages_received", input.size())
+            .entry("pages_completed", input.size())
+            .entry("process_nanos", either(greaterThanOrEqualTo(0)).or((Matcher<Integer>) (Matcher<?>) greaterThanOrEqualTo(0L)));
+
+        mapMatcher = extendStatusMatcher(mapMatcher, input, output);
+
+        assertMap(map, mapMatcher);
+    }
+
+    protected MapMatcher extendStatusMatcher(MapMatcher mapMatcher, List<Page> input, List<Page> output) {
+        return mapMatcher;
+    }
+}

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/OperatorTestCase.java

@@ -224,6 +224,7 @@ protected final void assertSimple(DriverContext context, int size) {
         var operator = simple().get(context);
         List<Page> results = drive(operator, input.iterator(), context);
         assertSimpleOutput(origInput, results);
+        assertOperatorStatus(operator, origInput, results);
         assertThat(context.breaker().getUsed(), equalTo(0L));
 
         // Release all result blocks. After this, all input blocks should be released as well, otherwise we have a leak.

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/SourceOperatorTestCase.java

@@ -7,14 +7,22 @@
 
 package org.elasticsearch.compute.test;
 
+import org.elasticsearch.compute.data.Page;
+
+import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.test.MapMatcher.assertMap;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
 
 public abstract class SourceOperatorTestCase extends AnyOperatorTestCase {
     @Override
-    protected void assertEmptyStatus(Map<String, Object> map) {
-        assertMap(map, matchesMap().extraOk().entry("pages_emitted", 0).entry("rows_emitted", 0));
+    protected void assertStatus(Map<String, Object> map, List<Page> input, List<Page> output) {
+        assertMap(
+            map,
+            matchesMap().extraOk()
+                .entry("pages_emitted", output.size())
+                .entry("rows_emitted", output.stream().mapToInt(Page::getPositionCount).sum())
+        );
     }
 }