SampleOperaetorTests + fix status

Author: jan-elastic

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

@@ -31,9 +31,8 @@ public class SampleOperator implements Operator {
     private final Deque<Page> outputPages;
     private final RandomSamplingQuery.RandomSamplingIterator randomSamplingIterator;
 
-    private int pagesCollected = 0;
-    private int pagesEmitted = 0;
-    private int rowsCollected = 0;
+    private int pagesProcessed = 0;
+    private int rowsReceived = 0;
     private int rowsEmitted = 0;
 
     private long collectNanos;
@@ -76,19 +75,19 @@ public boolean needsInput() {
      */
     @Override
     public void addInput(Page page) {
-        final var addStart = System.nanoTime();
+        long startTime = System.nanoTime();
         createOutputPage(page);
-        rowsCollected += page.getPositionCount();
-        pagesCollected++;
+        rowsReceived += page.getPositionCount();
         page.releaseBlocks();
-        collectNanos += System.nanoTime() - addStart;
+        pagesProcessed++;
+        collectNanos += System.nanoTime() - startTime;
     }
 
     private void createOutputPage(Page page) {
         final int[] sampledPositions = new int[page.getPositionCount()];
         int sampledIdx = 0;
-        for (int i = randomSamplingIterator.docID(); i - rowsCollected < page.getPositionCount(); i = randomSamplingIterator.nextDoc()) {
-            sampledPositions[sampledIdx++] = i - rowsCollected;
+        for (int i = randomSamplingIterator.docID(); i - rowsReceived < page.getPositionCount(); i = randomSamplingIterator.nextDoc()) {
+            sampledPositions[sampledIdx++] = i - rowsReceived;
         }
         if (sampledIdx > 0) {
             outputPages.add(page.filter(Arrays.copyOf(sampledPositions, sampledIdx)));
@@ -119,7 +118,6 @@ public Page getOutput() {
             page = null;
         } else {
             page = outputPages.removeFirst();
-            pagesEmitted++;
             rowsEmitted += page.getPositionCount();
         }
         emitNanos += System.nanoTime() - emitStart;
@@ -139,15 +137,15 @@ public void close() {
 
     @Override
     public String toString() {
-        return "SampleOperator[sampled = " + rowsEmitted + "/" + rowsCollected + "]";
+        return "SampleOperator[sampled = " + rowsEmitted + "/" + rowsReceived + "]";
     }
 
     @Override
     public Operator.Status status() {
-        return new Status(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted);
+        return new Status(collectNanos, emitNanos, pagesProcessed, rowsReceived, rowsEmitted);
     }
 
-    private record Status(long collectNanos, long emitNanos, int pagesCollected, int pagesEmitted, int rowsCollected, int rowsEmitted)
+    private record Status(long collectNanos, long emitNanos, int pagesProcessed, int rowsReceived, int rowsEmitted)
         implements
             Operator.Status {
 
@@ -158,23 +156,15 @@ private record Status(long collectNanos, long emitNanos, int pagesCollected, int
         );
 
         Status(StreamInput streamInput) throws IOException {
-            this(
-                streamInput.readVLong(),
-                streamInput.readVLong(),
-                streamInput.readVInt(),
-                streamInput.readVInt(),
-                streamInput.readVInt(),
-                streamInput.readVInt()
-            );
+            this(streamInput.readVLong(), streamInput.readVLong(), streamInput.readVInt(), streamInput.readVInt(), streamInput.readVInt());
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeVLong(collectNanos);
             out.writeVLong(emitNanos);
-            out.writeVInt(pagesCollected);
-            out.writeVInt(pagesEmitted);
-            out.writeVInt(rowsCollected);
+            out.writeVInt(pagesProcessed);
+            out.writeVInt(rowsReceived);
             out.writeVInt(rowsEmitted);
         }
 
@@ -194,9 +184,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             if (builder.humanReadable()) {
                 builder.field("emit_time", TimeValue.timeValueNanos(emitNanos));
             }
-            builder.field("pages_collected", pagesCollected);
-            builder.field("pages_emitted", pagesEmitted);
-            builder.field("rows_collected", rowsCollected);
+            builder.field("pages_processed", pagesProcessed);
+            builder.field("rows_received", rowsReceived);
             builder.field("rows_emitted", rowsEmitted);
             return builder.endObject();
         }
@@ -208,15 +197,14 @@ public boolean equals(Object o) {
             Status other = (Status) o;
             return collectNanos == other.collectNanos
                 && emitNanos == other.emitNanos
-                && pagesCollected == other.pagesCollected
-                && pagesEmitted == other.pagesEmitted
-                && rowsCollected == other.rowsCollected
+                && pagesProcessed == other.pagesProcessed
+                && rowsReceived == other.rowsReceived
                 && rowsEmitted == other.rowsEmitted;
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted);
+            return Objects.hash(collectNanos, emitNanos, pagesProcessed, rowsReceived, rowsEmitted);
         }
 
         @Override

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

@@ -0,0 +1,53 @@
+/*
+ * 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.operator;
+
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.test.OperatorTestCase;
+import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
+import org.hamcrest.Matcher;
+
+import java.util.List;
+import java.util.stream.LongStream;
+
+import static org.hamcrest.Matchers.closeTo;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.matchesPattern;
+
+public class SampleOperatorTests extends OperatorTestCase {
+
+    @Override
+    protected SourceOperator simpleInput(BlockFactory blockFactory, int size) {
+        return new SequenceLongBlockSourceOperator(blockFactory, LongStream.range(0, size));
+    }
+
+    @Override
+    protected void assertSimpleOutput(List<Page> input, List<Page> results) {
+        int inputCount = input.stream().mapToInt(Page::getPositionCount).sum();
+        int outputCount = results.stream().mapToInt(Page::getPositionCount).sum();
+        double meanExpectedOutputCount = 0.5 * inputCount;
+        double stdDevExpectedOutputCount = Math.sqrt(meanExpectedOutputCount);
+        assertThat((double) outputCount, closeTo(meanExpectedOutputCount, 10 * stdDevExpectedOutputCount));
+    }
+
+    @Override
+    protected Operator.OperatorFactory simple() {
+        return new SampleOperator.Factory(0.5, randomInt());
+    }
+
+    @Override
+    protected Matcher<String> expectedDescriptionOfSimple() {
+        return matchesPattern("SampleOperator\\[probability = 0.5, seed = -?\\d+]");
+    }
+
+    @Override
+    protected Matcher<String> expectedToStringOfSimple() {
+        return equalTo("SampleOperator[sampled = 0/0]");
+    }
+}