From e280a91aa1390aee728c2faacddbcecc9aff6174 Mon Sep 17 00:00:00 2001 From: Oleg Lvovitch Date: Fri, 24 Jul 2026 19:33:13 +0100 Subject: [PATCH 1/4] ESQL: add chunked reader for record boundary probes Record boundary probing needs a few hundred bytes from an arbitrary offset, but the length it must declare up front is unknown, so it declared the whole remainder of the object. That range is undrainable, so releasing it required abortStream, and an aborted response destroys its HTTP connection instead of returning it to the pool. Every probe then re-paid a TCP connect and TLS handshake. ChunkedStorageInputStream reads a range as bounded, fully consumed chunks via StorageObject#readBytes. Each chunk is a closed range read to completion, so it closes cleanly and its connection is reusable. Sizes grow geometrically from 32KB, so a normal record costs one request and a pathological one costs a logarithmic number. Chunk boundaries are invisible to consumers: splitter lookahead past a lone CR reads the first byte of the next chunk, and a splitter byte budget keeps counting across chunks. ChunkedProbeEquivalenceTests pins that, comparing boundaries against a plain stream at every offset with chunk sizes down to one byte. An off-by-one here would be a silently wrong row count, not a failure, so the equivalence is the safety argument. No call site uses this yet. --- .../ChunkedStorageInputStream.java | 147 ++++++++++++ .../ChunkedProbeEquivalenceTests.java | 214 ++++++++++++++++++ .../ChunkedStorageInputStreamTests.java | 139 ++++++++++++ .../datasources/RecordingStorageObject.java | 105 +++++++++ 4 files changed, 605 insertions(+) create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStream.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedProbeEquivalenceTests.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStreamTests.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/RecordingStorageObject.java diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStream.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStream.java new file mode 100644 index 0000000000000..f8e2626a33973 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStream.java @@ -0,0 +1,147 @@ +/* + * 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.xpack.esql.datasources; + +import org.elasticsearch.xpack.esql.datasources.spi.StorageObject; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.util.Objects; + +/** + * Reads a byte range of a {@link StorageObject} as a stream of bounded, fully-consumed chunks. + *

+ * Record-boundary probing needs a few hundred bytes from an arbitrary offset, but the length it must declare + * up front is unknown — a record can legitimately run to {@code maxRecordBytes}. Declaring the whole remainder + * of the object (the previous approach) made the request undrainable, so the only way to release it was + * {@link StorageObject#abortStream}, and an aborted response destroys its HTTP connection instead of returning + * it to the pool. Every probe then re-paid a TCP connect and TLS handshake, which is most of the cost of a + * cross-region probe and all of the reason ~150 serial probes took ~73s (esql-planning#1528). + *

+ * Fetching in bounded chunks removes the dilemma: each chunk is a closed range that is read to completion, so + * it closes cleanly and its connection is reusable — the same reason {@code S3StorageObject#fetchMetadata} + * drains its one-byte length probe rather than aborting it. Chunk sizes grow geometrically from + * {@link #FIRST_CHUNK_SIZE}, so a normal record costs a single request and a pathological one costs a + * logarithmic number, with total transfer within roughly twice the bytes actually consumed. + *

+ * Consumers see a plain {@link InputStream}: chunk boundaries are invisible, so a record splitter's lookahead + * (for instance peeking past a lone {@code CR} to test for {@code LF}) simply reads the first byte of the next + * chunk, and a splitter's own byte budget keeps counting across chunks unchanged. + */ +final class ChunkedStorageInputStream extends InputStream { + + /** + * Size of the first fetch. Large enough that a single request answers essentially any real record, small + * enough that the bytes wasted when only a line's prefix is needed stay far below the stride being probed. + */ + static final int FIRST_CHUNK_SIZE = 32 * 1024; + + /** Ceiling on chunk growth, so one probe of a pathological record cannot turn into a single huge request. */ + static final int MAX_CHUNK_SIZE = 4 * 1024 * 1024; + + private final StorageObject object; + private final long endExclusive; + private final int maxChunkSize; + + /** Absolute position of the next byte to fetch. */ + private long fetchPosition; + private int nextChunkSize; + + private byte[] chunk = new byte[0]; + private int chunkPosition; + private int chunkLength; + private boolean exhausted; + + ChunkedStorageInputStream(StorageObject object, long position, long endExclusive) { + this(object, position, endExclusive, FIRST_CHUNK_SIZE, MAX_CHUNK_SIZE); + } + + ChunkedStorageInputStream(StorageObject object, long position, long endExclusive, int firstChunkSize, int maxChunkSize) { + if (position < 0) { + throw new IllegalArgumentException("position must be non-negative, got: " + position); + } + if (firstChunkSize <= 0 || maxChunkSize < firstChunkSize) { + throw new IllegalArgumentException("require 0 < firstChunkSize <= maxChunkSize, got: " + firstChunkSize + ", " + maxChunkSize); + } + this.object = Objects.requireNonNull(object); + this.endExclusive = endExclusive; + this.fetchPosition = position; + this.nextChunkSize = firstChunkSize; + this.maxChunkSize = maxChunkSize; + } + + @Override + public int read() throws IOException { + if (ensureAvailable() == false) { + return -1; + } + return chunk[chunkPosition++] & 0xFF; + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + Objects.checkFromIndexSize(off, len, b.length); + if (len == 0) { + return 0; + } + if (ensureAvailable() == false) { + return -1; + } + int n = Math.min(len, chunkLength - chunkPosition); + System.arraycopy(chunk, chunkPosition, b, off, n); + chunkPosition += n; + return n; + } + + /** + * Ensures at least one unread byte is buffered, fetching the next chunk only when the current one is spent. + * A read that is satisfied from the current chunk never triggers a fetch, so the read that finds a record + * boundary cannot pull a speculative chunk across the wire. + * + * @return false once the range is exhausted + */ + private boolean ensureAvailable() throws IOException { + if (chunkPosition < chunkLength) { + return true; + } + if (exhausted || fetchPosition >= endExclusive) { + return false; + } + int size = (int) Math.min(nextChunkSize, endExclusive - fetchPosition); + if (chunk.length < size) { + chunk = new byte[size]; + } + ByteBuffer target = ByteBuffer.wrap(chunk, 0, size); + // Positional reads may return short (a FileChannel-backed provider is entitled to), so refill until the + // chunk is full or the object ends. + while (target.hasRemaining()) { + int read = object.readBytes(fetchPosition + target.position(), target); + if (read <= 0) { + break; + } + } + int filled = target.position(); + if (filled == 0) { + exhausted = true; + return false; + } + chunkPosition = 0; + chunkLength = filled; + fetchPosition += filled; + nextChunkSize = (int) Math.min((long) nextChunkSize * 2, maxChunkSize); + return true; + } + + /** + * No-op: every chunk is fully consumed by the time it is replaced, so this stream holds no provider + * resource between fetches and there is nothing to release — which is the point of reading in chunks. + */ + @Override + public void close() {} +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedProbeEquivalenceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedProbeEquivalenceTests.java new file mode 100644 index 0000000000000..b60ffe3dde2e2 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedProbeEquivalenceTests.java @@ -0,0 +1,214 @@ +/* + * 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.xpack.esql.datasources; + +import org.elasticsearch.common.breaker.NoopCircuitBreaker; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.util.BigArrays; +import org.elasticsearch.compute.data.BlockFactory; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.esql.datasource.csv.CsvFormatReader; +import org.elasticsearch.xpack.esql.datasource.ndjson.NdJsonFormatReader; +import org.elasticsearch.xpack.esql.datasources.spi.RecordSplitter; +import org.elasticsearch.xpack.esql.datasources.spi.SegmentableFormatReader; +import org.hamcrest.Matchers; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; + +/** + * The safety argument for {@link ChunkedStorageInputStream}. + *

+ * Reading in chunks must be invisible to record splitting: for every offset, probing through the chunked stream + * has to return exactly the boundary that probing through a plain whole-remainder stream returns. This is not a + * nice-to-have equivalence — a boundary that lands one byte off does not fail, it silently shifts where one + * split ends and the next begins, and the query returns a wrong row count. So the tests here run adversarially + * small chunks, which puts chunk edges on every interesting byte: inside a record, on a terminator, between the + * {@code CR} and {@code LF} of a {@code CRLF} pair, and on the last byte before the range ends. + */ +public class ChunkedProbeEquivalenceTests extends ESTestCase { + + private static final int MAX_RECORD_BYTES = SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES; + + private BlockFactory blockFactory() { + return BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build(); + } + + private SegmentableFormatReader plainCsvReader() { + return (SegmentableFormatReader) new CsvFormatReader(blockFactory()).withConfig(Map.of("mode", "plain")); + } + + private SegmentableFormatReader quotedCsvReader() { + return (SegmentableFormatReader) new CsvFormatReader(blockFactory()).withConfig(Map.of("mode", "quoted")); + } + + private SegmentableFormatReader ndjsonReader() { + return new NdJsonFormatReader(Settings.EMPTY, blockFactory(), List.of()); + } + + /** + * Probes every offset of a payload both ways and asserts the answers agree. Chunk sizes are deliberately + * tiny so the chunk grid sweeps across every byte of every record. + */ + private void assertBoundaryEquivalence(SegmentableFormatReader reader, byte[] payload, int maxRecordBytes) throws IOException { + RecordSplitter splitter = reader.recordSplitter(maxRecordBytes); + assertTrue("this corpus is only meaningful for a strided splitter", splitter.supportsStridedProbing()); + RecordingStorageObject object = new RecordingStorageObject(payload); + + for (int pos = 0; pos < payload.length; pos++) { + long viaPlainStream; + try (InputStream plain = object.newStream(pos, payload.length - pos)) { + viaPlainStream = splitter.findNextRecordBoundary(plain); + } + + int firstChunk = randomIntBetween(1, 17); + long viaChunkedStream; + try (InputStream chunked = new ChunkedStorageInputStream(object, pos, payload.length, firstChunk, firstChunk * 2)) { + viaChunkedStream = splitter.findNextRecordBoundary(chunked); + } + + assertEquals("boundary disagreed at offset " + pos + " with first chunk " + firstChunk, viaPlainStream, viaChunkedStream); + } + assertEquals("the chunked probe path must never abort", 0, object.abortCalls.get()); + } + + public void testNdjsonBoundaryEquivalenceAcrossEveryOffset() throws IOException { + StringBuilder ndjson = new StringBuilder(); + for (int i = 0; i < 60; i++) { + ndjson.append("{\"id\":").append(i).append(",\"name\":\"value_").append(i).append("\"}\n"); + } + assertBoundaryEquivalence(ndjsonReader(), ndjson.toString().getBytes(StandardCharsets.UTF_8), MAX_RECORD_BYTES); + } + + public void testNdjsonWithCrlfBoundaryEquivalence() throws IOException { + StringBuilder ndjson = new StringBuilder(); + for (int i = 0; i < 60; i++) { + ndjson.append("{\"id\":").append(i).append("}\r\n"); + } + assertBoundaryEquivalence(ndjsonReader(), ndjson.toString().getBytes(StandardCharsets.UTF_8), MAX_RECORD_BYTES); + } + + public void testNdjsonWithLoneCarriageReturnsBoundaryEquivalence() throws IOException { + StringBuilder ndjson = new StringBuilder(); + for (int i = 0; i < 60; i++) { + ndjson.append("{\"id\":").append(i).append("}\r"); + } + assertBoundaryEquivalence(ndjsonReader(), ndjson.toString().getBytes(StandardCharsets.UTF_8), MAX_RECORD_BYTES); + } + + public void testNdjsonWithMixedTerminatorsAndRaggedRecordsBoundaryEquivalence() throws IOException { + StringBuilder ndjson = new StringBuilder(); + for (int i = 0; i < 80; i++) { + ndjson.append("{\"id\":").append(i).append(",\"pad\":\"").append("x".repeat(randomIntBetween(0, 40))).append("\"}"); + switch (randomIntBetween(0, 2)) { + case 0 -> ndjson.append('\n'); + case 1 -> ndjson.append("\r\n"); + default -> ndjson.append('\r'); + } + } + assertBoundaryEquivalence(ndjsonReader(), ndjson.toString().getBytes(StandardCharsets.UTF_8), MAX_RECORD_BYTES); + } + + public void testPayloadWithNoTrailingTerminatorBoundaryEquivalence() throws IOException { + byte[] payload = "{\"id\":1}\n{\"id\":2}\n{\"id\":3-unterminated".getBytes(StandardCharsets.UTF_8); + assertBoundaryEquivalence(ndjsonReader(), payload, MAX_RECORD_BYTES); + } + + public void testPlainCsvBoundaryEquivalenceAcrossEveryOffset() throws IOException { + StringBuilder csv = new StringBuilder("id,name\n"); + for (int i = 0; i < 60; i++) { + csv.append(i).append(",value_").append(i).append('\n'); + } + assertBoundaryEquivalence(plainCsvReader(), csv.toString().getBytes(StandardCharsets.UTF_8), MAX_RECORD_BYTES); + } + + /** + * A record longer than the byte budget must reach the same verdict either way. The budget is counted by the + * splitter as it consumes bytes, so it has to survive those bytes arriving in many chunks rather than one. + */ + public void testRecordExceedingMaxRecordBytesReachesSameVerdict() throws IOException { + byte[] payload = ("{\"id\":1}\n" + "y".repeat(4096) + "\n{\"id\":2}\n").getBytes(StandardCharsets.UTF_8); + int budget = 512; + RecordSplitter splitter = ndjsonReader().recordSplitter(budget); + RecordingStorageObject object = new RecordingStorageObject(payload); + + int pos = 9; // first byte of the oversized record + long viaPlainStream; + try (InputStream plain = object.newStream(pos, payload.length - pos)) { + viaPlainStream = splitter.findNextRecordBoundary(plain); + } + long viaChunkedStream; + try (InputStream chunked = new ChunkedStorageInputStream(object, pos, payload.length, 7, 32)) { + viaChunkedStream = splitter.findNextRecordBoundary(chunked); + } + + assertEquals(RecordSplitter.RECORD_TOO_LARGE, viaPlainStream); + assertEquals("the budget must be counted identically across chunk boundaries", viaPlainStream, viaChunkedStream); + } + + /** The proven-probe and exact-walk shapes read the same stream, so they need the same guarantee. */ + public void testQuotedCsvProvenProbeAndExactWalkEquivalence() throws IOException { + StringBuilder csv = new StringBuilder("id,name\n"); + for (int i = 0; i < 60; i++) { + csv.append(i).append(",\"value ").append(i).append("\"\n"); + } + byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8); + + RecordSplitter splitter = quotedCsvReader().recordSplitter(MAX_RECORD_BYTES); + assumeTrue("quoted CSV must support proven probing for this test", splitter.supportsProvenProbing()); + RecordingStorageObject object = new RecordingStorageObject(payload); + + for (int pos = 1; pos < payload.length; pos++) { + long provenPlain; + try (InputStream plain = object.newStream(pos, payload.length - pos)) { + provenPlain = splitter.findProvenRecordBoundary(plain); + } + long provenChunked; + try (InputStream chunked = new ChunkedStorageInputStream(object, pos, payload.length, randomIntBetween(1, 13), 32)) { + provenChunked = splitter.findProvenRecordBoundary(chunked); + } + assertEquals("proven probe disagreed at offset " + pos, provenPlain, provenChunked); + + long walkPlain; + try (InputStream plain = object.newStream(0, payload.length)) { + walkPlain = splitter.findRecordStartAtOrAfter(plain, pos, () -> false); + } + long walkChunked; + try (InputStream chunked = new ChunkedStorageInputStream(object, 0, payload.length, randomIntBetween(1, 13), 32)) { + walkChunked = splitter.findRecordStartAtOrAfter(chunked, pos, () -> false); + } + assertEquals("exact walk disagreed at min skip " + pos, walkPlain, walkChunked); + } + assertEquals("the chunked probe path must never abort", 0, object.abortCalls.get()); + } + + /** The chunked probe must not read materially more than the plain probe consumed to find the same answer. */ + public void testProbeTransfersStayBoundedRelativeToTheRecordLength() throws IOException { + StringBuilder ndjson = new StringBuilder(); + for (int i = 0; i < 5_000; i++) { + ndjson.append("{\"id\":").append(i).append(",\"name\":\"value_").append(i).append("\"}\n"); + } + byte[] payload = ndjson.toString().getBytes(StandardCharsets.UTF_8); + RecordingStorageObject object = new RecordingStorageObject(payload); + RecordSplitter splitter = ndjsonReader().recordSplitter(MAX_RECORD_BYTES); + + try (InputStream chunked = new ChunkedStorageInputStream(object, payload.length / 2, payload.length)) { + assertThat(splitter.findNextRecordBoundary(chunked), Matchers.greaterThanOrEqualTo(0L)); + } + + assertEquals("a short record must be answered by one request", 1, object.readBytesCalls.size()); + assertThat( + "one request must not exceed the first chunk size", + object.readBytesCalls.get(0)[1], + Matchers.lessThanOrEqualTo((long) ChunkedStorageInputStream.FIRST_CHUNK_SIZE) + ); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStreamTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStreamTests.java new file mode 100644 index 0000000000000..9ae704705d47e --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ChunkedStorageInputStreamTests.java @@ -0,0 +1,139 @@ +/* + * 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.xpack.esql.datasources; + +import org.elasticsearch.test.ESTestCase; +import org.hamcrest.Matchers; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; + +/** + * Request-shape tests for {@link ChunkedStorageInputStream}. What matters here is not only that the bytes come + * out right, but that they arrive as bounded closed ranges that are read to completion — that property is what + * lets the provider pool the connection, and it is the whole reason this stream exists. + */ +public class ChunkedStorageInputStreamTests extends ESTestCase { + + public void testBoundaryInsideFirstChunkCostsOneRequest() throws IOException { + byte[] payload = "line one\nline two\n".getBytes(StandardCharsets.UTF_8); + RecordingStorageObject object = new RecordingStorageObject(payload); + + try (InputStream stream = new ChunkedStorageInputStream(object, 0, payload.length)) { + assertEquals('l', stream.read()); + } + + assertEquals("a read inside the first chunk must not fetch a second", 1, object.readBytesCalls.size()); + assertEquals("the probe path must never open a raw stream", 0, object.newStreamCalls.size()); + assertEquals("the probe path must never abort", 0, object.abortCalls.get()); + } + + public void testChunkSizesGrowGeometricallyAndAreCapped() throws IOException { + byte[] payload = new byte[600]; + RecordingStorageObject object = new RecordingStorageObject(payload); + + // firstChunk 32, cap 128: expect 32, 64, 128, 128, ... + try (InputStream stream = new ChunkedStorageInputStream(object, 0, payload.length, 32, 128)) { + assertEquals(payload.length, stream.readAllBytes().length); + } + + assertThat(object.readBytesCalls.get(0)[1], Matchers.equalTo(32L)); + assertThat(object.readBytesCalls.get(1)[1], Matchers.equalTo(64L)); + assertThat(object.readBytesCalls.get(2)[1], Matchers.equalTo(128L)); + assertThat(object.readBytesCalls.get(3)[1], Matchers.equalTo(128L)); + } + + public void testNoRequestReachesPastTheDeclaredEnd() throws IOException { + byte[] payload = new byte[1000]; + RecordingStorageObject object = new RecordingStorageObject(payload); + long end = 300; + + try (InputStream stream = new ChunkedStorageInputStream(object, 100, end, 64, 1024)) { + assertEquals(200, stream.readAllBytes().length); + } + + for (long[] call : object.readBytesCalls) { + assertThat( + "request [" + call[0] + "," + call[1] + ") must stay within the declared end", + call[0] + call[1], + Matchers.lessThanOrEqualTo(end) + ); + } + } + + public void testContentMatchesTheUnderlyingRangeAcrossChunkSizes() throws IOException { + byte[] payload = randomByteArrayOfLength(randomIntBetween(1, 4096)); + int from = randomIntBetween(0, payload.length - 1); + RecordingStorageObject object = new RecordingStorageObject(payload, randomIntBetween(1, 64)); + + byte[] read; + try ( + InputStream stream = new ChunkedStorageInputStream( + object, + from, + payload.length, + randomIntBetween(1, 128), + randomIntBetween(128, 512) + ) + ) { + read = stream.readAllBytes(); + } + + byte[] expected = new byte[payload.length - from]; + System.arraycopy(payload, from, expected, 0, expected.length); + assertArrayEquals("short positional reads must be refilled, not truncated", expected, read); + } + + public void testSingleByteAndBulkReadsAgree() throws IOException { + byte[] payload = randomByteArrayOfLength(randomIntBetween(1, 2048)); + RecordingStorageObject object = new RecordingStorageObject(payload); + + byte[] oneAtATime = new byte[payload.length]; + try (InputStream stream = new ChunkedStorageInputStream(object, 0, payload.length, 16, 64)) { + for (int i = 0; i < payload.length; i++) { + int b = stream.read(); + assertThat("stream ended early at " + i, b, Matchers.greaterThanOrEqualTo(0)); + oneAtATime[i] = (byte) b; + } + assertEquals("stream must be exhausted", -1, stream.read()); + } + assertArrayEquals(payload, oneAtATime); + } + + public void testEmptyRangeReadsAsEof() throws IOException { + RecordingStorageObject object = new RecordingStorageObject(new byte[100]); + + try (InputStream stream = new ChunkedStorageInputStream(object, 50, 50)) { + assertEquals(-1, stream.read()); + assertEquals(-1, stream.read(new byte[8], 0, 8)); + } + + assertEquals("an empty range must not issue a request", 0, object.readBytesCalls.size()); + } + + public void testEndBeyondObjectStopsAtObjectEnd() throws IOException { + byte[] payload = new byte[10]; + RecordingStorageObject object = new RecordingStorageObject(payload); + + try (InputStream stream = new ChunkedStorageInputStream(object, 0, 1_000_000)) { + assertEquals(payload.length, stream.readAllBytes().length); + assertEquals(-1, stream.read()); + } + } + + public void testZeroLengthReadReturnsZeroWithoutFetching() throws IOException { + RecordingStorageObject object = new RecordingStorageObject(new byte[100]); + + try (InputStream stream = new ChunkedStorageInputStream(object, 0, 100)) { + assertEquals(0, stream.read(new byte[8], 0, 0)); + } + + assertEquals(0, object.readBytesCalls.size()); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/RecordingStorageObject.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/RecordingStorageObject.java new file mode 100644 index 0000000000000..e73e0aa53cd62 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/RecordingStorageObject.java @@ -0,0 +1,105 @@ +/* + * 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.xpack.esql.datasources; + +import org.elasticsearch.xpack.esql.datasources.spi.StorageObject; +import org.elasticsearch.xpack.esql.datasources.spi.StoragePath; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * In-memory {@link StorageObject} that records every request it serves, so a test can assert on the shape of the + * I/O rather than only on its result. + *

+ * This is the instrument for the connection-reuse claim behind {@link ChunkedStorageInputStream}: a request is + * poolable exactly when it is a closed range that is read to completion, so the counters a test cares about are + * how many requests were issued, how long each one declared itself to be, and whether any of them had to be + * aborted. Aborts are the failure signal — the probe path is expected to produce none. + *

+ * {@code maxBytesPerReadCall} makes positional reads return short, which providers backed by a + * {@code FileChannel} are entitled to do, so callers that assume a single call fills the buffer are caught here + * rather than in production. + */ +final class RecordingStorageObject implements StorageObject { + + private final byte[] data; + private final int maxBytesPerReadCall; + private final StoragePath path; + + /** One {@code [position, requestedLength]} entry per {@link #readBytes} call, in call order. */ + final List readBytesCalls = new CopyOnWriteArrayList<>(); + /** One {@code [position, requestedLength]} entry per {@link #newStream(long, long)} call, in call order. */ + final List newStreamCalls = new CopyOnWriteArrayList<>(); + final AtomicInteger abortCalls = new AtomicInteger(); + + RecordingStorageObject(byte[] data) { + this(data, Integer.MAX_VALUE); + } + + RecordingStorageObject(byte[] data, int maxBytesPerReadCall) { + if (maxBytesPerReadCall <= 0) { + throw new IllegalArgumentException("maxBytesPerReadCall must be positive, got: " + maxBytesPerReadCall); + } + this.data = data; + this.maxBytesPerReadCall = maxBytesPerReadCall; + this.path = StoragePath.of("s3://bucket/recording.data"); + } + + @Override + public int readBytes(long position, ByteBuffer target) { + readBytesCalls.add(new long[] { position, target.remaining() }); + if (position >= data.length) { + return -1; + } + int available = (int) Math.min(data.length - position, target.remaining()); + int n = Math.min(available, maxBytesPerReadCall); + target.put(data, (int) position, n); + return n; + } + + @Override + public InputStream newStream(long position, long length) { + newStreamCalls.add(new long[] { position, length }); + int from = (int) Math.min(position, data.length); + int to = length == READ_TO_END ? data.length : (int) Math.min(position + length, data.length); + return new ByteArrayInputStream(data, from, Math.max(0, to - from)); + } + + @Override + public void abortStream(InputStream stream) throws IOException { + abortCalls.incrementAndGet(); + stream.close(); + } + + @Override + public long length() { + return data.length; + } + + @Override + public Instant lastModified() { + return Instant.EPOCH; + } + + @Override + public boolean exists() { + return true; + } + + @Override + public StoragePath path() { + return path; + } +} From 51f9264771d448a6d4566ce02b33878706332f6d Mon Sep 17 00:00:00 2001 From: Oleg Lvovitch Date: Fri, 24 Jul 2026 19:35:54 +0100 Subject: [PATCH 2/4] ESQL: probe record boundaries without aborting the read computeRecordAlignedMacroSplitStarts opened a range covering the rest of the file for every stride probe, read a few hundred bytes to the next record boundary, and aborted. The range was undrainable, so abort was the only exit, and an aborted response destroys its connection rather than returning it to the pool. Each of the ~150 probes over a 9.8GB ndjson file therefore re-paid a TCP connect and TLS handshake, which is most of the ~0.5s per probe seen cross-region. All three probe shapes (strided, proven, exact walk) now read through ChunkedStorageInputStream, so each request is a closed range consumed in full and closing it hands back a poolable connection. Boundaries are unchanged. The existing boundary-value tests pass untouched; the two tests that pinned the abort now assert its absence, since asserting we abort was asserting the defect. Does not close the issue on its own: ~150 probes remain serial, which is a floor no per-probe saving gets under. --- .../esql/datasources/FileSplitProvider.java | 17 +++++++-------- .../datasources/FileSplitProviderTests.java | 19 ++++++++++------- .../StorageObjectAbortChainTests.java | 21 +++++++------------ 3 files changed, 26 insertions(+), 31 deletions(-) diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java index 214c225652f3d..0fa8814d24a8c 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java @@ -55,7 +55,6 @@ import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual; import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals; -import java.io.Closeable; import java.io.IOException; import java.io.InputStream; import java.util.ArrayList; @@ -1078,10 +1077,11 @@ static List computeRecordAlignedMacroSplitStarts( } long boundary; if (strided) { - InputStream stream = storageObject.newStream(pos, remaining); - // Abort rather than close: findNextRecordBoundary reads only a prefix of the range - // (fileLength - pos bytes), but close() on providers like S3 drains the remainder. - try (Closeable abortOnExit = () -> storageObject.abortStream(stream)) { + // Probe through bounded chunks: each is a closed range read to completion, so closing it + // releases a connection the provider can pool. Declaring the whole remainder instead (the + // previous shape) left abort as the only exit, and an abort destroys the connection, so every + // probe re-paid a connect and TLS handshake. + try (InputStream stream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { long skipped = splitter.findNextRecordBoundary(stream); if (skipped == RecordSplitter.RECORD_TOO_LARGE || skipped < 0) { break; @@ -1090,8 +1090,7 @@ static List computeRecordAlignedMacroSplitStarts( } } else { long probed; - InputStream probeStream = storageObject.newStream(pos, remaining); - try (Closeable abortOnExit = () -> storageObject.abortStream(probeStream)) { + try (InputStream probeStream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { probed = splitter.findProvenRecordBoundary(probeStream); } if (probed >= 0) { @@ -1099,10 +1098,8 @@ static List computeRecordAlignedMacroSplitStarts( } else if (probed == RecordSplitter.AMBIGUOUS) { // Bounded probe could not prove a boundary near pos; fall back to an exact walk from the last // proven record start. minSkip is stream-relative (pos - exactCursor) and always > 0. - long walkRemaining = fileLength - exactCursor; - InputStream walkStream = storageObject.newStream(exactCursor, walkRemaining); long start; - try (Closeable abortOnExit = () -> storageObject.abortStream(walkStream)) { + try (InputStream walkStream = new ChunkedStorageInputStream(storageObject, exactCursor, fileLength)) { start = splitter.findRecordStartAtOrAfter(walkStream, pos - exactCursor, isCancelled); } if (start == RecordSplitter.RECORD_TOO_LARGE || start < 0) { diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java index 1ecbd10eaa9e0..9eb464e15ce11 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java @@ -78,7 +78,7 @@ import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.instanceOf; -import static org.hamcrest.Matchers.lessThan; +import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.eq; @@ -1019,11 +1019,13 @@ private static Set trueRecordStarts(RecordSplitter splitter, byte[] payloa } /** - * Regression guard: {@link FileSplitProvider#computeRecordAlignedMacroSplitStarts} opens a - * range stream for each stride probe, reads only enough bytes to find the next record - * boundary, then must call {@link StorageObject#abortStream} — not a draining {@code close()}. + * Regression guard on the shape of the I/O, not just its result: each stride probe must read bounded + * chunks that it consumes in full, so the provider can pool the connection. Neither of the two failure + * modes is acceptable — draining a huge range wastes the transfer, and aborting to avoid the drain + * destroys the connection and makes the next probe re-handshake, which is what made cross-region + * discovery cost ~0.5s per probe. */ - public void testComputeRecordAlignedMacroSplitStartsDoesNotDrainStream() throws IOException { + public void testComputeRecordAlignedMacroSplitStartsUsesPooledChunkReads() throws IOException { var blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build(); StringBuilder csv = new StringBuilder("id,name\n"); @@ -1050,11 +1052,12 @@ public void testComputeRecordAlignedMacroSplitStartsDoesNotDrainStream() throws ); assertThat("expected multiple macro-split boundaries", starts.size(), greaterThan(1)); - assertTrue("each boundary probe must abort the underlying stream", tracking.abortCalls.get() >= starts.size() - 1); + assertEquals("boundary probes must not abort — an aborted response is not a poolable connection", 0, tracking.abortCalls.get()); + assertFalse("no probe may reach the abort path", tracking.aborted.get()); assertThat( - "boundary probes must not drain the range streams; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", + "boundary probes must read bounded chunks; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", tracking.bytesConsumed.get(), - lessThan(fileLength / 2) + lessThanOrEqualTo((long) starts.size() * ChunkedStorageInputStream.FIRST_CHUNK_SIZE) ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java index 5e2d45614a661..8dc227d6dc420 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java @@ -88,12 +88,12 @@ public void testAbortPropagatesThroughDecoratorChainWithoutDrain() throws IOExce } /** - * Regression guard for {@link FileSplitProvider#computeRecordAlignedMacroSplitStarts}: each - * macro-split boundary probe opens {@code newStream(pos, remaining)} and reads only a prefix, - * so cleanup must abort (not drain) through the same decorator chain used for uncompressed - * text files on object storage. + * Regression guard for {@link FileSplitProvider#computeRecordAlignedMacroSplitStarts}: each macro-split + * boundary probe reads bounded chunks that it consumes in full, so cleanup neither drains a large + * remainder nor aborts — through the same decorator chain used for uncompressed text files on object + * storage, whose {@code readBytes} delegation is what carries the chunked reads down to the provider. */ - public void testMacroSplitDiscoveryAbortPropagatesThroughDecoratorChainWithoutDrain() throws IOException { + public void testMacroSplitDiscoveryProbesFullyConsumeClosedRangesThroughDecoratorChain() throws IOException { StringBuilder csv = new StringBuilder("id,name\n"); for (int i = 0; i < 200_000; i++) { csv.append(i).append(",value_").append(i).append('\n'); @@ -123,15 +123,10 @@ public void testMacroSplitDiscoveryAbortPropagatesThroughDecoratorChainWithoutDr ); assertThat("expected multiple macro-split boundaries", starts.size(), Matchers.greaterThan(1)); - assertTrue("each probe must abort the raw stream", tracking.abortCalls.get() >= starts.size() - 1); + assertEquals("probes must not abort — an aborted response is not a poolable connection", 0, tracking.abortCalls.get()); + assertFalse("no probe may reach the abort path", tracking.aborted.get()); assertThat( - "macro-split probes must not drain range streams; consumed " - + tracking.bytesConsumed.get() - + " of " - + fileLength - + " bytes across " - + tracking.abortCalls.get() - + " probes", + "macro-split probes must read bounded chunks; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", tracking.bytesConsumed.get(), Matchers.lessThan(fileLength / 2) ); From b06aae31cfde3880e1d69c357acf4fca57decb11 Mon Sep 17 00:00:00 2001 From: Oleg Lvovitch Date: Fri, 24 Jul 2026 19:38:28 +0100 Subject: [PATCH 3/4] ESQL: chunk the read-time segment probes too computeSegments carries a near-verbatim copy of the discovery probe loop, run per macro-split on the data node to sub-segment it. It had the same defect: a range declared to the end of the split, read for a few hundred bytes, then aborted, destroying the connection. It runs ahead of any parsing dispatch, so on a 9.8GB scan roughly 15 probes per split become dead time in front of every split, on the order of 2000 reconnects across the scan. Both probe loops now share ChunkedStorageInputStream. At read time the object is a RangeStorageObject, whose readBytes already offset translates and clamps to the split view, so the chunked reads need no coordinate handling. The ticket query does not reach this path: a pushed row limit makes openWithParallelism bail before segmenting. This is for scans and aggregates. Segment boundaries are unchanged; the datasources suite passes with only the two abort-pinning tests inverted. --- .../ParallelParsingCoordinator.java | 17 +++++++---------- .../ParallelParsingCoordinatorTests.java | 16 +++++++++------- .../StorageObjectAbortChainTests.java | 18 +++++++----------- 3 files changed, 23 insertions(+), 28 deletions(-) diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinator.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinator.java index dce90d79e73f8..b6293c7955f31 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinator.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinator.java @@ -26,7 +26,6 @@ import org.elasticsearch.xpack.esql.datasources.spi.StorageObject; import org.elasticsearch.xpack.esql.datasources.spi.StripeColumnScope; -import java.io.Closeable; import java.io.IOException; import java.io.InputStream; import java.util.ArrayList; @@ -604,10 +603,11 @@ public static List computeSegments( } long boundary; if (strided) { - InputStream stream = storageObject.newStream(pos, remaining); - // Abort rather than close: findNextRecordBoundary reads only a prefix of the range - // (fileLength - pos bytes), but close() on providers like S3 drains the remainder. - try (Closeable abortOnExit = () -> storageObject.abortStream(stream)) { + // Probe through bounded chunks: each is a closed range read to completion, so closing it + // releases a connection the provider can pool. Declaring the whole remainder instead (the + // previous shape) left abort as the only exit, and an abort destroys the connection, so every + // probe re-paid a connect and TLS handshake. + try (InputStream stream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { long skipped = splitter.findNextRecordBoundary(stream); if (skipped < 0) { break; @@ -616,8 +616,7 @@ public static List computeSegments( } } else { long probed; - InputStream probeStream = storageObject.newStream(pos, remaining); - try (Closeable abortOnExit = () -> storageObject.abortStream(probeStream)) { + try (InputStream probeStream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { probed = splitter.findProvenRecordBoundary(probeStream); } if (probed >= 0) { @@ -625,10 +624,8 @@ public static List computeSegments( } else if (probed == RecordSplitter.AMBIGUOUS) { // Exact walk from the last proven boundary. This path is range-bounded and read at planning // time, so it relies on the existing read-time cancellation rather than a new supplier. - long walkRemaining = fileLength - exactCursor; - InputStream walkStream = storageObject.newStream(exactCursor, walkRemaining); long start; - try (Closeable abortOnExit = () -> storageObject.abortStream(walkStream)) { + try (InputStream walkStream = new ChunkedStorageInputStream(storageObject, exactCursor, fileLength)) { start = splitter.findRecordStartAtOrAfter(walkStream, pos - exactCursor, () -> false); } if (start == RecordSplitter.RECORD_TOO_LARGE || start < 0) { diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinatorTests.java index 5fd6678ca0c85..2e4762ba46125 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinatorTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/ParallelParsingCoordinatorTests.java @@ -130,11 +130,12 @@ public void testComputeSegmentsAlignsToBoundaries() throws IOException { } /** - * Regression guard: {@link ParallelParsingCoordinator#computeSegments} opens a range stream for - * each nominal split probe, reads only enough bytes to find the next record boundary, then must - * call {@link StorageObject#abortStream} — not a draining {@code close()}. + * Regression guard on the shape of the I/O, not just its result: each nominal split probe must read + * bounded chunks that it consumes in full, so the provider can pool the connection. Draining a huge + * remainder wastes the transfer, and aborting to avoid the drain destroys the connection — this loop + * runs per macro-split at read time, so a per-probe reconnect is paid on every scan. */ - public void testComputeSegmentsDoesNotDrainStream() throws IOException { + public void testComputeSegmentsUsesPooledChunkReads() throws IOException { BlockFactory blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build(); StringBuilder csv = new StringBuilder("id,name\n"); @@ -159,11 +160,12 @@ public void testComputeSegmentsDoesNotDrainStream() throws IOException { ); assertThat("expected multiple parse segments", segments.size(), Matchers.greaterThan(1)); - assertTrue("each segment probe must abort the underlying stream", tracking.abortCalls.get() >= segments.size() - 1); + assertEquals("segment probes must not abort — an aborted response is not a poolable connection", 0, tracking.abortCalls.get()); + assertFalse("no probe may reach the abort path", tracking.aborted.get()); assertThat( - "segment probes must not drain the range streams; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", + "segment probes must read bounded chunks; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", tracking.bytesConsumed.get(), - Matchers.lessThan(fileLength / 2) + Matchers.lessThanOrEqualTo((long) segments.size() * ChunkedStorageInputStream.FIRST_CHUNK_SIZE) ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java index 8dc227d6dc420..11e720b2573fc 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/StorageObjectAbortChainTests.java @@ -133,10 +133,11 @@ public void testMacroSplitDiscoveryProbesFullyConsumeClosedRangesThroughDecorato } /** - * Regression guard for {@link ParallelParsingCoordinator#computeSegments}: in-file parallel parsing - * probes record boundaries through the same decorator chain used for uncompressed object-store reads. + * Regression guard for {@link ParallelParsingCoordinator#computeSegments}: in-file parallel parsing probes + * record boundaries in bounded chunks it consumes in full, through the same decorator chain used for + * uncompressed object-store reads — so cleanup neither drains a large remainder nor aborts. */ - public void testComputeSegmentsAbortPropagatesThroughDecoratorChainWithoutDrain() throws IOException { + public void testComputeSegmentsProbesFullyConsumeClosedRangesThroughDecoratorChain() throws IOException { StringBuilder csv = new StringBuilder("id,name\n"); for (int i = 0; i < 200_000; i++) { csv.append(i).append(",value_").append(i).append('\n'); @@ -158,15 +159,10 @@ public void testComputeSegmentsAbortPropagatesThroughDecoratorChainWithoutDrain( List segments = ParallelParsingCoordinator.computeSegments(csvReader, chain, fileLength, 4, csvReader.minimumSegmentSize()); assertThat("expected multiple parse segments", segments.size(), Matchers.greaterThan(1)); - assertTrue("each probe must abort the raw stream", tracking.abortCalls.get() >= segments.size() - 1); + assertEquals("probes must not abort — an aborted response is not a poolable connection", 0, tracking.abortCalls.get()); + assertFalse("no probe may reach the abort path", tracking.aborted.get()); assertThat( - "segment probes must not drain range streams; consumed " - + tracking.bytesConsumed.get() - + " of " - + fileLength - + " bytes across " - + tracking.abortCalls.get() - + " probes", + "segment probes must read bounded chunks; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes", tracking.bytesConsumed.get(), Matchers.lessThan(fileLength / 2) ); From ab0b2f3af3e98250b4d7b8c21fbceda210104a54 Mon Sep 17 00:00:00 2001 From: Oleg Lvovitch Date: Fri, 24 Jul 2026 19:46:32 +0100 Subject: [PATCH 4/4] ESQL: probe the stride grid in parallel on a single file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Discovery fans out across files, so a dataset of one file gets no parallelism at all — and one large file is exactly where boundary probing dominates planning. ~150 probes over a 9.8GB ndjson file ran back to back on one thread, each a round trip to a distant store. A strided splitter resolves a boundary at any offset without knowing the previous one, so its probes are independent. The strided path now fixes its grid up front at stride multiples and, when discovery is not already using the pool for files, probes the grid through BoundedParallelGather. The two branches are complements by construction, so a probe gather is never nested inside a file gather. This replaces a chain that advanced by lastBoundary + stride. Both emit true record starts and tile the file; the grid does not accumulate drift, so a split can be shorter than a stride by up to the record straddling its grid point. A record spanning several grid points answers each with the same boundary, so repeats are dropped, and a probe that finds none truncates the grid, leaving the remainder as one trailing split, as the chain did. Probes return the splitter sentinel rather than throwing: a throw would fast-fail the gather and discard boundaries siblings had already found. Cancellation still throws, which is the point. Boundaries are identical with and without an executor, and that is asserted rather than assumed. --- .../esql/datasources/FileSplitProvider.java | 184 +++++++++++++++--- .../datasources/FileSplitProviderTests.java | 153 +++++++++++++++ 2 files changed, 312 insertions(+), 25 deletions(-) diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java index 0fa8814d24a8c..c2fe8a80c35de 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/datasources/FileSplitProvider.java @@ -381,14 +381,19 @@ public SplitDiscoveryResult discoverSplits(SplitDiscoveryContext context) { if (executor != null && tasks.size() > 1) { perFileSplits = BoundedParallelGather.gather( tasks, - task -> processFileForSplits(task, hoistedProvider, isCancelled), + task -> processFileForSplits(task, hoistedProvider, isCancelled, null), MAX_PARALLEL_SPLIT_DISCOVERY, executor ); } else { + // Files are the usual unit of parallelism, but a dataset of one file gets none from that — and a + // single large file is exactly where boundary probing dominates planning. Lend the pool to the + // probes instead. The two branches are complements by construction, so probes never nest a + // second blocking gather inside a task already waiting on the first. + Executor probeExecutor = tasks.size() == 1 ? executor : null; perFileSplits = new ArrayList<>(tasks.size()); for (FileTask task : tasks) { - perFileSplits.add(processFileForSplits(task, hoistedProvider, isCancelled)); + perFileSplits.add(processFileForSplits(task, hoistedProvider, isCancelled, probeExecutor)); } } } catch (IOException e) { @@ -456,18 +461,29 @@ private static Map attributesToTypeMap(List attribu * otherwise falls back to the registry for per-call provider resolution. * This method is safe to call concurrently from multiple threads. */ - private List processFileForSplits(FileTask task, @Nullable StorageProvider hoistedProvider, BooleanSupplier isCancelled) - throws IOException { + private List processFileForSplits( + FileTask task, + @Nullable StorageProvider hoistedProvider, + BooleanSupplier isCancelled, + @Nullable Executor probeExecutor + ) throws IOException { if (isCancelled.getAsBoolean()) { throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE); } // Carry the cancellation signal as ambient thread-local state so the synchronous retry/throttle // backoff inside the footer reads below can abort a parked sleep on cancel. - return StorageRetryCancellation.callWithCancellation(isCancelled, () -> computeFileSplits(task, hoistedProvider, isCancelled)); + return StorageRetryCancellation.callWithCancellation( + isCancelled, + () -> computeFileSplits(task, hoistedProvider, isCancelled, probeExecutor) + ); } - private List computeFileSplits(FileTask task, @Nullable StorageProvider hoistedProvider, BooleanSupplier isCancelled) - throws IOException { + private List computeFileSplits( + FileTask task, + @Nullable StorageProvider hoistedProvider, + BooleanSupplier isCancelled, + @Nullable Executor probeExecutor + ) throws IOException { List fileSplits = new ArrayList<>(); // Resolve the config-aware reader once and reuse it for both the sequential-whole-file gate and the @@ -558,7 +574,8 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr fileSplits, hoistedProvider, configuredReader, - isCancelled + isCancelled, + probeExecutor )) { return fileSplits; } @@ -927,7 +944,8 @@ private boolean tryNewlineAlignedMacroSplits( List splits, @Nullable StorageProvider hoistedProvider, @Nullable FormatReader reader, - BooleanSupplier isCancelled + BooleanSupplier isCancelled, + @Nullable Executor probeExecutor ) throws IOException { if (formatRegistry == null || storageRegistry == null || targetStrideBytes <= 0 || fileLength <= targetStrideBytes) { return false; @@ -954,7 +972,8 @@ private boolean tryNewlineAlignedMacroSplits( fileLength, targetStrideBytes, maxRecordBytes, - isCancelled + isCancelled, + probeExecutor ); if (starts.size() <= 1) { return false; @@ -1037,8 +1056,8 @@ private static boolean legacyUnstampedWholeFile(FileSplit split) { } /** - * Absolute byte offsets where each macro-split starts (always begins with {@code 0}), mirroring - * {@link ParallelParsingCoordinator#computeSegments} stride semantics with {@code targetStrideBytes}. + * Probes sequentially. See the overload taking a probe {@link Executor} for the boundary semantics, which + * are identical either way. */ static List computeRecordAlignedMacroSplitStarts( SegmentableFormatReader reader, @@ -1047,6 +1066,45 @@ static List computeRecordAlignedMacroSplitStarts( long targetStrideBytes, int maxRecordBytes, BooleanSupplier isCancelled + ) throws IOException { + return computeRecordAlignedMacroSplitStarts( + reader, + storageObject, + fileLength, + targetStrideBytes, + maxRecordBytes, + isCancelled, + null + ); + } + + /** + * Absolute byte offsets where each macro-split starts (always begins with {@code 0}), one per + * {@code targetStrideBytes} stride of the file. + *

+ * A strided splitter can probe any offset directly, so its stride grid is fixed up front at multiples of + * {@code targetStrideBytes} and every probe is independent — which is what lets {@code probeExecutor}, when + * supplied, run them concurrently. Each probe still costs a network round trip, so on a single large file + * over a distant object store the serial version is the dominant cost of planning. Boundaries do not depend + * on whether an executor was supplied. + *

+ * The grid replaces an earlier chain that advanced by {@code lastBoundary + stride}. Both emit true record + * starts and both tile the file; the grid simply does not accumulate drift, so a split can come out shorter + * than a stride by up to the length of the record straddling its grid point, where the chain was always at + * least a stride. A record long enough to span several grid points yields the same boundary from each, so + * repeats are dropped rather than emitted as empty splits. + *

+ * A non-strided (but proven-probing) splitter is genuinely sequential — each probe's fallback walks from the + * previous proven record start — so it keeps the chained loop and ignores {@code probeExecutor}. + */ + static List computeRecordAlignedMacroSplitStarts( + SegmentableFormatReader reader, + StorageObject storageObject, + long fileLength, + long targetStrideBytes, + int maxRecordBytes, + BooleanSupplier isCancelled, + @Nullable Executor probeExecutor ) throws IOException { List boundaries = new ArrayList<>(); boundaries.add(0L); @@ -1066,6 +1124,9 @@ static List computeRecordAlignedMacroSplitStarts( + "] supports neither strided nor proven probing and cannot be macro-split" ); } + if (strided) { + return stridedMacroSplitStarts(splitter, storageObject, fileLength, targetStrideBytes, minSegment, isCancelled, probeExecutor); + } // The last proven record start, i.e. the base offset the exact walk streams from when the probe is // AMBIGUOUS. The file start is always a record start, so it seeds at 0. long exactCursor = 0L; @@ -1076,19 +1137,7 @@ static List computeRecordAlignedMacroSplitStarts( break; } long boundary; - if (strided) { - // Probe through bounded chunks: each is a closed range read to completion, so closing it - // releases a connection the provider can pool. Declaring the whole remainder instead (the - // previous shape) left abort as the only exit, and an abort destroys the connection, so every - // probe re-paid a connect and TLS handshake. - try (InputStream stream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { - long skipped = splitter.findNextRecordBoundary(stream); - if (skipped == RecordSplitter.RECORD_TOO_LARGE || skipped < 0) { - break; - } - boundary = pos + skipped; - } - } else { + { long probed; try (InputStream probeStream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { probed = splitter.findProvenRecordBoundary(probeStream); @@ -1130,6 +1179,91 @@ static List computeRecordAlignedMacroSplitStarts( return boundaries; } + /** + * Macro-split starts for a strided splitter: probe a fixed grid of stride multiples, concurrently when an + * executor is available. Probes are independent because a strided splitter can resolve a record boundary at + * any offset without knowing the previous one. + */ + private static List stridedMacroSplitStarts( + RecordSplitter splitter, + StorageObject storageObject, + long fileLength, + long targetStrideBytes, + long minSegment, + BooleanSupplier isCancelled, + @Nullable Executor probeExecutor + ) throws IOException { + List gridPositions = new ArrayList<>(); + for (long pos = targetStrideBytes; pos < fileLength && fileLength - pos >= minSegment; pos += targetStrideBytes) { + gridPositions.add(pos); + } + + List probed; + if (probeExecutor == null || gridPositions.size() <= 1) { + probed = new ArrayList<>(gridPositions.size()); + for (long pos : gridPositions) { + probed.add(probeStridedBoundary(splitter, storageObject, fileLength, pos, isCancelled)); + } + } else { + try { + probed = BoundedParallelGather.gather( + gridPositions, + pos -> probeStridedBoundary(splitter, storageObject, fileLength, pos, isCancelled), + MAX_PARALLEL_SPLIT_DISCOVERY, + probeExecutor + ); + } catch (IOException | RuntimeException e) { + throw e; + } catch (Exception e) { + throw new IOException("Failed to probe record-aligned macro-split boundaries", e); + } + } + + List boundaries = new ArrayList<>(); + boundaries.add(0L); + for (long boundary : probed) { + // A probe that found no boundary (end of data, or a record over the byte budget) stops the grid + // rather than skipping a point: the chained walk this replaced could never see past such a probe, + // and stopping leaves the rest of the file as one trailing split, which is correct if less parallel. + if (boundary < 0 || boundary >= fileLength || fileLength - boundary < minSegment) { + break; + } + // A record spanning several grid points answers each of them with its own end, so the same boundary + // arrives more than once. Emitting it twice would mean a zero-length split. + if (boundary <= boundaries.get(boundaries.size() - 1)) { + continue; + } + boundaries.add(boundary); + } + return boundaries; + } + + /** + * Resolves the record boundary at or after one grid point, returning the splitter's negative sentinel + * unchanged when there is none. Probes must not throw for a missing boundary: a throw would fast-fail the + * whole gather and discard the boundaries its siblings already found. Cancellation is the deliberate + * exception — there, abandoning the remaining probes is the point. + */ + private static long probeStridedBoundary( + RecordSplitter splitter, + StorageObject storageObject, + long fileLength, + long pos, + BooleanSupplier isCancelled + ) throws IOException { + if (isCancelled.getAsBoolean()) { + throw new TaskCancelledException("split discovery cancelled"); + } + // Probe through bounded chunks: each is a closed range read to completion, so closing it releases a + // connection the provider can pool. Declaring the whole remainder instead (the previous shape) left + // abort as the only exit, and an abort destroys the connection, so every probe re-paid a connect and + // TLS handshake. + try (InputStream stream = new ChunkedStorageInputStream(storageObject, pos, fileLength)) { + long skipped = splitter.findNextRecordBoundary(stream); + return skipped < 0 ? skipped : pos + skipped; + } + } + private boolean tryIndexedSplits( IndexedDecompressionCodec indexedCodec, StoragePath filePath, diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java index 9eb464e15ce11..2593f08ffd410 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java @@ -73,6 +73,9 @@ import java.util.Map; import java.util.Set; import java.util.TreeSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BooleanSupplier; @@ -2523,4 +2526,154 @@ private static Literal intLiteral(int value) { private static Attribute refAttr(String name) { return new ReferenceAttribute(SRC, name, DataType.KEYWORD); } + + /** + * A payload comfortably larger than the 1 MiB {@link SegmentableFormatReader#minimumSegmentSize}, since a + * grid point is only kept when at least that much file remains after it — a smaller corpus yields no grid + * at all and would make these tests assert nothing. + */ + private static byte[] plainCsvPayload(int targetBytes) { + StringBuilder csv = new StringBuilder("id,name\n"); + for (int i = 0; csv.length() < targetBytes; i++) { + csv.append(i).append(",value_").append(i).append('\n'); + } + return csv.toString().getBytes(StandardCharsets.UTF_8); + } + + private SegmentableFormatReader plainCsvReader() { + var blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build(); + return (SegmentableFormatReader) new CsvFormatReader(blockFactory).withConfig(Map.of("mode", "plain")); + } + + /** + * Probing the stride grid concurrently must not change where the splits land. Determinism is what makes the + * fan-out reviewable: it turns "does parallelism corrupt the boundaries" into a question the sequential run + * answers, rather than one that depends on interleaving. + */ + public void testStridedGridBoundariesMatchWithAndWithoutProbeExecutor() throws Exception { + byte[] payload = plainCsvPayload(10 * 1024 * 1024); + SegmentableFormatReader reader = plainCsvReader(); + StorageObject object = new RecordingStorageObject(payload); + long stride = payload.length / 8; + + List sequential = FileSplitProvider.computeRecordAlignedMacroSplitStarts( + reader, + object, + payload.length, + stride, + SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES, + () -> false, + null + ); + + ExecutorService executor = Executors.newFixedThreadPool(8); + List concurrent; + try { + concurrent = FileSplitProvider.computeRecordAlignedMacroSplitStarts( + reader, + object, + payload.length, + stride, + SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES, + () -> false, + executor + ); + } finally { + executor.shutdownNow(); + } + + assertThat("expected several macro-split boundaries", sequential.size(), greaterThan(2)); + assertEquals("concurrent probing must produce identical boundaries", sequential, concurrent); + + Set trueStarts = trueRecordStarts(reader.recordSplitter(SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES), payload); + long previous = -1; + for (long start : concurrent) { + assertThat("boundaries must strictly increase", start, greaterThan(previous)); + assertTrue("boundary " + start + " must be a true record start", trueStarts.contains(start)); + previous = start; + } + } + + /** + * A record longer than the stride answers several grid points with the same boundary. Emitting it once per + * grid point would produce zero-length splits, so the repeats must collapse. + */ + public void testStridedGridDropsRepeatsFromRecordSpanningGridPoints() throws Exception { + StringBuilder csv = new StringBuilder("id,name\n"); + csv.append("1,").append("x".repeat(2_000_000)).append('\n'); + for (int i = 0; csv.length() < 12 * 1024 * 1024; i++) { + csv.append(i).append(",value_").append(i).append('\n'); + } + byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8); + SegmentableFormatReader reader = plainCsvReader(); + + // Small enough that the single huge record straddles many consecutive grid points. + long stride = 1024 * 1024; + List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts( + reader, + new RecordingStorageObject(payload), + payload.length, + stride, + SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES, + () -> false, + null + ); + + assertEquals("boundaries must be unique", new TreeSet<>(starts).size(), starts.size()); + long previous = -1; + for (long start : starts) { + assertThat("boundaries must strictly increase", start, greaterThan(previous)); + previous = start; + } + } + + /** + * A grid point whose record exceeds the byte budget stops the grid, leaving the rest of the file as one + * trailing split — the shape the chained walk produced, since it could never see past such a probe. + */ + public void testStridedGridStopsAtFirstGridPointOverRecordBudget() throws Exception { + StringBuilder csv = new StringBuilder("id,name\n"); + for (int i = 0; csv.length() < 4 * 1024 * 1024; i++) { + csv.append(i).append(",value_").append(i).append('\n'); + } + int oversizedAt = csv.length(); + csv.append("big,").append("x".repeat(3 * 1024 * 1024)).append('\n'); + for (int i = 0; csv.length() < 12 * 1024 * 1024; i++) { + csv.append(i).append(",tail_").append(i).append('\n'); + } + byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8); + + List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts( + plainCsvReader(), + new RecordingStorageObject(payload), + payload.length, + 1024 * 1024, + 64 * 1024, + () -> false, + null + ); + + assertThat("boundaries before the oversized record are still emitted", starts.size(), greaterThan(1)); + for (long start : starts) { + assertThat("the grid must not continue past the oversized record", start, lessThanOrEqualTo((long) oversizedAt)); + } + } + + public void testStridedGridProbeCancellationIsObserved() { + byte[] payload = plainCsvPayload(10 * 1024 * 1024); + AtomicBoolean cancelled = new AtomicBoolean(); + + expectThrows( + TaskCancelledException.class, + () -> FileSplitProvider.computeRecordAlignedMacroSplitStarts( + plainCsvReader(), + new RecordingStorageObject(payload), + payload.length, + payload.length / 8, + SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES, + () -> cancelled.getAndSet(true), + null + ) + ); + } }