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/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..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
@@ -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;
@@ -382,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) {
@@ -457,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
@@ -559,7 +574,8 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
fileSplits,
hoistedProvider,
configuredReader,
- isCancelled
+ isCancelled,
+ probeExecutor
)) {
return fileSplits;
}
@@ -928,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;
@@ -955,7 +972,8 @@ private boolean tryNewlineAlignedMacroSplits(
fileLength,
targetStrideBytes,
maxRecordBytes,
- isCancelled
+ isCancelled,
+ probeExecutor
);
if (starts.size() <= 1) {
return false;
@@ -1038,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,
@@ -1048,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);
@@ -1067,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;
@@ -1077,21 +1137,9 @@ static List computeRecordAlignedMacroSplitStarts(
break;
}
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)) {
- long skipped = splitter.findNextRecordBoundary(stream);
- if (skipped == RecordSplitter.RECORD_TOO_LARGE || skipped < 0) {
- break;
- }
- boundary = pos + skipped;
- }
- } 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 +1147,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) {
@@ -1133,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/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/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/FileSplitProviderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/datasources/FileSplitProviderTests.java
index 1ecbd10eaa9e0..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,12 +73,15 @@
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;
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 +1022,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 +1055,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)
);
}
@@ -2520,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
+ )
+ );
+ }
}
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/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;
+ }
+}
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..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
@@ -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,25 +123,21 @@ 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)
);
}
/**
- * 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');
@@ -163,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)
);