Merge branch 'main' into adding-settings-to-data-streams

Author: masseyke

docs/changelog/125570.yaml

@@ -0,0 +1,5 @@
+pr: 125570
+summary: ES|QL random sampling
+area: Machine Learning
+type: feature
+issues: []

modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/ScaledFloatFieldBlockLoaderTests.java

@@ -9,8 +9,8 @@
 
 package org.elasticsearch.index.mapper.extras;
 
+import org.elasticsearch.datageneration.FieldType;
 import org.elasticsearch.index.mapper.NumberFieldBlockLoaderTestCase;
-import org.elasticsearch.logsdb.datageneration.FieldType;
 import org.elasticsearch.plugins.Plugin;
 
 import java.util.Collection;

modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java

@@ -47,7 +47,9 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.ClusterSettings;
@@ -63,6 +65,7 @@
 import org.elasticsearch.http.HttpBodyTracer;
 import org.elasticsearch.http.HttpServerTransport;
 import org.elasticsearch.http.HttpTransportSettings;
+import org.elasticsearch.index.IndexingPressure;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.rest.BaseRestHandler;
@@ -73,13 +76,17 @@
 import org.elasticsearch.rest.RestResponse;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.ClusterServiceUtils;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.MockLog;
 import org.elasticsearch.test.junit.annotations.TestLogging;
+import org.elasticsearch.test.rest.ObjectPath;
 import org.elasticsearch.transport.Transports;
 import org.elasticsearch.transport.netty4.Netty4Utils;
+import org.elasticsearch.xcontent.json.JsonXContent;
 
 import java.nio.channels.ClosedChannelException;
+import java.nio.charset.StandardCharsets;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -102,6 +109,7 @@
 import static io.netty.handler.codec.http.HttpVersion.HTTP_1_1;
 import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
@@ -361,7 +369,12 @@ public void test413TooLargeOnExpect100Continue() throws Exception {
 
     // ensures that oversized chunked encoded request has maxContentLength limit and returns 413
     public void testOversizedChunkedEncoding() throws Exception {
-        try (var clientContext = newClientContext(t -> {/* ignore exception from e.g. server closing socket */})) {
+        try (
+            var clientContext = newClientContext(
+                internalCluster().getRandomNodeName(),
+                t -> {/* ignore exception from e.g. server closing socket */}
+            )
+        ) {
             var opaqueId = clientContext.newOpaqueId();
             final var requestBodyStream = new HttpChunkedInput(
                 new ChunkedStream(new ByteBufInputStream(Unpooled.wrappedBuffer(randomByteArrayOfLength(MAX_CONTENT_LENGTH + 1)))),
@@ -501,6 +514,72 @@ private void assertHttpBodyLogging(Consumer<ClientContext> test) throws Exceptio
         }
     }
 
+    public void testBulkIndexingRequestSplitting() throws Exception {
+        final var watermarkBytes = between(100, 200);
+        final var tinyNode = internalCluster().startCoordinatingOnlyNode(
+            Settings.builder()
+                .put(IndexingPressure.SPLIT_BULK_LOW_WATERMARK.getKey(), ByteSizeValue.ofBytes(watermarkBytes))
+                .put(IndexingPressure.SPLIT_BULK_LOW_WATERMARK_SIZE.getKey(), ByteSizeValue.ofBytes(watermarkBytes))
+                .build()
+        );
+
+        try (var clientContext = newClientContext(tinyNode, cause -> ExceptionsHelper.maybeDieOnAnotherThread(new AssertionError(cause)))) {
+            final var request = new DefaultHttpRequest(HTTP_1_1, POST, "/_bulk");
+            request.headers().add(CONTENT_TYPE, APPLICATION_JSON);
+            HttpUtil.setTransferEncodingChunked(request, true);
+
+            final var channel = clientContext.channel();
+            channel.writeAndFlush(request);
+
+            final var indexName = randomIdentifier();
+            final var indexCreatedListener = ClusterServiceUtils.addTemporaryStateListener(
+                cs -> Iterators.filter(
+                    cs.metadata().indicesAllProjects().iterator(),
+                    indexMetadata -> indexMetadata.getIndex().getName().equals(indexName)
+                ).hasNext()
+            );
+
+            indexCreatedListener.addListener(ActionListener.running(() -> logger.info("--> index created")));
+
+            final var valueLength = between(10, 30);
+            final var docSizeBytes = "{'field':''}".length() + valueLength;
+            final var itemCount = between(watermarkBytes / docSizeBytes + 1, 300); // enough to split at least once
+            assertThat(itemCount * docSizeBytes, greaterThan(watermarkBytes));
+            for (int i = 0; i < itemCount; i++) {
+                channel.write(new DefaultHttpContent(Unpooled.wrappedBuffer(Strings.format("""
+                    {"index":{"_index":"%s"}}
+                    {"field":"%s"}
+                    """, indexName, randomAlphaOfLength(valueLength)).getBytes(StandardCharsets.UTF_8))));
+            }
+
+            channel.flush();
+            safeAwait(indexCreatedListener); // index must be created before we finish sending the request
+
+            channel.writeAndFlush(new DefaultLastHttpContent());
+            final var response = clientContext.getNextResponse();
+            try {
+                assertEquals(RestStatus.OK.getStatus(), response.status().code());
+                final ObjectPath responseBody;
+                final var copy = response.content().copy(); // Netty4Utils doesn't handle direct buffers, so copy to heap first
+                try {
+                    responseBody = ObjectPath.createFromXContent(JsonXContent.jsonXContent, Netty4Utils.toBytesReference(copy));
+                } finally {
+                    copy.release();
+                }
+                assertFalse(responseBody.evaluate("errors"));
+                assertEquals(itemCount, responseBody.evaluateArraySize("items"));
+                for (int i = 0; i < itemCount; i++) {
+                    assertEquals(
+                        RestStatus.CREATED.getStatus(),
+                        (int) asInstanceOf(int.class, responseBody.evaluateExact("items", Integer.toString(i), "index", "status"))
+                    );
+                }
+            } finally {
+                response.release();
+            }
+        }
+    }
+
     static FullHttpRequest fullHttpRequest(String opaqueId, ByteBuf content) {
         var request = new DefaultFullHttpRequest(HTTP_1_1, POST, ControlServerRequestPlugin.ROUTE, Unpooled.wrappedBuffer(content));
         request.headers().add(CONTENT_LENGTH, content.readableBytes());
@@ -539,11 +618,13 @@ protected boolean addMockHttpTransport() {
     private static final LongSupplier idGenerator = new AtomicLong()::getAndIncrement;
 
     private ClientContext newClientContext() throws Exception {
-        return newClientContext(cause -> ExceptionsHelper.maybeDieOnAnotherThread(new AssertionError(cause)));
+        return newClientContext(
+            internalCluster().getRandomNodeName(),
+            cause -> ExceptionsHelper.maybeDieOnAnotherThread(new AssertionError(cause))
+        );
     }
 
-    private ClientContext newClientContext(Consumer<Throwable> exceptionHandler) throws Exception {
-        var nodeName = internalCluster().getRandomNodeName();
+    private ClientContext newClientContext(String nodeName, Consumer<Throwable> exceptionHandler) throws Exception {
         var clientResponseQueue = new LinkedBlockingDeque<FullHttpResponse>(16);
         final var httpServerTransport = internalCluster().getInstance(HttpServerTransport.class, nodeName);
         var remoteAddr = randomFrom(httpServerTransport.boundAddress().boundAddresses());
@@ -556,7 +637,7 @@ private ClientContext newClientContext(Consumer<Throwable> exceptionHandler) thr
                 protected void initChannel(SocketChannel ch) {
                     var p = ch.pipeline();
                     p.addLast(new HttpClientCodec());
-                    p.addLast(new HttpObjectAggregator(ByteSizeUnit.KB.toIntBytes(4)));
+                    p.addLast(new HttpObjectAggregator(ByteSizeUnit.MB.toIntBytes(4)));
                     p.addLast(new SimpleChannelInboundHandler<FullHttpResponse>() {
                         @Override
                         protected void channelRead0(ChannelHandlerContext ctx, FullHttpResponse msg) {

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

@@ -226,7 +226,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion SYNONYMS_REFRESH_PARAM = def(9_060_0_00);
     public static final TransportVersion DOC_FIELDS_AS_LIST = def(9_061_0_00);
     public static final TransportVersion DENSE_VECTOR_OFF_HEAP_STATS = def(9_062_00_0);
-    public static final TransportVersion SETTINGS_IN_DATA_STREAMS = def(9_063_00_0);
+    public static final TransportVersion RANDOM_SAMPLER_QUERY_BUILDER = def(9_063_0_00);
+    public static final TransportVersion SETTINGS_IN_DATA_STREAMS = def(9_064_00_0);
 
     /*
      * 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);