ESQL: Add a random sample command

docs/changelog/123879.yaml

@@ -0,0 +1,5 @@
+pr: 123879
+summary: Add a random sample command
+area: ES|QL
+type: feature
+issues: []

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

@@ -184,6 +184,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion JINA_AI_EMBEDDING_TYPE_SUPPORT_ADDED = def(9_020_0_00);
     public static final TransportVersion RE_REMOVE_MIN_COMPATIBLE_SHARD_NODE = def(9_021_0_00);
     public static final TransportVersion INCLUDE_INDEX_MODE_IN_GET_DATA_STREAM = def(9_022_0_00);
+    public static final TransportVersion RANDOM_SAMPLER_QUERY_BUILDER = def(9_023_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -134,6 +134,7 @@
 import org.elasticsearch.search.aggregations.bucket.sampler.UnmappedSampler;
 import org.elasticsearch.search.aggregations.bucket.sampler.random.InternalRandomSampler;
 import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplerAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplingQueryBuilder;
 import org.elasticsearch.search.aggregations.bucket.terms.DoubleTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.LongRareTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.LongTerms;
@@ -1186,6 +1187,9 @@ private void registerQueryParsers(List<SearchPlugin> plugins) {
         registerQuery(new QuerySpec<>(ExactKnnQueryBuilder.NAME, ExactKnnQueryBuilder::new, parser -> {
             throw new IllegalArgumentException("[exact_knn] queries cannot be provided directly");
         }));
+        registerQuery(
+            new QuerySpec<>(RandomSamplingQueryBuilder.NAME, RandomSamplingQueryBuilder::new, RandomSamplingQueryBuilder::fromXContent)
+        );
 
         registerFromPlugin(plugins, SearchPlugin::getQueries, this::registerQuery);
     }

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

@@ -44,14 +44,34 @@ 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 strictly between 0.0 and 1.0, was [" + p + "]");
+        }
+    }
+
+    public double probability() {
+        return p;
+    }
+
+    public int seed() {
+        return seed;
+    }
+
+    public int hash() {
+        return hash;
+    }
+
     @Override
     public String toString(String field) {
         return "RandomSamplingQuery{" + "p=" + p + ", seed=" + seed + ", hash=" + hash + '}';
@@ -98,13 +118,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,149 @@
+/*
+ * 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.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplingQuery.checkProbabilityRange;
+import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
+import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
+
+public class RandomSamplingQueryBuilder extends AbstractQueryBuilder<RandomSamplingQueryBuilder> {
+
+    public static final String NAME = "random_sampling";
+    static final ParseField PROBABILITY = new ParseField("query");
+    static final ParseField SEED = new ParseField("seed");
+    static final ParseField HASH = new ParseField("hash");
+
+    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(StreamInput in) throws IOException {
+        super(in);
+        this.probability = in.readDouble();
+        this.seed = in.readInt();
+        this.hash = in.readInt();
+    }
+
+    public RandomSamplingQueryBuilder hash(Integer hash) {
+        this.hash = hash;
+        return this;
+    }
+
+    public double probability() {
+        return probability;
+    }
+
+    public int seed() {
+        return seed;
+    }
+
+    public int hash() {
+        return hash;
+    }
+
+    @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.getPreferredName(), probability);
+        builder.field(SEED.getPreferredName(), seed);
+        builder.field(HASH.getPreferredName(), hash);
+        builder.endObject();
+    }
+
+    private static final ConstructingObjectParser<RandomSamplingQueryBuilder, Void> PARSER = new ConstructingObjectParser<>(
+        NAME,
+        false,
+        args -> {
+            var randomSamplingQueryBuilder = new RandomSamplingQueryBuilder((double) args[0]);
+            if (args[1] != null) {
+                randomSamplingQueryBuilder.seed((int) args[1]);
+            }
+            if (args[2] != null) {
+                randomSamplingQueryBuilder.hash((int) args[2]);
+            }
+            return randomSamplingQueryBuilder;
+        }
+    );
+
+    static {
+        PARSER.declareDouble(constructorArg(), PROBABILITY);
+        PARSER.declareInt(optionalConstructorArg(), SEED);
+        PARSER.declareInt(optionalConstructorArg(), HASH);
+    }
+
+    public static RandomSamplingQueryBuilder fromXContent(XContentParser parser) throws IOException {
+        return PARSER.apply(parser, null);
+    }
+
+    @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.RANDOM_SAMPLER_QUERY_BUILDER;
+    }
+}

server/src/test/java/org/elasticsearch/search/SearchModuleTests.java

@@ -444,6 +444,7 @@ public CheckedBiConsumer<ShardSearchRequest, StreamOutput, IOException> getReque
         "range",
         "regexp",
         "knn_score_doc",
+        "random_sampling",
         "script",
         "script_score",
         "simple_query_string",

server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplingQueryBuilderTests.java

@@ -0,0 +1,75 @@
+/*
+ * 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.index.query.SearchExecutionContext;
+import org.elasticsearch.test.AbstractQueryTestCase;
+import org.elasticsearch.xcontent.XContentParseException;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class RandomSamplingQueryBuilderTests extends AbstractQueryTestCase<RandomSamplingQueryBuilder> {
+
+    @Override
+    protected RandomSamplingQueryBuilder doCreateTestQueryBuilder() {
+        double probability = randomDoubleBetween(0.0, 1.0, false);
+        var builder = new RandomSamplingQueryBuilder(probability);
+        if (randomBoolean()) {
+            builder.seed(randomInt());
+        }
+        if (randomBoolean()) {
+            builder.hash(randomInt());
+        }
+        return builder;
+    }
+
+    @Override
+    protected void doAssertLuceneQuery(RandomSamplingQueryBuilder queryBuilder, Query query, SearchExecutionContext context)
+        throws IOException {
+        var rsQuery = asInstanceOf(RandomSamplingQuery.class, query);
+        assertThat(rsQuery.probability(), equalTo(queryBuilder.probability()));
+        assertThat(rsQuery.seed(), equalTo(queryBuilder.seed()));
+        assertThat(rsQuery.hash(), equalTo(queryBuilder.hash()));
+    }
+
+    @Override
+    protected boolean supportsBoost() {
+        return false;
+    }
+
+    @Override
+    protected boolean supportsQueryName() {
+        return false;
+    }
+
+    @Override
+    public void testUnknownField() {
+        var json = "{ \""
+            + RandomSamplingQueryBuilder.NAME
+            + "\" : {\"bogusField\" : \"someValue\", \""
+            + RandomSamplingQueryBuilder.PROBABILITY.getPreferredName()
+            + "\" : \""
+            + randomBoolean()
+            + "\", \""
+            + RandomSamplingQueryBuilder.SEED.getPreferredName()
+            + "\" : \""
+            + randomInt()
+            + "\", \""
+            + RandomSamplingQueryBuilder.HASH.getPreferredName()
+            + "\" : \""
+            + randomInt()
+            + "\" } }";
+        var e = expectThrows(XContentParseException.class, () -> parseQuery(json));
+        assertTrue(e.getMessage().contains("bogusField"));
+    }
+}

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.compute.operator;
 
-import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -69,20 +68,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);
         }
     }