Add a random sample commadn

This adds RANDOM_SAMPLE <probability> <seed>? command.

Author: bpintea

server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplingQuery.java

@@ -44,14 +44,22 @@ public final class RandomSamplingQuery extends Query {
      *                  can be generated
      */
     public RandomSamplingQuery(double p, int seed, int hash) {
-        if (p <= 0.0 || p >= 1.0) {
-            throw new IllegalArgumentException("RandomSampling probability must be between 0.0 and 1.0, was [" + p + "]");
-        }
+        checkProbabilityRange(p);
         this.p = p;
         this.seed = seed;
         this.hash = hash;
     }
 
+    /**
+     * Verifies that the probability is within the (0.0, 1.0) range.
+     * @throws IllegalArgumentException in case of an invalid probability.
+     */
+    public static void checkProbabilityRange(double p) throws IllegalArgumentException {
+        if (p <= 0.0 || p >= 1.0) {
+            throw new IllegalArgumentException("RandomSampling probability must be between 0.0 and 1.0, was [" + p + "]");
+        }
+    }
+
     @Override
     public String toString(String field) {
         return "RandomSamplingQuery{" + "p=" + p + ", seed=" + seed + ", hash=" + hash + '}';
@@ -98,13 +106,13 @@ public void visit(QueryVisitor visitor) {
     /**
      * A DocIDSetIter that skips a geometrically random number of documents
      */
-    static class RandomSamplingIterator extends DocIdSetIterator {
+    public static class RandomSamplingIterator extends DocIdSetIterator {
         private final int maxDoc;
         private final double p;
         private final FastGeometric distribution;
         private int doc = -1;
 
-        RandomSamplingIterator(int maxDoc, double p, IntSupplier rng) {
+        public RandomSamplingIterator(int maxDoc, double p, IntSupplier rng) {
             this.maxDoc = maxDoc;
             this.p = p;
             this.distribution = new FastGeometric(rng, p);

server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplingQueryBuilder.java

@@ -0,0 +1,97 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.aggregations.bucket.sampler.random;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.Randomness;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplingQuery.checkProbabilityRange;
+
+public class RandomSamplingQueryBuilder extends AbstractQueryBuilder<RandomSamplingQueryBuilder> {
+
+    public static final String NAME = "random_sampling";
+
+    private final double probability;
+    private int seed = Randomness.get().nextInt();
+    private int hash = 0;
+
+    public RandomSamplingQueryBuilder(double probability) {
+        checkProbabilityRange(probability);
+        this.probability = probability;
+    }
+
+    public RandomSamplingQueryBuilder seed(int seed) {
+        checkProbabilityRange(probability);
+        this.seed = seed;
+        return this;
+    }
+
+    public RandomSamplingQueryBuilder hash(Integer hash) {
+        this.hash = hash;
+        return this;
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeDouble(probability);
+        out.writeInt(seed);
+        out.writeInt(hash);
+    }
+
+    @Override
+    protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(NAME);
+        builder.field("probability", probability);
+        builder.field("seed", seed);
+        builder.field("hash", hash);
+        builder.endObject();
+
+    }
+
+    @Override
+    protected Query doToQuery(SearchExecutionContext context) throws IOException {
+        return new RandomSamplingQuery(probability, seed, hash);
+    }
+
+    @Override
+    protected boolean doEquals(RandomSamplingQueryBuilder other) {
+        return probability == other.probability && seed == other.seed && hash == other.hash;
+    }
+
+    @Override
+    protected int doHashCode() {
+        return Objects.hash(probability, seed, hash);
+    }
+
+    /**
+     * Returns the name of the writeable object
+     */
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    /**
+     * The minimal version of the recipient this object can be sent to
+     */
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.ZERO;
+    }
+}

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/MetadataAttribute.java

@@ -172,6 +172,10 @@ public static boolean isSupported(String name) {
         return ATTRIBUTES_MAP.containsKey(name);
     }
 
+    public static boolean isScoreAttribute(Expression a) {
+        return a instanceof MetadataAttribute ma && ma.name().equals(SCORE);
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (false == super.equals(obj)) {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Page.java

@@ -293,4 +293,21 @@ public Page projectBlocks(int[] blockMapping) {
             }
         }
     }
+
+    public Page filter(int... positions) {
+        Block[] filteredBlocks = new Block[blocks.length];
+        boolean success = false;
+        try {
+            for (int i = 0; i < blocks.length; i++) {
+                filteredBlocks[i] = getBlock(i).filter(positions);
+            }
+            success = true;
+        } finally {
+            releaseBlocks();
+            if (success == false) {
+                Releasables.closeExpectNoException(filteredBlocks);
+            }
+        }
+        return new Page(filteredBlocks);
+    }
 }

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

@@ -69,20 +69,7 @@ protected Page process(Page page) {
             }
             positions = Arrays.copyOf(positions, rowCount);
 
-            Block[] filteredBlocks = new Block[page.getBlockCount()];
-            boolean success = false;
-            try {
-                for (int i = 0; i < page.getBlockCount(); i++) {
-                    filteredBlocks[i] = page.getBlock(i).filter(positions);
-                }
-                success = true;
-            } finally {
-                page.releaseBlocks();
-                if (success == false) {
-                    Releasables.closeExpectNoException(filteredBlocks);
-                }
-            }
-            return new Page(filteredBlocks);
+            return page.filter(positions);
         }
     }