Simplify RandomSampleOperator

Author: jan-elastic

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

@@ -7,9 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import com.carrotsearch.hppc.BitMixer;
-
-import org.apache.lucene.search.DocIdSetIterator;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
@@ -22,43 +19,32 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.List;
+import java.util.Deque;
+import java.util.LinkedList;
 import java.util.Objects;
 import java.util.SplittableRandom;
 
 public class RandomSampleOperator implements Operator {
 
-    // The threshold for the number of rows to collect in a batch before starting sampling it.
-    private static final int ROWS_BATCH_THRESHOLD = 10_000;
-    // How many batches can be to keep in memory and still accept new input Pages.
-    // Besides these many buffered batches, the operator holds an additional batch that's being sampled.
-    private static final int MAX_BUFFERED_BATCHES = 1;
-
-    private final double probability;
-    private final int seed;
-
-    private boolean collecting = true;
-    private boolean isFinished = false;
-    private final PageBatching pageBatching;
-    private BatchSampling currentSampling;
+    private boolean finished;
+    private final Deque<Page> outputPages;
+    private final RandomSamplingQuery.RandomSamplingIterator randomSamplingIterator;
 
     private int pagesCollected = 0;
     private int pagesEmitted = 0;
     private int rowsCollected = 0;
     private int rowsEmitted = 0;
-    private int batchesSampled = 0;
 
     private long collectNanos;
     private long emitNanos;
 
     public RandomSampleOperator(double probability, int seed) {
-        this.probability = probability;
-        this.seed = seed;
-        // TODO derive the threshold from the probability and a max cap
-        pageBatching = new PageBatching(ROWS_BATCH_THRESHOLD, MAX_BUFFERED_BATCHES);
+        finished = false;
+        outputPages = new LinkedList<>();
+        SplittableRandom random = new SplittableRandom(seed);
+        randomSamplingIterator = new RandomSamplingQuery.RandomSamplingIterator(Integer.MAX_VALUE, probability, random::nextInt);
+        randomSamplingIterator.nextDoc();
     }
 
     public record Factory(double probability, int seed) implements OperatorFactory {
@@ -79,7 +65,7 @@ public String describe() {
      */
     @Override
     public boolean needsInput() {
-        return collecting && pageBatching.capacityAvailable();
+        return finished == false;
     }
 
     /**
@@ -91,80 +77,64 @@ public boolean needsInput() {
     @Override
     public void addInput(Page page) {
         final var addStart = System.nanoTime();
-        collect(page);
+        createOutputPage(page);
+        rowsCollected += page.getPositionCount();
+        pagesCollected++;
+        page.releaseBlocks();
         collectNanos += System.nanoTime() - addStart;
     }
 
-    private void collect(Page page) {
-        pagesCollected++;
-        rowsCollected += page.getPositionCount();
-        pageBatching.addPage(page);
+    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;
+        }
+        if (sampledIdx > 0) {
+            outputPages.add(page.filter(Arrays.copyOf(sampledPositions, sampledIdx)));
+        }
     }
 
     /**
      * notifies the operator that it won't receive any more input pages
      */
     @Override
     public void finish() {
-        if (collecting && rowsCollected > 0) { // finish() can be called multiple times
-            pageBatching.flush();
-        }
-        collecting = false;
+        finished = true;
     }
 
     /**
      * whether the operator has finished processing all input pages and made the corresponding output pages available
      */
     @Override
     public boolean isFinished() {
-        return isFinished;
+        return finished && outputPages.isEmpty();
     }
 
-    /**
-     * returns non-null if output page available. Only called when isFinished() == false
-     *
-     * @throws UnsupportedOperationException if the operator is a {@link SinkOperator}
-     */
     @Override
     public Page getOutput() {
         final var emitStart = System.nanoTime();
-        Page page = emit();
+        Page page;
+        if (outputPages.isEmpty()) {
+            page = null;
+        } else {
+            page = outputPages.removeFirst();
+            pagesEmitted++;
+            rowsEmitted += page.getPositionCount();
+        }
         emitNanos += System.nanoTime() - emitStart;
         return page;
     }
 
-    private Page emit() {
-        if (currentSampling == null) {
-            if (pageBatching.hasNext() == false) {
-                if (collecting == false) {
-                    isFinished = true;
-                }
-                return null; // not enough pages on the input yet
-            }
-            final var currentBatch = pageBatching.next();
-            currentSampling = new BatchSampling(currentBatch, probability, seed);
-            batchesSampled++;
-        }
-
-        final var page = currentSampling.next();
-        if (page != null) {
-            rowsEmitted += page.getPositionCount();
-            pagesEmitted++;
-            return page;
-        } // else: current batch is exhausted
-
-        currentSampling.close();
-        currentSampling = null;
-        return emit();
-    }
-
     /**
      * notifies the operator that it won't be used anymore (i.e. none of the other methods called),
      * and its resources can be cleaned up
      */
     @Override
     public void close() {
-        pageBatching.close();
+        for (Page page : outputPages) {
+            page.releaseBlocks();
+        }
     }
 
     @Override
@@ -174,175 +144,12 @@ public String toString() {
 
     @Override
     public Operator.Status status() {
-        return new Status(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted, batchesSampled);
-    }
-
-    private static class SamplingIterator {
-
-        private final RandomSamplingQuery.RandomSamplingIterator samplingIterator;
-        private int nextDoc = -1;
-
-        SamplingIterator(int maxDoc, double probability, int seed) {
-            final SplittableRandom random = new SplittableRandom(BitMixer.mix(seed));
-            samplingIterator = new RandomSamplingQuery.RandomSamplingIterator(maxDoc, probability, random::nextInt);
-            advance();
-        }
-
-        boolean hasNext() {
-            return nextDoc != DocIdSetIterator.NO_MORE_DOCS;
-        }
-
-        int next() {
-            return nextDoc;
-        }
-
-        void advance() {
-            assert hasNext() : "No more docs to sample";
-            nextDoc = samplingIterator.nextDoc();
-        }
-    }
-
-    private record PagesBatch(ArrayDeque<Page> batch, int rowCount) {}
-
-    private static class PageBatching {
-
-        private final int collectingRowThreshold;
-        private final int maxBufferedBatches;
-
-        private final List<PagesBatch> batches = new ArrayList<>();
-
-        private int collectingBatchRowCount = 0;
-        private ArrayDeque<Page> collectingBatch = new ArrayDeque<>();
-
-        PageBatching(int collectingRowThreshold, int maxBufferedBatches) {
-            this.collectingRowThreshold = collectingRowThreshold;
-            this.maxBufferedBatches = maxBufferedBatches;
-        }
-
-        void addPage(Page page) {
-            collectingBatch.add(page);
-            collectingBatchRowCount += page.getPositionCount();
-            if (collectingBatchRowCount >= collectingRowThreshold) {
-                rotate();
-            }
-        }
-
-        private void rotate() {
-            batches.add(new PagesBatch(collectingBatch, collectingBatchRowCount));
-            collectingBatch = new ArrayDeque<>();
-            collectingBatchRowCount = 0;
-        }
-
-        boolean hasNext() {
-            return batches.isEmpty() == false;
-        }
-
-        PagesBatch next() {
-            return batches.removeFirst();
-        }
-
-        public boolean capacityAvailable() {
-            return batches.size() < maxBufferedBatches;
-        }
-
-        void flush() {
-            while (batches.isEmpty() == false) {
-                var batch = batches.removeFirst();
-                collectingBatch.addAll(batch.batch);
-                collectingBatchRowCount += batch.rowCount;
-            }
-            if (collectingBatch.isEmpty() == false) {
-                rotate();
-            }
-        }
-
-        void close() {
-            assert batches.isEmpty() : "There are still available batches";
-            assert collectingBatch.isEmpty() : "Current batch has not been rotated";
-        }
+        return new Status(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted);
     }
 
-    private static class BatchSampling {
-
-        private final ArrayDeque<Page> pagesDeque;
-        private final SamplingIterator samplingIterator;
-
-        private int rowsProcessed = 0;
-
-        BatchSampling(PagesBatch batch, double probability, int seed) {
-            pagesDeque = batch.batch;
-            samplingIterator = new SamplingIterator(batch.rowCount, probability, seed);
-        }
-
-        Page next() {
-            while (pagesDeque.isEmpty() == false) {
-                final var page = pagesDeque.poll();
-                final int positionCount = page.getPositionCount();
-                final int[] sampledPositions = new int[positionCount];
-                int sampledIdx = 0;
-
-                while (true) {
-                    if (samplingIterator.hasNext()) {
-                        var docOffset = samplingIterator.next() - rowsProcessed;
-                        if (docOffset < positionCount) {
-                            sampledPositions[sampledIdx++] = docOffset;
-                            samplingIterator.advance();
-                        } else {
-                            // position falls outside the current page
-                            break;
-                        }
-                    } else {
-                        // no more docs to sample
-                        drainPages();
-                        break;
-                    }
-                }
-                rowsProcessed += positionCount;
-
-                if (sampledIdx > 0) {
-                    var filter = Arrays.copyOf(sampledPositions, sampledIdx);
-                    return page.filter(filter);
-                } // else: fetch a new page (if any left)
-
-                releasePage(page);
-            }
-
-            return null;
-        }
-
-        private void drainPages() {
-            Page page;
-            do {
-                page = pagesDeque.poll();
-            } while (releasePage(page));
-        }
-
-        /**
-         * Returns true if there was a non-null page that was released.
-         */
-        private static boolean releasePage(Page page) {
-            if (page != null) {
-                page.releaseBlocks();
-                return true;
-            }
-            return false;
-        }
-
-        void close() {
-            assert pagesDeque.isEmpty() : "There are still unreleased pages";
-            assert samplingIterator.hasNext() == false : "There are still docs to sample";
-        }
-    }
-
-    private record Status(
-        long collectNanos,
-        long emitNanos,
-        int pagesCollected,
-        int pagesEmitted,
-        int rowsCollected,
-        int rowsEmitted,
-        int batchesSampled
-    ) implements Operator.Status {
+    private record Status(long collectNanos, long emitNanos, int pagesCollected, int pagesEmitted, int rowsCollected, int rowsEmitted)
+        implements
+            Operator.Status {
 
         public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
             Operator.Status.class,
@@ -357,7 +164,6 @@ private record Status(
                 streamInput.readVInt(),
                 streamInput.readVInt(),
                 streamInput.readVInt(),
-                streamInput.readVInt(),
                 streamInput.readVInt()
             );
         }
@@ -370,7 +176,6 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeVInt(pagesEmitted);
             out.writeVInt(rowsCollected);
             out.writeVInt(rowsEmitted);
-            out.writeVInt(batchesSampled);
         }
 
         @Override
@@ -393,7 +198,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field("pages_emitted", pagesEmitted);
             builder.field("rows_collected", rowsCollected);
             builder.field("rows_emitted", rowsEmitted);
-            builder.field("batches_sampled", batchesSampled);
             return builder.endObject();
         }
 
@@ -407,13 +211,12 @@ public boolean equals(Object o) {
                 && pagesCollected == other.pagesCollected
                 && pagesEmitted == other.pagesEmitted
                 && rowsCollected == other.rowsCollected
-                && rowsEmitted == other.rowsEmitted
-                && batchesSampled == other.batchesSampled;
+                && rowsEmitted == other.rowsEmitted;
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted, batchesSampled);
+            return Objects.hash(collectNanos, emitNanos, pagesCollected, pagesEmitted, rowsCollected, rowsEmitted);
         }
 
         @Override