conf
*/
static final String COMPRESSED_OFFSET_SPLIT_KEY = "_compressed_offset_split";
- /** Maximum parallel per-file I/O tasks during split discovery (Parquet footer reads, etc.). */
+ /**
+ * Ceiling on concurrent I/O tasks during split discovery, applied separately to the per-file planning pass
+ * (Parquet footer reads, etc.) and to the record-boundary probes that follow it. The two passes run one after
+ * the other, so this bounds in-flight reads at any instant rather than being multiplied between them.
+ */
static final int MAX_PARALLEL_SPLIT_DISCOVERY = 16;
+ /**
+ * Bytes read from each stride offset when probing for the next record boundary during newline-aligned
+ * macro-split discovery. A row-oriented record (an NDJSON line, a CSV row) is far smaller than this, so its
+ * terminating newline is found well within the window. An offset whose record does span the whole window
+ * yields no split start rather than reading further, which costs that one split: the offsets after it are
+ * probed as usual and the splits either side of it merge into one. Bounding the probe to a fixed window keeps
+ * each boundary probe a small, predictable ranged GET instead of a range opened to end-of-file. A probe drains
+ * the window in full (see {@link #probeStridedBoundary}) so the connection is pooled for the next probe to
+ * reuse; the window is bounded so that drain stays cheap.
+ *
+ * Doubles as the floor on the strided probe stride (see {@link #newlineMacroSplitCandidate}): a stride below
+ * the window would put several probe offsets inside one window's worth of file, each drained in full, while
+ * materializing a probe task per stride offset — {@code fileLength / stride} of them, unbounded below. Flooring
+ * the stride to the window caps a file's probe count at {@code fileLength / window} and keeps every probe's
+ * window disjoint from its neighbours'.
+ */
+ static final long MACRO_SPLIT_PROBE_WINDOW_BYTES = 256 * 1024;
+
private static final String DISCOVERY_CANCELLED_MESSAGE = "ES|QL external split discovery cancelled";
private final long targetSplitSizeBytes;
@@ -374,22 +397,24 @@ public SplitDiscoveryResult discoverSplits(SplitDiscoveryContext context) {
return new SplitDiscoveryResult(List.of(), 0, exhaustivelyPruned);
}
- // Phase 2: I/O-bound split discovery — parallelize when executor is available.
+ // Phase 2: I/O-bound split planning, parallelized across files when an executor is available. Files
+ // whose record boundaries still need probing come back as deferred descriptors; everything else
+ // (Parquet footers, block-aligned compressed, range splits, whole-file) finishes here.
final StorageProvider hoistedProvider = sharedProvider;
final BooleanSupplier isCancelled = context.isCancelled();
- List> perFileSplits;
+ List planResults;
try {
if (executor != null && tasks.size() > 1) {
- perFileSplits = BoundedParallelGather.gather(
+ planResults = BoundedParallelGather.gather(
tasks,
task -> processFileForSplits(task, hoistedProvider, isCancelled),
MAX_PARALLEL_SPLIT_DISCOVERY,
executor
);
} else {
- perFileSplits = new ArrayList<>(tasks.size());
+ planResults = new ArrayList<>(tasks.size());
for (FileTask task : tasks) {
- perFileSplits.add(processFileForSplits(task, hoistedProvider, isCancelled));
+ planResults.add(processFileForSplits(task, hoistedProvider, isCancelled));
}
}
} catch (IOException e) {
@@ -400,16 +425,143 @@ public SplitDiscoveryResult discoverSplits(SplitDiscoveryContext context) {
throw new RuntimeException("Failed to discover splits", e);
}
- // Flatten per-file split lists into a single ordered list.
+ // Phase 3: probe the deferred files' record boundaries. Every deferred file's stride offsets go into one
+ // flat batch under a single concurrency budget, so the number of in-flight probe reads is bounded by that
+ // budget no matter how many files are being probed. Probing per file instead would multiply the per-file
+ // budget by the number of files in flight.
+ Map> boundariesByPlan = probeDeferredBoundaries(planResults, isCancelled);
+
+ // Splits are emitted in file order: walking the plan results keeps a deferred file's macro-splits in the
+ // same position its file occupied in the file list.
List splits = new ArrayList<>();
- for (List fileSplits : perFileSplits) {
- splits.addAll(fileSplits);
+ for (int i = 0; i < planResults.size(); i++) {
+ PlanResult planResult = planResults.get(i);
+ if (planResult.deferred() == null) {
+ splits.addAll(planResult.splits());
+ } else {
+ List starts = boundariesByPlan.get(i);
+ if (starts == null) {
+ // A file is only deferred when it has offsets to probe, so the probe phase answers for every
+ // deferred file. Fail loud rather than on a null if that ever stops holding.
+ throw new IllegalStateException("no probed boundaries for deferred file " + planResult.deferred().filePath());
+ }
+ splits.addAll(buildNewlineMacroSplits(planResult.deferred(), starts));
+ }
}
// Each surviving task produces at least one split, so the task count is the number of
// distinct files that are actually scanned after coordinator-side pruning.
return new SplitDiscoveryResult(splits, tasks.size());
}
+ /**
+ * Probes the record boundaries of every deferred file, keyed by the file's index in {@code planResults}.
+ *
+ * With an executor, all files' stride offsets are gathered into one flat batch so a single concurrency budget
+ * ({@link #probeConcurrency()}) bounds the in-flight probe reads across the whole query, rather than one budget
+ * per file multiplied by the files in flight. Without one, each file is walked serially. Both produce the same
+ * boundaries.
+ */
+ private Map> probeDeferredBoundaries(List planResults, BooleanSupplier isCancelled) {
+ List deferredIndices = new ArrayList<>();
+ int probeCount = 0;
+ for (int i = 0; i < planResults.size(); i++) {
+ DeferredNewlineSplits deferred = planResults.get(i).deferred();
+ if (deferred != null) {
+ deferredIndices.add(i);
+ probeCount += deferred.positions().size();
+ }
+ }
+ if (probeCount == 0) {
+ return Map.of();
+ }
+ // A cancel between planning and probing should be seen before any probe read is issued.
+ if (isCancelled.getAsBoolean()) {
+ throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE);
+ }
+
+ try {
+ Map> boundariesByPlan;
+ if (executor == null) {
+ boundariesByPlan = new HashMap<>(deferredIndices.size());
+ for (int planIndex : deferredIndices) {
+ DeferredNewlineSplits deferred = planResults.get(planIndex).deferred();
+ boundariesByPlan.put(
+ planIndex,
+ probeStridedBoundariesSerially(
+ deferred.splitter(),
+ deferred.storageObject(),
+ deferred.fileLength(),
+ deferred.positions(),
+ deferred.minSegment(),
+ isCancelled
+ )
+ );
+ }
+ } else {
+ List probeTasks = new ArrayList<>(probeCount);
+ for (int planIndex : deferredIndices) {
+ DeferredNewlineSplits deferred = planResults.get(planIndex).deferred();
+ for (long position : deferred.positions()) {
+ probeTasks.add(new ProbeTask(planIndex, deferred, position));
+ }
+ }
+ List outcomes = BoundedParallelGather.gather(
+ probeTasks,
+ probe -> runProbe(probe, isCancelled),
+ probeConcurrency(),
+ executor
+ );
+
+ // Results come back in input order, and each file's offsets were added in ascending order, so
+ // grouping by file preserves the ascending order the reduction requires.
+ Map> outcomesByPlan = new HashMap<>();
+ for (int i = 0; i < probeTasks.size(); i++) {
+ outcomesByPlan.computeIfAbsent(probeTasks.get(i).planIndex(), k -> new ArrayList<>()).add(outcomes.get(i));
+ }
+ boundariesByPlan = new HashMap<>(outcomesByPlan.size());
+ for (Map.Entry> entry : outcomesByPlan.entrySet()) {
+ boundariesByPlan.put(entry.getKey(), reduceProbeOutcomes(entry.getValue()));
+ }
+ }
+ // A cancel landing after the last probe has read is seen by no probe, so check once more here rather
+ // than returning a split set the caller will discard.
+ if (isCancelled.getAsBoolean()) {
+ throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE);
+ }
+ return boundariesByPlan;
+ } catch (RuntimeException e) {
+ throw e;
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to discover splits", e);
+ }
+ }
+
+ /** Probes the one stride offset a task carries, against the file the task belongs to. */
+ private static ProbeOutcome runProbe(ProbeTask probe, BooleanSupplier isCancelled) throws IOException {
+ DeferredNewlineSplits deferred = probe.deferred();
+ return probeStridedBoundary(
+ deferred.splitter(),
+ deferred.storageObject(),
+ probe.position(),
+ deferred.fileLength(),
+ deferred.minSegment(),
+ isCancelled
+ );
+ }
+
+ /**
+ * How many boundary probes may be in flight at once across the whole query.
+ *
+ * Bounded by {@link #MAX_PARALLEL_SPLIT_DISCOVERY}, but clamped to the node's blob-store concurrency: every
+ * probe holds one of those permits for as long as its stream is open, so submitting more probes than there
+ * are permits would only park discovery threads waiting to acquire one. A configured concurrency of {@code 0}
+ * disables permit limiting altogether rather than meaning "no concurrency", so the ceiling applies as-is.
+ */
+ int probeConcurrency() {
+ int permits = ExternalSourceSettings.blobStoreConcurrency(settings);
+ return permits > 0 ? Math.min(MAX_PARALLEL_SPLIT_DISCOVERY, permits) : MAX_PARALLEL_SPLIT_DISCOVERY;
+ }
+
/**
* Throws {@link TaskCancelledException} when the originating query has been cancelled, so that a
* long-running split discovery (e.g. thousands of Parquet footer reads) aborts promptly. Mirrors
@@ -444,6 +596,50 @@ private record FileTask(
@Nullable Map inferredFileTypes
) {}
+ /**
+ * A file that will be macro-split at record boundaries, carrying everything needed to probe those boundaries
+ * and then build its splits. Planning produces this without reading any file bytes so that the probes of
+ * every such file can be pooled into one bounded, concurrent batch instead of running per file.
+ *
+ * {@code positions} holds the fixed stride offsets to probe, and is empty for a splitter that cannot be
+ * probed at a fixed offset (quoted/escaped CSV/TSV) and must use the sequential proven walk instead.
+ *
+ * {@code splitter} is shared by every probe of this file, which is safe because record splitters hold no
+ * per-scan state (they are configured with a max record size and nothing else).
+ */
+ private record DeferredNewlineSplits(
+ StoragePath filePath,
+ long fileLength,
+ @Nullable String format,
+ Map config,
+ Map partitionValues,
+ @Nullable ColumnMapping columnMapping,
+ @Nullable List readSchema,
+ StorageObject storageObject,
+ RecordSplitter splitter,
+ long minSegment,
+ long targetStrideBytes,
+ int maxRecordBytes,
+ List positions
+ ) {}
+
+ /**
+ * The outcome of planning one file: either its final splits, or a descriptor whose record boundaries still
+ * need probing. Deferring the probing lets every candidate file's probes share a single concurrency budget.
+ */
+ private record PlanResult(List splits, @Nullable DeferredNewlineSplits deferred) {
+ static PlanResult of(List splits) {
+ return new PlanResult(splits, null);
+ }
+
+ static PlanResult needsProbing(DeferredNewlineSplits deferred) {
+ return new PlanResult(List.of(), deferred);
+ }
+ }
+
+ /** One stride offset to probe, tied back to the file (by plan index) whose boundaries it contributes to. */
+ private record ProbeTask(int planIndex, DeferredNewlineSplits deferred, long position) {}
+
private static Map attributesToTypeMap(List attributes) {
Map types = new HashMap<>(attributes.size());
for (Attribute a : attributes) {
@@ -457,7 +653,7 @@ 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)
+ private PlanResult processFileForSplits(FileTask task, @Nullable StorageProvider hoistedProvider, BooleanSupplier isCancelled)
throws IOException {
if (isCancelled.getAsBoolean()) {
throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE);
@@ -467,7 +663,7 @@ private List processFileForSplits(FileTask task, @Nullable Storag
return StorageRetryCancellation.callWithCancellation(isCancelled, () -> computeFileSplits(task, hoistedProvider, isCancelled));
}
- private List computeFileSplits(FileTask task, @Nullable StorageProvider hoistedProvider, BooleanSupplier isCancelled)
+ private PlanResult computeFileSplits(FileTask task, @Nullable StorageProvider hoistedProvider, BooleanSupplier isCancelled)
throws IOException {
List fileSplits = new ArrayList<>();
@@ -500,7 +696,7 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
task.readSchema()
)
);
- return fileSplits;
+ return PlanResult.of(fileSplits);
}
// Try block-aligned splitting for splittable compressed files (e.g. .ndjson.bz2).
@@ -517,7 +713,7 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
fileSplits,
hoistedProvider
)) {
- return fileSplits;
+ return PlanResult.of(fileSplits);
}
if (tryRangeAwareSplits(
@@ -534,7 +730,7 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
fileSplits,
hoistedProvider
)) {
- return fileSplits;
+ return PlanResult.of(fileSplits);
}
Map config = task.config();
@@ -546,7 +742,7 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
List readSchema = task.readSchema();
long effectiveTargetSplitBytes = resolveTargetSplitSize(config);
- if (tryNewlineAlignedMacroSplits(
+ DeferredNewlineSplits candidate = newlineMacroSplitCandidate(
filePath,
fileLength,
format,
@@ -556,12 +752,20 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
readSchema,
effectiveTargetSplitBytes,
task.maxRecordBytes(),
- fileSplits,
hoistedProvider,
- configuredReader,
- isCancelled
- )) {
- return fileSplits;
+ configuredReader
+ );
+ if (candidate != null) {
+ // A strided candidate's probes are deferred so they can share one concurrency budget with every
+ // other file's probes. A provable-but-not-strided one must walk sequentially, so resolve it here.
+ if (candidate.positions().isEmpty() == false) {
+ return PlanResult.needsProbing(candidate);
+ }
+ if (candidate.splitter().supportsStridedProbing() == false) {
+ return PlanResult.of(buildNewlineMacroSplits(candidate, resolveProvenMacroSplitStarts(candidate, isCancelled)));
+ }
+ // Strided, but no stride offset is far enough from end-of-file to be worth splitting at.
+ return PlanResult.of(buildNewlineMacroSplits(candidate, List.of(0L)));
}
// Whole-file split when macro splitting does not apply (small files, unsupported formats, or single aligned span).
@@ -578,7 +782,34 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
readSchema
)
);
- return fileSplits;
+ return PlanResult.of(fileSplits);
+ }
+
+ /**
+ * Resolves a non-strided candidate's macro-split starts with the sequential proven walk.
+ *
+ * A splitter that can be probed neither at a fixed offset nor by proving a record start must have been routed
+ * to a whole-file split upstream; if one arrives here that gate failed, so fail loud rather than emit
+ * mis-aligned macro-splits that silently mis-count rows.
+ */
+ private static List resolveProvenMacroSplitStarts(DeferredNewlineSplits deferred, BooleanSupplier isCancelled)
+ throws IOException {
+ RecordSplitter splitter = deferred.splitter();
+ if (splitter.supportsProvenProbing() == false) {
+ throw new IllegalStateException(
+ "record splitter ["
+ + splitter.getClass().getName()
+ + "] supports neither strided nor proven probing and cannot be macro-split"
+ );
+ }
+ return computeProvenMacroSplitStarts(
+ splitter,
+ deferred.storageObject(),
+ deferred.fileLength(),
+ deferred.targetStrideBytes(),
+ deferred.minSegment(),
+ isCancelled
+ );
}
/**
@@ -586,8 +817,8 @@ private List computeFileSplits(FileTask task, @Nullable StoragePr
* (no {@code formatRegistry}, no object name, or an unknown extension). Config-aware so a {@code WITH}
* override (e.g. {@code mode=plain}, {@code quote=none}) selects the same reader/splitter the read path
* will actually use: {@code byExtension} alone yields the extension default (quoted for {@code .csv}),
- * whose non-strided splitter would trip {@link #computeRecordAlignedMacroSplitStarts}' guard for a
- * plain-mode file. {@code withConfig} returns {@code null} only for test mocks; the base reader is used
+ * whose non-strided splitter would send a plain-mode file down the sequential proven walk instead of the
+ * strided one. {@code withConfig} returns {@code null} only for test mocks; the base reader is used
* in that case. The compression suffix is stripped by {@link FormatNameResolver}, so this resolves the
* inner text reader for compressed files (e.g. {@code .csv.bz2}) too.
*/
@@ -912,10 +1143,16 @@ private boolean tryRangeAwareSplits(
}
/**
- * Macro-splits supported line-oriented formats at record boundaries near {@code targetStrideBytes},
- * enabling multiple workers per file without mid-record cuts.
+ * Decides whether a file can be macro-split at record boundaries near {@code targetStrideBytes} and, if so,
+ * returns everything needed to probe those boundaries and build the splits. Performs no I/O: for a
+ * strided splitter the probe positions are pure arithmetic, so the caller is free to run the probes later and
+ * concurrently with other files' probes. Returns {@code null} when the file is not a macro-split candidate.
+ *
+ * A non-strided but provable splitter (quoted/escaped CSV/TSV) cannot defer: its walk is sequential, so the
+ * descriptor carries no positions and the caller must resolve it with the serial proven walk.
*/
- private boolean tryNewlineAlignedMacroSplits(
+ @Nullable
+ private DeferredNewlineSplits newlineMacroSplitCandidate(
StoragePath filePath,
long fileLength,
@Nullable String format,
@@ -925,41 +1162,78 @@ private boolean tryNewlineAlignedMacroSplits(
@Nullable List readSchema,
long targetStrideBytes,
int maxRecordBytes,
- List splits,
@Nullable StorageProvider hoistedProvider,
- @Nullable FormatReader reader,
- BooleanSupplier isCancelled
+ @Nullable FormatReader reader
) throws IOException {
if (formatRegistry == null || storageRegistry == null || targetStrideBytes <= 0 || fileLength <= targetStrideBytes) {
- return false;
+ return null;
}
if (isNewlineMacroSplitCandidateExtension(format) == false) {
- return false;
+ return null;
}
// Reuses the reader resolved once in processFileForSplits (config-aware; see resolveConfiguredReader).
if (reader == null) {
- return false;
+ return null;
}
if (reader instanceof CompressionDelegatingFormatReader) {
- return false;
+ return null;
}
if (reader instanceof SegmentableFormatReader == false) {
- return false;
+ return null;
}
SegmentableFormatReader segmentableReader = (SegmentableFormatReader) reader;
StorageProvider provider = resolveProvider(filePath, config, hoistedProvider);
StorageObject object = provider.newObject(filePath, fileLength);
- List starts = computeRecordAlignedMacroSplitStarts(
- segmentableReader,
- object,
+ RecordSplitter splitter = segmentableReader.recordSplitter(maxRecordBytes);
+ long minSegment = segmentableReader.minimumSegmentSize();
+ // A strided splitter probes fixed offsets, so its positions are known here without reading anything.
+ // Anything else keeps the sequential walk and carries no positions. The stride is floored to the probe
+ // window so a tiny target_split_size cannot materialize a probe per stride offset (unbounded below);
+ // see MACRO_SPLIT_PROBE_WINDOW_BYTES.
+ List positions = splitter.supportsStridedProbing()
+ ? macroSplitProbePositions(fileLength, Math.max(targetStrideBytes, MACRO_SPLIT_PROBE_WINDOW_BYTES), minSegment)
+ : List.of();
+ return new DeferredNewlineSplits(
+ filePath,
fileLength,
+ format,
+ config,
+ partitionValues,
+ columnMapping,
+ readSchema,
+ object,
+ splitter,
+ minSegment,
targetStrideBytes,
maxRecordBytes,
- isCancelled
+ positions
);
+ }
+
+ /**
+ * Builds a candidate file's splits from its resolved macro-split starts: one contiguous split per boundary,
+ * the last extending to end-of-file, each stamped so the read side can tell where it sits in the file.
+ * Falls back to a single whole-file split when probing found no usable boundary.
+ */
+ private static List buildNewlineMacroSplits(DeferredNewlineSplits deferred, List starts) {
+ long fileLength = deferred.fileLength();
+ Map config = deferred.config();
if (starts.size() <= 1) {
- return false;
+ return List.of(
+ FileSplit.withReadSchema(
+ "file",
+ deferred.filePath(),
+ 0,
+ fileLength,
+ deferred.format(),
+ wholeFileSplitConfig(config),
+ deferred.partitionValues(),
+ deferred.columnMapping(),
+ deferred.readSchema()
+ )
+ );
}
+ List splits = new ArrayList<>(starts.size());
for (int i = 0; i < starts.size(); i++) {
long start = starts.get(i);
long end = (i + 1 < starts.size()) ? starts.get(i + 1) : fileLength;
@@ -973,10 +1247,20 @@ private boolean tryNewlineAlignedMacroSplits(
splitConfig.put(LAST_SPLIT_KEY, "true");
}
splits.add(
- FileSplit.withReadSchema("file", filePath, start, length, format, splitConfig, partitionValues, columnMapping, readSchema)
+ FileSplit.withReadSchema(
+ "file",
+ deferred.filePath(),
+ start,
+ length,
+ deferred.format(),
+ splitConfig,
+ deferred.partitionValues(),
+ deferred.columnMapping(),
+ deferred.readSchema()
+ )
);
}
- return true;
+ return splits;
}
static boolean isNewlineMacroSplitCandidateExtension(@Nullable String format) {
@@ -987,7 +1271,7 @@ static boolean isNewlineMacroSplitCandidateExtension(@Nullable String format) {
return ".ndjson".equals(f) || ".jsonl".equals(f) || ".json".equals(f) || ".csv".equals(f) || ".tsv".equals(f);
}
- /** Whether this leaf split came from {@link #tryNewlineAlignedMacroSplits}. */
+ /** Whether this leaf split came from {@link #buildNewlineMacroSplits}. */
public static boolean isRecordAlignedMacroSplit(FileSplit split) {
return split != null && "true".equals(split.config().get(RECORD_ALIGNED_MACRO_SPLIT_KEY));
}
@@ -1038,35 +1322,162 @@ 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}.
+ * The fixed stride offsets a strided macro-split walk probes: {@code k * targetStrideBytes} for
+ * {@code k = 1, 2, ...} while the offset is inside the file and leaves at least {@code minSegment} bytes
+ * behind it. Pure arithmetic and no I/O, so the planning phase can compute a file's probe positions before
+ * deciding how (or where) to run the probes.
+ *
+ * Fixed offsets are what make the probes independent of one another. The read-time segmentation in
+ * {@link ParallelParsingCoordinator#computeSegments} re-anchors each stride on the boundary the previous probe
+ * found, so its offsets drift relative to these; both partition a file at record boundaries, they just do not
+ * pick the same ones.
*/
- static List computeRecordAlignedMacroSplitStarts(
- SegmentableFormatReader reader,
+ static List macroSplitProbePositions(long fileLength, long targetStrideBytes, long minSegment) {
+ List positions = new ArrayList<>();
+ for (long pos = targetStrideBytes; pos < fileLength; pos += targetStrideBytes) {
+ if (fileLength - pos < minSegment) {
+ break;
+ }
+ positions.add(pos);
+ }
+ return positions;
+ }
+
+ /**
+ * Reduces probe outcomes, which must be in ascending position order, to macro-split start offsets seeded
+ * with the file start. An offset that yielded no boundary contributes nothing, so the macro-splits either
+ * side of it merge into one that spans it: one unsplittable stretch of the file costs one split rather than
+ * every split after it. A boundary that does not advance past the previous one is dropped, which is what
+ * adjacent stride offsets landing inside the same record produce.
+ */
+ static List reduceProbeOutcomes(List outcomes) {
+ List boundaries = new ArrayList<>();
+ boundaries.add(0L);
+ for (ProbeOutcome outcome : outcomes) {
+ if (outcome.found() && outcome.boundary() > boundaries.get(boundaries.size() - 1)) {
+ boundaries.add(outcome.boundary());
+ }
+ }
+ return boundaries;
+ }
+
+ /**
+ * Probes each position on the calling thread. Every probe is independent of every other, so this produces the
+ * same boundaries as gathering the same positions concurrently and reducing the results, which is what lets a
+ * node without a discovery executor fall back to this walk without changing how a file is split.
+ */
+ static List probeStridedBoundariesSerially(
+ RecordSplitter splitter,
+ StorageObject storageObject,
+ long fileLength,
+ List positions,
+ long minSegment,
+ BooleanSupplier isCancelled
+ ) throws IOException {
+ List outcomes = new ArrayList<>(positions.size());
+ for (long pos : positions) {
+ outcomes.add(probeStridedBoundary(splitter, storageObject, pos, fileLength, minSegment, isCancelled));
+ }
+ return reduceProbeOutcomes(outcomes);
+ }
+
+ /**
+ * The result of probing one stride offset: either the record boundary to split at, or {@link #NONE} when the
+ * offset yields no split start, because no boundary lies within its probe window or because the boundary it
+ * found would leave too little of the file behind it.
+ *
+ * {@link #NONE} is local to its own offset: the macro-splits on either side of it merge into one, and every
+ * other offset's boundary still stands. Nothing about one offset's outcome constrains another's, which is what
+ * lets the probes run in any order.
+ */
+ record ProbeOutcome(boolean found, long boundary) {
+ static final ProbeOutcome NONE = new ProbeOutcome(false, -1L);
+
+ static ProbeOutcome at(long boundary) {
+ return new ProbeOutcome(true, boundary);
+ }
+ }
+
+ /**
+ * Finds the first record boundary at or after {@code pos} by reading a bounded window there.
+ *
+ * The probe reads at most {@link #MACRO_SPLIT_PROBE_WINDOW_BYTES}, then drains the rest of that window and
+ * closes the stream rather than aborting it: a partially-read body is dropped by providers like S3 along
+ * with its pooled HTTP connection, so the next probe would pay a fresh connection setup. Draining a small
+ * bounded window instead returns the connection to the pool for the next probe to reuse, which is worth far
+ * more than the few extra bytes transferred. On failure or cancellation the stream is aborted instead, so
+ * the connection and its storage permit are released immediately rather than after a pointless drain.
+ *
+ * Runs under a {@link StorageRetryCancellation} scope so a probe parked in retry/throttle backoff aborts on
+ * cancel instead of sleeping out its retry budget. The scope is thread-local, so it must be installed here,
+ * on whichever thread runs the probe, rather than once around the whole discovery.
+ */
+ static ProbeOutcome probeStridedBoundary(
+ RecordSplitter splitter,
+ StorageObject storageObject,
+ long pos,
+ long fileLength,
+ long minSegment,
+ BooleanSupplier isCancelled
+ ) throws IOException {
+ if (isCancelled.getAsBoolean()) {
+ throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE);
+ }
+ return StorageRetryCancellation.callWithCancellation(isCancelled, () -> {
+ long probeWindow = Math.min(MACRO_SPLIT_PROBE_WINDOW_BYTES, fileLength - pos);
+ InputStream stream = storageObject.newStream(pos, probeWindow);
+ long skipped;
+ // Not try-with-resources: close() drains on providers like S3, which is what we want after a
+ // successful scan but not on the failure path, where abortStream releases the connection at once.
+ boolean drained = false;
+ try {
+ skipped = splitter.findNextRecordBoundary(stream);
+ // A query cancelled while this probe was scanning has no follow-up probe to hand the pooled
+ // connection to, so there is nothing to buy by draining the rest of the window.
+ if (isCancelled.getAsBoolean()) {
+ throw new TaskCancelledException(DISCOVERY_CANCELLED_MESSAGE);
+ }
+ stream.transferTo(OutputStream.nullOutputStream());
+ drained = true;
+ } finally {
+ if (drained) {
+ stream.close();
+ } else {
+ storageObject.abortStream(stream);
+ }
+ }
+ // No boundary within the window: its end was reached, or the record exceeds the splitter's maximum.
+ // Either way this offset yields no split start, and the split before it runs on through the record
+ // that swallowed the window until the next offset that does find one.
+ if (skipped == RecordSplitter.RECORD_TOO_LARGE || skipped < 0) {
+ return ProbeOutcome.NONE;
+ }
+ long boundary = pos + skipped;
+ // Splitting this close to the end would leave a runt final split. Every higher offset resolves to a
+ // boundary at least this far in, so they all yield nothing too and the last split extends to EOF.
+ if (boundary >= fileLength || fileLength - boundary < minSegment) {
+ return ProbeOutcome.NONE;
+ }
+ return ProbeOutcome.at(boundary);
+ });
+ }
+
+ /**
+ * Macro-split starts for a splitter that cannot be probed at a fixed offset (quoted/escaped CSV/TSV) but can
+ * prove a record start ({@link RecordSplitter#supportsProvenProbing()}). Each iteration's position depends on
+ * the boundary the previous one found, and an {@code AMBIGUOUS} probe walks forward from the last proven
+ * record start, so this walk is inherently sequential and cannot be parallelized like the strided one.
+ */
+ static List computeProvenMacroSplitStarts(
+ RecordSplitter splitter,
StorageObject storageObject,
long fileLength,
long targetStrideBytes,
- int maxRecordBytes,
+ long minSegment,
BooleanSupplier isCancelled
) throws IOException {
List boundaries = new ArrayList<>();
boundaries.add(0L);
- long minSegment = reader.minimumSegmentSize();
- RecordSplitter splitter = reader.recordSplitter(maxRecordBytes);
- boolean strided = splitter.supportsStridedProbing();
- boolean proven = splitter.supportsProvenProbing();
- // A strided splitter (plain CSV/TSV, NDJSON) probes any stride offset directly. A non-strided splitter
- // (quoted/escaped CSV/TSV) can still be macro-split when it supports proven probing: the file start is a
- // known record start, so exactCursor seeds at 0 and every emitted boundary is a proven record start.
- // A splitter that is neither must have been routed to a whole-file split upstream; if one arrives here
- // that gate failed, so fail loud rather than emit mis-aligned macro-splits that silently mis-count.
- if (strided == false && proven == false) {
- throw new IllegalStateException(
- "record splitter ["
- + splitter.getClass().getName()
- + "] supports neither strided nor proven probing and cannot be macro-split"
- );
- }
// 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,43 +1488,30 @@ 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)) {
- probed = splitter.findProvenRecordBoundary(probeStream);
+ long probed;
+ InputStream probeStream = storageObject.newStream(pos, remaining);
+ try (Closeable abortOnExit = () -> storageObject.abortStream(probeStream)) {
+ probed = splitter.findProvenRecordBoundary(probeStream);
+ }
+ if (probed >= 0) {
+ boundary = pos + probed;
+ } 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)) {
+ start = splitter.findRecordStartAtOrAfter(walkStream, pos - exactCursor, isCancelled);
}
- if (probed >= 0) {
- boundary = pos + probed;
- } 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)) {
- start = splitter.findRecordStartAtOrAfter(walkStream, pos - exactCursor, isCancelled);
- }
- if (start == RecordSplitter.RECORD_TOO_LARGE || start < 0) {
- break;
- }
- boundary = exactCursor + start;
- } else {
- // findProvenRecordBoundary only ever returns a boundary (>= 0) or AMBIGUOUS.
- assert false : "findProvenRecordBoundary returned an unexpected sentinel: " + probed;
+ if (start == RecordSplitter.RECORD_TOO_LARGE || start < 0) {
break;
}
+ boundary = exactCursor + start;
+ } else {
+ // findProvenRecordBoundary only ever returns a boundary (>= 0) or AMBIGUOUS.
+ assert false : "findProvenRecordBoundary returned an unexpected sentinel: " + probed;
+ break;
}
if (boundary >= fileLength) {
break;
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..631650463a66c 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
@@ -15,6 +15,7 @@
import org.elasticsearch.compute.data.BlockFactory;
import org.elasticsearch.compute.data.Page;
import org.elasticsearch.compute.operator.CloseableIterator;
+import org.elasticsearch.core.Nullable;
import org.elasticsearch.tasks.TaskCancelledException;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
@@ -64,19 +65,28 @@
import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
+import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BooleanSupplier;
import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.lessThan;
import static org.mockito.ArgumentMatchers.any;
@@ -726,8 +736,11 @@ private void assertNewlineAlignedMacroSplitsDisjointAndMarked(
return -1L;
});
+ // The stride is floored to the probe window, so both it and the payload are sized in windows: several
+ // strides' worth of file, with slack past the last offset so the final split is not a runt.
+ long stride = FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES;
StringBuilder sb = new StringBuilder();
- for (int i = 0; i < 4000; i++) {
+ while (sb.length() < 4 * stride + stride / 2) {
sb.append(lineContent);
}
byte[] payload = sb.toString().getBytes(StandardCharsets.UTF_8);
@@ -740,7 +753,6 @@ private void assertNewlineAlignedMacroSplitsDisjointAndMarked(
StorageProviderRegistry storageRegistry = createPayloadStorageRegistry(payload);
- long stride = 3000;
FileSplitProvider splitter = new FileSplitProvider(
stride,
new DecompressionCodecRegistry(),
@@ -921,6 +933,628 @@ private List discoverRealDelimitedSplits(
return splitter.discoverSplits(ctx).splits();
}
+ /**
+ * Concurrency budgets are per node, not per file: every deferred file's boundary probes share one budget, so
+ * the ceiling never multiplies by the number of files being probed. A configured concurrency of {@code 0}
+ * turns permit limiting off rather than meaning "no concurrency", so it must not collapse the budget.
+ */
+ public void testProbeConcurrencyIsClampedToBlobStoreConcurrency() {
+ assertEquals(4, probeConcurrencyFor(Settings.builder().put("esql.external.max_concurrent_requests", 4).build()));
+ int ceiling = FileSplitProvider.MAX_PARALLEL_SPLIT_DISCOVERY;
+ assertEquals(ceiling, probeConcurrencyFor(Settings.builder().put("esql.external.max_concurrent_requests", ceiling).build()));
+ assertEquals(
+ FileSplitProvider.MAX_PARALLEL_SPLIT_DISCOVERY,
+ probeConcurrencyFor(Settings.builder().put("esql.external.max_concurrent_requests", 200).build())
+ );
+ assertEquals(
+ "permit limiting disabled must not disable concurrency",
+ FileSplitProvider.MAX_PARALLEL_SPLIT_DISCOVERY,
+ probeConcurrencyFor(Settings.builder().put("esql.external.max_concurrent_requests", 0).build())
+ );
+ }
+
+ private static int probeConcurrencyFor(Settings settings) {
+ return new FileSplitProvider(1024, new DecompressionCodecRegistry(), null, null, settings).probeConcurrency();
+ }
+
+ /**
+ * With a discovery executor the boundary probes of a multi-file query run concurrently, and they produce
+ * exactly the splits the serial path produces. The two must agree: probing fixed stride offsets makes each
+ * probe independent of the others, so the order they complete in cannot change the boundary set.
+ */
+ public void testStridedProbesRunConcurrentlyAndAgreeWithSerialDiscovery() throws Exception {
+ Map payloads = Map.of("one.csv", delimitedPayload("a,b,c\n"), "two.csv", delimitedPayload("d,e,f\n"));
+ // Two probes per file (stride and minimum segment are both 1 MiB against a ~3.5 MiB payload).
+ StreamTracking tracking = new StreamTracking(2);
+
+ List serial = discoverPlainCsvSplits(payloads, CSV_MIN_SEGMENT_BYTES, null, null);
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ List parallel;
+ try {
+ parallel = discoverPlainCsvSplits(payloads, CSV_MIN_SEGMENT_BYTES, executor, tracking);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertThat("both files must macro-split", serial.size(), greaterThan(2));
+ assertEquals("parallel probing must not change the split set", describe(serial), describe(parallel));
+ assertThat("probes must overlap when an executor is available", tracking.peakInFlight.get(), greaterThan(1));
+ }
+
+ /**
+ * The motivating case is one very large file, not many files, so a single file's own boundary probes must run
+ * concurrently. This also checks the boundaries themselves against ground truth on a payload with varying line
+ * lengths, where an off-byte in the probe arithmetic would land a split mid-record: every emitted split must
+ * begin at a real record start, and the splits together must cover the file exactly once.
+ */
+ public void testSingleLargeFileProbesConcurrentlyOnTrueRecordStarts() throws Exception {
+ StringBuilder csv = new StringBuilder();
+ int row = 0;
+ while (csv.length() < 4 * CSV_MIN_SEGMENT_BYTES) {
+ // Line lengths cycle so stride offsets land at varying depths into a record.
+ csv.append(row).append(",value,").append("x".repeat(row % 37)).append('\n');
+ row++;
+ }
+ byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8);
+ Set trueStarts = trueRecordStarts(stridedSplitter(), payload);
+ StreamTracking tracking = new StreamTracking(2);
+
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ List splits;
+ try {
+ splits = discoverPlainCsvSplits(Map.of("big.csv", payload), 256 * 1024, executor, tracking);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertThat("one large file must still macro-split several ways", splits.size(), greaterThan(2));
+ assertThat("a single file's probes must overlap", tracking.peakInFlight.get(), greaterThan(1));
+ long expectedOffset = 0;
+ for (ExternalSplit split : splits) {
+ FileSplit fileSplit = (FileSplit) split;
+ assertEquals("splits must be contiguous", expectedOffset, fileSplit.offset());
+ assertTrue("split at " + fileSplit.offset() + " must begin on a record start", trueStarts.contains(fileSplit.offset()));
+ assertThat("no empty split", fileSplit.length(), greaterThan(0L));
+ expectedOffset += fileSplit.length();
+ }
+ assertEquals("splits must cover the whole file", payload.length, expectedOffset);
+ }
+
+ /**
+ * A file whose every record spans the probe window offers no boundary at any offset, so it is read whole. The
+ * price of an offset that finds nothing is one read: it says nothing about the offsets after it, so none of
+ * their reads may be skipped on the strength of it.
+ */
+ public void testAFileWithNoBoundaryInAnyProbeWindowIsReadWhole() throws Exception {
+ int window = Math.toIntExact(FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES);
+ long stride = 2L * window;
+ // Records of four windows against a stride of two windows: each record terminator lands an odd number of
+ // windows into the file while every probe covers an even one, so no probe window holds a terminator.
+ byte[] payload = ("y".repeat(4 * window - 1) + "\n").repeat(4).getBytes(StandardCharsets.UTF_8);
+ StreamTracking tracking = new StreamTracking(1);
+
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ List splits;
+ try {
+ splits = discoverPlainCsvSplits(Map.of("long-lines.csv", payload), stride, executor, tracking);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertEquals("a file with no usable boundary is read whole", 1, splits.size());
+ int probes = FileSplitProvider.macroSplitProbePositions(payload.length, stride, CSV_MIN_SEGMENT_BYTES).size();
+ assertThat("the payload must offer several offsets to probe", probes, greaterThan(1));
+ assertEquals("an offset that finds nothing must not suppress the others", probes, tracking.opens.get());
+ }
+
+ /**
+ * A target split size below the probe window is floored to the window: below it, probe offsets stop being
+ * disjoint (several land inside one window's worth of file, each drained in full) and the number of probe
+ * tasks materialized during planning is {@code fileLength / stride}, unbounded below. Flooring caps a file's
+ * probes at {@code fileLength / window}, so a tiny stride must probe exactly as a window-sized one does.
+ */
+ public void testTargetStrideBelowTheProbeWindowIsFlooredToTheWindow() {
+ Map payloads = Map.of("tiny-stride.csv", delimitedPayload("a,b,c\n"));
+
+ List floored = discoverPlainCsvSplits(payloads, 1024, null, null);
+ List atWindow = discoverPlainCsvSplits(payloads, FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES, null, null);
+
+ assertThat("the file must still macro-split", floored.size(), greaterThan(1));
+ assertEquals("a sub-window stride must split exactly like a window-sized one", describe(atWindow), describe(floored));
+ }
+
+ /**
+ * A record too long for the probe window costs the one split its offset would have started, and no more: the
+ * offsets past it are probed as usual and their boundaries stand, so the splits either side of the record
+ * merge into one that spans it. Serial and concurrent discovery must agree, which they do because an offset's
+ * outcome depends on nothing but its own read.
+ */
+ public void testARecordExceedingTheWindowMidFileCostsOnlyItsOwnSplit() throws Exception {
+ int window = Math.toIntExact(FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES);
+ long stride = 2L * window;
+ StringBuilder csv = new StringBuilder();
+ // Normal rows up to just short of the second stride offset, so the probe there falls inside the long row
+ // that follows and finds no boundary within its window.
+ while (csv.length() < 2 * stride - window) {
+ csv.append("a,b,c\n");
+ }
+ long longRowStart = csv.length();
+ csv.append("z".repeat(3 * window)).append('\n');
+ while (csv.length() < 8 * stride) {
+ csv.append("a,b,c\n");
+ }
+ byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8);
+ Map payloads = Map.of("mid-file.csv", payload);
+
+ List serial = discoverPlainCsvSplits(payloads, stride, null, null);
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ List parallel;
+ try {
+ parallel = discoverPlainCsvSplits(payloads, stride, executor, null);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertEquals("an offset's outcome must not depend on when it is probed", describe(serial), describe(parallel));
+ // One offset lands inside the long row and yields nothing; the file start and the rest of the offsets each
+ // start a split, so the long row costs exactly one of them.
+ int probes = FileSplitProvider.macroSplitProbePositions(payload.length, stride, CSV_MIN_SEGMENT_BYTES).size();
+ assertEquals("only the offset inside the long record loses its split", probes, serial.size());
+ FileSplit spanning = (FileSplit) serial.get(1);
+ assertThat("one split must span the record no probe could split", spanning.offset(), lessThan(longRowStart));
+ assertThat(spanning.offset() + spanning.length(), greaterThan(longRowStart + 3L * window));
+ }
+
+ /**
+ * Probe concurrency is a per-query budget, so more probes than the budget must not all be in flight at once.
+ * The budget follows the node's blob-store concurrency, which is what actually limits open streams.
+ */
+ public void testProbesInFlightStayWithinTheConcurrencyBudget() throws Exception {
+ Settings settings = Settings.builder().put("esql.external.max_concurrent_requests", 2).build();
+ byte[] payload = delimitedPayload("a,b,c,d,e\n");
+ // A 128 KiB stride over a ~3.5 MiB payload leaves far more probe offsets than the budget of 2.
+ StreamTracking tracking = new StreamTracking(2);
+
+ ExecutorService executor = Executors.newFixedThreadPool(8);
+ try {
+ discoverPlainCsvSplits(Map.of("wide.csv", payload), 128 * 1024, executor, tracking, settings);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertThat("more probes than the budget must have run", tracking.opens.get(), greaterThan(2));
+ assertEquals("in-flight probes must not exceed the budget", 2, tracking.peakInFlight.get());
+ }
+
+ /**
+ * The budget covers the whole query, so probing several files together must not multiply it. Probing one file
+ * at a time under its own budget would allow files times budget streams open at once, which is what pooling
+ * every deferred file's offsets into one batch prevents.
+ */
+ public void testProbesOfDifferentFilesShareOneConcurrencyBudget() throws Exception {
+ Settings budgetOfTwo = Settings.builder().put("esql.external.max_concurrent_requests", 2).build();
+ Map payloads = new HashMap<>();
+ for (int file = 0; file < 4; file++) {
+ payloads.put("f" + file + ".csv", delimitedPayload("a,b,c," + file + "\n"));
+ }
+ // A 128 KiB stride over ~3.5 MiB gives each file far more probe offsets than the budget of 2.
+ StreamTracking tracking = new StreamTracking(2);
+
+ ExecutorService executor = Executors.newFixedThreadPool(8);
+ try {
+ discoverPlainCsvSplits(payloads, 128 * 1024, executor, tracking, budgetOfTwo);
+ } finally {
+ executor.shutdown();
+ }
+
+ assertThat("every file must contribute probes", tracking.opens.get(), greaterThan(payloads.size() * 2));
+ assertEquals("one budget across all files, not one per file", 2, tracking.peakInFlight.get());
+ }
+
+ /**
+ * Planning a strided file reads nothing, so a cancel arriving while it runs is not seen by any read. It must be
+ * seen between the phases instead: a cancelled query must not even dispatch its probes, let alone read.
+ */
+ public void testCancellationBetweenPlanningAndProbingDispatchesNoProbe() throws Exception {
+ byte[] payload = delimitedPayload("a,b,c\n");
+ StreamTracking tracking = new StreamTracking(1);
+ // Planning hands out the file's storage object without reading it, which is the window under test.
+ BooleanSupplier cancelOncePlanned = () -> tracking.objects.get() > 0;
+
+ ExecutorService pool = Executors.newFixedThreadPool(2);
+ AtomicInteger dispatched = new AtomicInteger();
+ Executor countingExecutor = command -> {
+ dispatched.incrementAndGet();
+ pool.execute(command);
+ };
+ try {
+ expectThrows(
+ TaskCancelledException.class,
+ () -> discoverPlainCsvSplits(
+ Map.of("planned.csv", payload),
+ 256 * 1024,
+ countingExecutor,
+ tracking,
+ Settings.EMPTY,
+ cancelOncePlanned
+ )
+ );
+ } finally {
+ pool.shutdown();
+ }
+ // The single planned file is planned on the calling thread, so anything dispatched would be a probe.
+ assertEquals("a cancelled query must not dispatch probes", 0, dispatched.get());
+ assertEquals("a cancel seen before probing must not read", 0, tracking.opens.get());
+ }
+
+ /**
+ * A probe read that fails must fail the query. Probes run on other threads and each one only ever adds a
+ * boundary, so an I/O error that did not travel back out of the gather would leave the file quietly split at
+ * the boundaries the surviving probes happened to find.
+ */
+ public void testAFailedProbeReadFailsDiscovery() throws Exception {
+ byte[] payload = delimitedPayload("a,b,c\n");
+ // The first stream any probe opens fails; planning reads nothing, so no other read can be the first.
+ StreamTracking tracking = new StreamTracking(1, 1);
+
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ RuntimeException failure;
+ try {
+ failure = expectThrows(
+ RuntimeException.class,
+ () -> discoverPlainCsvSplits(Map.of("broken.csv", payload), 256 * 1024, executor, tracking)
+ );
+ } finally {
+ executor.shutdown();
+ }
+
+ assertThat(failure.getCause(), instanceOf(IOException.class));
+ assertEquals("connection reset", failure.getCause().getMessage());
+ }
+
+ /**
+ * A cancel landing between two probes is seen by the next one before it opens a stream, so a query cancelled
+ * part-way through a file's offsets fails rather than reading out the rest of them.
+ */
+ public void testCancellationBetweenProbesIsSeenBeforeTheNextRead() throws Exception {
+ byte[] payload = delimitedPayload("a,b,c\n");
+ StreamTracking tracking = new StreamTracking(1);
+ // Cancelled once the first probe has finished with its stream, so the cancel lands between two probes
+ // rather than inside one.
+ BooleanSupplier cancelAfterFirstProbe = () -> tracking.closes.get() > 0;
+
+ // A single thread makes the ordering deterministic: no second probe starts before the first one closes.
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ expectThrows(
+ TaskCancelledException.class,
+ () -> discoverPlainCsvSplits(
+ Map.of("cancelled.csv", payload),
+ 256 * 1024,
+ executor,
+ tracking,
+ Settings.EMPTY,
+ cancelAfterFirstProbe
+ )
+ );
+ } finally {
+ executor.shutdown();
+ }
+ assertEquals("the probes after the cancel must not read", 1, tracking.opens.get());
+ }
+
+ /**
+ * A cancel arriving once probing has started stops the query rather than reading out every remaining offset.
+ */
+ public void testCancellationDuringProbingAbortsDiscovery() throws Exception {
+ byte[] payload = delimitedPayload("a,b,c\n");
+ StreamTracking tracking = new StreamTracking(1);
+ // Cancelled as soon as the first probe opens its stream, so the cancel lands inside the probe phase
+ // rather than during planning.
+ BooleanSupplier cancelOnFirstRead = () -> tracking.opens.get() > 0;
+
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ expectThrows(
+ TaskCancelledException.class,
+ () -> discoverPlainCsvSplits(
+ Map.of("cancelled.csv", payload),
+ 256 * 1024,
+ executor,
+ tracking,
+ Settings.EMPTY,
+ cancelOnFirstRead
+ )
+ );
+ } finally {
+ executor.shutdown();
+ }
+ assertEquals("discovery must stop at the probe that saw the cancel", 1, tracking.opens.get());
+ }
+
+ /**
+ * Splits are emitted in file order even though only some files' boundaries are probed in the deferred phase:
+ * a small file finishes during planning while a large one comes back later, and the large one's macro-splits
+ * must still land in its file's position rather than after every planned file.
+ */
+ public void testSplitOrderFollowsFileOrderWhenOnlySomeFilesAreProbed() throws Exception {
+ byte[] small = "x,y,z\n".repeat(4).getBytes(StandardCharsets.UTF_8);
+ Map payloads = Map.of("a-small.csv", small, "b-large.csv", delimitedPayload("a,b,c\n"), "c-small.csv", small);
+
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ List splits;
+ try {
+ splits = discoverPlainCsvSplits(payloads, CSV_MIN_SEGMENT_BYTES, executor, null);
+ } finally {
+ executor.shutdown();
+ }
+
+ List objectOrder = new ArrayList<>();
+ for (ExternalSplit split : splits) {
+ String objectName = ((FileSplit) split).path().objectName();
+ if (objectOrder.isEmpty() || objectOrder.get(objectOrder.size() - 1).equals(objectName) == false) {
+ objectOrder.add(objectName);
+ }
+ }
+ assertEquals(List.of("a-small.csv", "b-large.csv", "c-small.csv"), objectOrder);
+ assertThat("the large file must contribute several macro-splits", splits.size(), greaterThan(3));
+ }
+
+ /** A plain-CSV payload of ~3.5 MiB: above two minimum segments, so it yields several macro-splits. */
+ private static byte[] delimitedPayload(String lineContent) {
+ StringBuilder sb = new StringBuilder();
+ while (sb.length() < 3 * CSV_MIN_SEGMENT_BYTES + CSV_MIN_SEGMENT_BYTES / 2) {
+ sb.append(lineContent);
+ }
+ return sb.toString().getBytes(StandardCharsets.UTF_8);
+ }
+
+ /** Offsets, lengths, and position stamps of a split list, in order: enough to tell two split sets apart. */
+ private static List describe(List splits) {
+ List described = new ArrayList<>(splits.size());
+ for (ExternalSplit split : splits) {
+ FileSplit fileSplit = (FileSplit) split;
+ described.add(
+ fileSplit.path()
+ + "["
+ + fileSplit.offset()
+ + ","
+ + fileSplit.length()
+ + ",first="
+ + FileSplitProvider.isFirstInFile(fileSplit)
+ + ",last="
+ + FileSplitProvider.isLastInFile(fileSplit)
+ + "]"
+ );
+ }
+ return described;
+ }
+
+ private static List discoverPlainCsvSplits(
+ Map payloads,
+ long targetStrideBytes,
+ @Nullable Executor executor,
+ @Nullable StreamTracking tracking
+ ) {
+ return discoverPlainCsvSplits(payloads, targetStrideBytes, executor, tracking, Settings.EMPTY);
+ }
+
+ private static List discoverPlainCsvSplits(
+ Map payloads,
+ long targetStrideBytes,
+ @Nullable Executor executor,
+ @Nullable StreamTracking tracking,
+ Settings settings
+ ) {
+ return discoverPlainCsvSplits(payloads, targetStrideBytes, executor, tracking, settings, () -> false);
+ }
+
+ /**
+ * Runs full split discovery over one or more plain-mode CSV files, optionally on an executor and against a
+ * stream-tracking storage provider.
+ */
+ private static List discoverPlainCsvSplits(
+ Map payloads,
+ long targetStrideBytes,
+ @Nullable Executor executor,
+ @Nullable StreamTracking tracking,
+ Settings settings,
+ BooleanSupplier isCancelled
+ ) {
+ FormatReaderRegistry formatRegistry = new FormatReaderRegistry(new DecompressionCodecRegistry());
+ formatRegistry.registerLazy(
+ "csv",
+ (s, bf) -> new CsvFormatReader(bf, CsvFormatOptions.DEFAULT, "csv", List.of(".csv")),
+ Settings.EMPTY,
+ null
+ );
+ formatRegistry.registerExtension(".csv", "csv");
+ formatRegistry.byName("csv");
+
+ FileSplitProvider provider = new FileSplitProvider(
+ targetStrideBytes,
+ new DecompressionCodecRegistry(),
+ createMultiFileStorageRegistry(payloads, tracking),
+ formatRegistry,
+ settings,
+ executor
+ );
+
+ List entries = new ArrayList<>();
+ for (String objectName : new TreeSet<>(payloads.keySet())) {
+ entries.add(new StorageEntry(StoragePath.of("s3://b/" + objectName), payloads.get(objectName).length, Instant.EPOCH));
+ }
+ FileList fileList = GlobExpander.fileListOf(entries, "s3://b/*.csv");
+ SplitDiscoveryContext ctx = new SplitDiscoveryContext(
+ null,
+ fileList,
+ Map.of(),
+ Map.of("mode", "plain"),
+ PartitionMetadata.EMPTY,
+ List.of(),
+ ExternalSchema.EMPTY,
+ null,
+ SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES,
+ isCancelled,
+ DeclaredReadSpec.NONE
+ );
+ return provider.discoverSplits(ctx).splits();
+ }
+
+ /**
+ * Counts storage streams opened and how many were open at once. Each opened stream waits for
+ * {@code expectedOverlap} peers to arrive, bounded by a timeout so a serialized run still finishes and simply
+ * reports a peak of 1. Without that wait, in-memory reads finish so fast that an assertion about overlapping
+ * probes would depend on scheduling luck.
+ */
+ private static final class StreamTracking {
+ private static final long OVERLAP_WAIT_MILLIS = 5_000;
+ private static final int NEVER_FAIL = 0;
+
+ final AtomicInteger peakInFlight = new AtomicInteger();
+ final AtomicInteger opens = new AtomicInteger();
+ final AtomicInteger closes = new AtomicInteger();
+ /** Storage objects handed out, which planning does before any stream is opened. */
+ final AtomicInteger objects = new AtomicInteger();
+ private final AtomicInteger inFlight = new AtomicInteger();
+ private final CountDownLatch overlap;
+ /** Ordinal of the stream whose reads fail, standing in for a mid-probe I/O error, or {@link #NEVER_FAIL}. */
+ private final int failReadsOnOpen;
+
+ StreamTracking(int expectedOverlap) {
+ this(expectedOverlap, NEVER_FAIL);
+ }
+
+ StreamTracking(int expectedOverlap, int failReadsOnOpen) {
+ this.overlap = new CountDownLatch(expectedOverlap);
+ this.failReadsOnOpen = failReadsOnOpen;
+ }
+
+ /** @return the ordinal of the stream just opened, counting from one */
+ int opened() throws InterruptedException {
+ int ordinal = opens.incrementAndGet();
+ int current = inFlight.incrementAndGet();
+ peakInFlight.accumulateAndGet(current, Math::max);
+ overlap.countDown();
+ overlap.await(OVERLAP_WAIT_MILLIS, TimeUnit.MILLISECONDS);
+ return ordinal;
+ }
+
+ void closed() {
+ closes.incrementAndGet();
+ inFlight.decrementAndGet();
+ }
+ }
+
+ /** Storage provider serving a distinct payload per object name, so a multi-file test can vary file sizes. */
+ private static StorageProviderRegistry createMultiFileStorageRegistry(Map payloads, @Nullable StreamTracking tracking) {
+ StorageProviderRegistry registry = new StorageProviderRegistry(Settings.EMPTY);
+ StorageProvider provider = new StorageProvider() {
+ @Override
+ public StorageObject newObject(StoragePath path) {
+ return newObject(path, payloadFor(path).length);
+ }
+
+ @Override
+ public StorageObject newObject(StoragePath path, long length) {
+ return newObject(path, length, Instant.EPOCH);
+ }
+
+ @Override
+ public StorageObject newObject(StoragePath path, long length, Instant lastModified) {
+ byte[] payload = payloadFor(path);
+ if (tracking != null) {
+ tracking.objects.incrementAndGet();
+ }
+ return new StorageObject() {
+ @Override
+ public InputStream newStream() {
+ return trackedStream(new ByteArrayInputStream(payload));
+ }
+
+ @Override
+ public InputStream newStream(long position, long len) {
+ return trackedStream(new ByteArrayInputStream(payload, Math.toIntExact(position), Math.toIntExact(len)));
+ }
+
+ @Override
+ public long length() {
+ return payload.length;
+ }
+
+ @Override
+ public Instant lastModified() {
+ return lastModified;
+ }
+
+ @Override
+ public boolean exists() {
+ return true;
+ }
+
+ @Override
+ public StoragePath path() {
+ return path;
+ }
+ };
+ }
+
+ private InputStream trackedStream(InputStream delegate) {
+ if (tracking == null) {
+ return delegate;
+ }
+ int ordinal;
+ try {
+ ordinal = tracking.opened();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError("interrupted while tracking stream", e);
+ }
+ if (ordinal == tracking.failReadsOnOpen) {
+ return new InputStream() {
+ @Override
+ public int read() throws IOException {
+ throw new IOException("connection reset");
+ }
+ };
+ }
+ return new FilterInputStream(delegate) {
+ @Override
+ public void close() throws IOException {
+ tracking.closed();
+ super.close();
+ }
+ };
+ }
+
+ private byte[] payloadFor(StoragePath path) {
+ byte[] payload = payloads.get(path.objectName());
+ assertNotNull("no payload registered for " + path, payload);
+ return payload;
+ }
+
+ @Override
+ public StorageIterator listObjects(StoragePath prefix, boolean recursive) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public boolean exists(StoragePath path) {
+ return payloads.containsKey(path.objectName());
+ }
+
+ @Override
+ public List supportedSchemes() {
+ return List.of("s3");
+ }
+
+ @Override
+ public void close() {}
+ };
+ registry.registerFactory("s3", StorageProviderFactory.noConfigKeys(() -> provider));
+ return registry;
+ }
+
/**
* Full {@link FileSplitProvider#discoverSplits} path for a quoted CSV file whose quoted fields carry
* embedded newlines and {@code ""}-escaped quotes: it emits multiple record-aligned macro-splits (proving
@@ -947,7 +1581,7 @@ public void testDiscoverSplitsMacroSplitsQuotedCsv() {
/**
* Quoted CSV (the default {@code .csv} mode, whose quoted fields may embed newlines) macro-splits via the
- * proven-probe path: {@link FileSplitProvider#computeRecordAlignedMacroSplitStarts} emits boundaries for
+ * proven-probe path: {@link FileSplitProvider#computeProvenMacroSplitStarts} emits boundaries for
* a non-strided but proven-capable splitter. Every emitted boundary must be a true record start, checked
* against the trusted sequential scanner {@link RecordSplitter#findNextRecordBoundary} looped from the file
* start (its prefix sums are the true record starts), and the boundaries must be strictly increasing. The
@@ -982,12 +1616,12 @@ public void testRecordAlignedMacroSplitDiscoveryProvesQuotedCsvBoundaries() thro
Set trueStarts = trueRecordStarts(csvReader.recordSplitter(SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES), payload);
long stride = fileLength / 4;
- List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts(
- csvReader,
+ List starts = FileSplitProvider.computeProvenMacroSplitStarts(
+ csvReader.recordSplitter(SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES),
obj,
fileLength,
stride,
- SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES,
+ csvReader.minimumSegmentSize(),
() -> false
);
@@ -1019,11 +1653,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()}.
+ * Each strided stride probe opens a bounded window (at most
+ * {@link FileSplitProvider#MACRO_SPLIT_PROBE_WINDOW_BYTES}), finds the next record boundary within it, then
+ * drains the window's remaining bytes and closes so the HTTP connection returns to the pool for the next
+ * probe to reuse. It must not abort (an aborted partial body drops the connection, forcing a fresh handshake
+ * per probe), and it must not open a range to end-of-file (which would drain far more than the window).
*/
- public void testComputeRecordAlignedMacroSplitStartsDoesNotDrainStream() throws IOException {
+ public void testSerialStridedProbesDrainBoundedProbeWindows() throws IOException {
var blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build();
StringBuilder csv = new StringBuilder("id,name\n");
@@ -1036,43 +1672,299 @@ public void testComputeRecordAlignedMacroSplitStartsDoesNotDrainStream() throws
DrainSimulatingStorageObject.Tracking tracking = new DrainSimulatingStorageObject.Tracking();
StorageObject object = DrainSimulatingStorageObject.create(payload, tracking);
- // Plain mode: the drain contract is format-agnostic, but macro-split discovery now refuses non-strided
- // splitters (default/quoted CSV), which are read whole-file instead. Plain CSV keeps strided probing.
+ // Plain mode keeps strided probing; macro-split discovery refuses non-strided (default/quoted) CSV,
+ // which is read whole-file instead.
var csvReader = (SegmentableFormatReader) new CsvFormatReader(blockFactory).withConfig(Map.of("mode", "plain"));
long stride = fileLength / 4;
- List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts(
- csvReader,
- object,
- fileLength,
- stride,
- SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES,
- () -> false
- );
+ List starts = serialStridedStarts(csvReader, object, fileLength, stride, SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES);
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("strided probes pool the connection by draining, never abort", 0, tracking.abortCalls.get());
+ assertTrue("probe streams must be closed", tracking.closed.get());
+ // Each successful probe drains its full window, so at least (boundaries) * window bytes are read: proof
+ // the window is drained (an abort-without-drain would read only the few bytes up to the first newline).
+ long window = FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES;
+ long probes = starts.size() - 1;
+ assertThat(tracking.bytesConsumed.get(), greaterThanOrEqualTo(probes * window));
+ // But far below the whole file: the probe reads a bounded window, never a range to end-of-file.
assertThat(
- "boundary probes must not drain the range streams; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes",
+ "boundary probes must read bounded windows; consumed " + tracking.bytesConsumed.get() + " of " + fileLength + " bytes",
tracking.bytesConsumed.get(),
lessThan(fileLength / 2)
);
}
- public void testRecordAlignedMacroSplitDiscoveryStopsOnMaxRecordSize() throws IOException {
+ // strided macro-split discovery: fixed probe offsets, independent probes
+
+ /**
+ * A strided (newline-terminated) record splitter. Plain-mode CSV turns quoting off, so every newline
+ * terminates a record and the splitter can be probed at any offset, exactly like NDJSON's.
+ */
+ private static RecordSplitter stridedSplitter() {
+ var blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build();
+ var reader = (SegmentableFormatReader) new CsvFormatReader(blockFactory).withConfig(Map.of("mode", "plain"));
+ return reader.recordSplitter(SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES);
+ }
+
+ /**
+ * A strided walk probes fixed multiples of the stride rather than re-anchoring on each boundary it finds.
+ * That independence is what lets the probes run concurrently: every position is known before any read.
+ */
+ public void testMacroSplitProbePositionsAreFixedStrideMultiples() {
+ assertEquals(List.of(100L, 200L, 300L, 400L), FileSplitProvider.macroSplitProbePositions(500, 100, 1));
+ }
+
+ /** A stride offset with less than a minimum segment behind it would produce a runt split, so it is not probed. */
+ public void testMacroSplitProbePositionsStopBeforeAShortTail() {
+ assertEquals(List.of(100L, 200L), FileSplitProvider.macroSplitProbePositions(500, 100, 250));
+ assertEquals(List.of(), FileSplitProvider.macroSplitProbePositions(500, 100, 500));
+ }
+
+ /** No offset is inside a file at or below one stride, which is the caller's "too small to split" case. */
+ public void testMacroSplitProbePositionsEmptyForFileWithinOneStride() {
+ assertEquals(List.of(), FileSplitProvider.macroSplitProbePositions(100, 100, 1));
+ assertEquals(List.of(), FileSplitProvider.macroSplitProbePositions(60, 100, 1));
+ }
+
+ /** The file start is always a split start, and each probed boundary follows in probe order. */
+ public void testReduceProbeOutcomesSeedsTheFileStart() {
+ List boundaries = FileSplitProvider.reduceProbeOutcomes(
+ List.of(FileSplitProvider.ProbeOutcome.at(120), FileSplitProvider.ProbeOutcome.at(240))
+ );
+ assertEquals(List.of(0L, 120L, 240L), boundaries);
+ }
+
+ /**
+ * An offset that yielded no boundary costs the one split it would have started: the boundaries after it still
+ * stand, and the splits either side of it merge into the one that spans it.
+ */
+ public void testReduceProbeOutcomesSkipsAnOffsetThatFoundNoBoundary() {
+ List boundaries = FileSplitProvider.reduceProbeOutcomes(
+ List.of(FileSplitProvider.ProbeOutcome.at(120), FileSplitProvider.ProbeOutcome.NONE, FileSplitProvider.ProbeOutcome.at(360))
+ );
+ assertEquals(List.of(0L, 120L, 360L), boundaries);
+ }
+
+ /** Two stride offsets landing inside one record resolve to the same boundary; only one split starts there. */
+ public void testReduceProbeOutcomesDropsBoundariesThatDoNotAdvance() {
+ List boundaries = FileSplitProvider.reduceProbeOutcomes(
+ List.of(FileSplitProvider.ProbeOutcome.at(120), FileSplitProvider.ProbeOutcome.at(120), FileSplitProvider.ProbeOutcome.at(240))
+ );
+ assertEquals(List.of(0L, 120L, 240L), boundaries);
+ }
+
+ /**
+ * A record longer than the probe window leaves no boundary to split at, so the probe yields none rather than
+ * reading on to end-of-file. Bounding the read is what keeps a probe a small, predictable ranged GET.
+ */
+ public void testProbeStridedBoundaryYieldsNoBoundaryWhenNoneInWindow() throws IOException {
+ int window = Math.toIntExact(FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES);
+ byte[] payload = new byte[4 * window];
+ Arrays.fill(payload, (byte) 'x');
+ payload[payload.length - 1] = '\n';
+ StorageObject object = createInMemoryStorageObject(payload, StoragePath.of("mem://long-line.ndjson"));
+
+ FileSplitProvider.ProbeOutcome outcome = FileSplitProvider.probeStridedBoundary(
+ stridedSplitter(),
+ object,
+ window,
+ payload.length,
+ 1,
+ () -> false
+ );
+
+ assertEquals("a record spanning the whole window leaves nothing to split at", FileSplitProvider.ProbeOutcome.NONE, outcome);
+ }
+
+ /** A boundary too close to end-of-file would leave a runt split, so the probe yields none instead. */
+ public void testProbeStridedBoundaryYieldsNoBoundaryWhenTailIsBelowMinimumSegment() throws IOException {
+ byte[] payload = "aaaa\nbbbb\ncccc\n".getBytes(StandardCharsets.UTF_8);
+ StorageObject object = createInMemoryStorageObject(payload, StoragePath.of("mem://short.ndjson"));
+ RecordSplitter splitter = stridedSplitter();
+
+ assertEquals(
+ FileSplitProvider.ProbeOutcome.at(10),
+ FileSplitProvider.probeStridedBoundary(splitter, object, 5, payload.length, 5, () -> false)
+ );
+ // With a 6-byte minimum segment, the same boundary leaves only 5 of the 15 bytes behind it.
+ assertEquals(
+ FileSplitProvider.ProbeOutcome.NONE,
+ FileSplitProvider.probeStridedBoundary(splitter, object, 5, payload.length, 6, () -> false)
+ );
+ }
+
+ /**
+ * A probe that fails mid-read aborts its stream rather than draining it: the read is not going to be reused,
+ * so the connection (and the storage permit it holds) should be released at once instead of after pointlessly
+ * transferring the rest of the window.
+ */
+ public void testProbeStridedBoundaryAbortsAFailedRead() {
+ AtomicInteger abortCalls = new AtomicInteger();
+ StorageObject failing = new StorageObject() {
+ @Override
+ public InputStream newStream() {
+ return newStream(0, 1);
+ }
+
+ @Override
+ public InputStream newStream(long position, long length) {
+ return new InputStream() {
+ @Override
+ public int read() throws IOException {
+ throw new IOException("connection reset");
+ }
+ };
+ }
+
+ @Override
+ public void abortStream(InputStream stream) throws IOException {
+ abortCalls.incrementAndGet();
+ stream.close();
+ }
+
+ @Override
+ public long length() {
+ return 4096;
+ }
+
+ @Override
+ public Instant lastModified() {
+ return Instant.EPOCH;
+ }
+
+ @Override
+ public boolean exists() {
+ return true;
+ }
+
+ @Override
+ public StoragePath path() {
+ return StoragePath.of("mem://failing.ndjson");
+ }
+ };
+
+ expectThrows(
+ IOException.class,
+ () -> FileSplitProvider.probeStridedBoundary(stridedSplitter(), failing, 1024, 4096, 1, () -> false)
+ );
+ assertEquals("a failed probe must abort, not drain", 1, abortCalls.get());
+ }
+
+ /** A probe checks for cancellation before opening its stream, so a cancelled query issues no further reads. */
+ public void testProbeStridedBoundaryFailsFastOnCancellation() {
+ AtomicInteger streamsOpened = new AtomicInteger();
+ byte[] payload = "aaaa\nbbbb\ncccc\n".getBytes(StandardCharsets.UTF_8);
+ StorageObject counting = new StorageObject() {
+ @Override
+ public InputStream newStream() {
+ streamsOpened.incrementAndGet();
+ return new ByteArrayInputStream(payload);
+ }
+
+ @Override
+ public InputStream newStream(long position, long length) {
+ streamsOpened.incrementAndGet();
+ return new ByteArrayInputStream(payload, (int) position, (int) length);
+ }
+
+ @Override
+ public long length() {
+ return payload.length;
+ }
+
+ @Override
+ public Instant lastModified() {
+ return Instant.EPOCH;
+ }
+
+ @Override
+ public boolean exists() {
+ return true;
+ }
+
+ @Override
+ public StoragePath path() {
+ return StoragePath.of("mem://cancelled.ndjson");
+ }
+ };
+
+ expectThrows(
+ TaskCancelledException.class,
+ () -> FileSplitProvider.probeStridedBoundary(stridedSplitter(), counting, 5, payload.length, 1, () -> true)
+ );
+ assertEquals("a cancelled probe must not read", 0, streamsOpened.get());
+ }
+
+ /**
+ * A record longer than the splitter's maximum cannot be measured by a probe, so the offset inside it yields no
+ * boundary and the split that began before it runs on past the record. The offsets beyond the record are
+ * probed as usual, so a record the splitter refuses to span costs the one split its offset would have started.
+ */
+ public void testMacroSplitDiscoverySkipsAnOffsetInsideAnOversizedRecord() throws IOException {
var blockFactory = BlockFactory.builder(BigArrays.NON_RECYCLING_INSTANCE).breaker(new NoopCircuitBreaker("test")).build();
- StringBuilder csv = new StringBuilder("ok\n").append("x".repeat(128)).append('\n');
- while (csv.length() < 2 * 1024 * 1024) {
+ int maxRecordBytes = 16;
+ long stride = 256 * 1024;
+ // A record straddling the first stride offset, long enough that the probe there runs past maxRecordBytes
+ // before reaching its terminator. The filler rows divide into the offset exactly, so it starts where meant.
+ long oversizedStart = stride - 64;
+ StringBuilder csv = new StringBuilder();
+ while (csv.length() < oversizedStart) {
+ csv.append("tail\n");
+ }
+ assertEquals("the filler must end exactly where the oversized record starts", oversizedStart, csv.length());
+ csv.append("x".repeat(127)).append('\n');
+ long oversizedEnd = csv.length();
+ while (csv.length() < 2 * CSV_MIN_SEGMENT_BYTES) {
csv.append("tail\n");
}
byte[] payload = csv.toString().getBytes(StandardCharsets.UTF_8);
StorageObject object = createInMemoryStorageObject(payload, StoragePath.of("mem://test.csv"));
- // Plain mode: max-record-size stop is format-agnostic; macro-split discovery now refuses non-strided
+ // Plain mode: the max-record-size verdict is format-agnostic, but macro-split discovery refuses non-strided
// (default/quoted) CSV. Plain CSV keeps strided probing.
var csvReader = (SegmentableFormatReader) new CsvFormatReader(blockFactory).withConfig(Map.of("mode", "plain"));
- List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts(csvReader, object, payload.length, 4, 16, () -> false);
+ List starts = serialStridedStarts(csvReader, object, payload.length, stride, maxRecordBytes);
- assertEquals(List.of(0L), starts);
+ // A maximum the record does fit in leaves every offset able to start a split, which is the count to
+ // measure the oversized record's cost against.
+ List unrestricted = serialStridedStarts(
+ csvReader,
+ object,
+ payload.length,
+ stride,
+ SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES
+ );
+ assertThat("the payload must offer several offsets to probe", unrestricted.size(), greaterThan(2));
+ assertEquals("only the offset inside the oversized record loses its split", unrestricted.size() - 1, starts.size());
+ for (long start : starts) {
+ assertFalse(
+ "no split may start inside a record the splitter cannot span: " + start,
+ start > oversizedStart && start < oversizedEnd
+ );
+ }
+ }
+
+ /**
+ * The serial form of the strided walk a node without a discovery executor falls back to: the fixed stride
+ * offsets of {@code reader}'s splitter, probed on the calling thread and reduced the same way the concurrent
+ * gather reduces them.
+ */
+ private static List serialStridedStarts(
+ SegmentableFormatReader reader,
+ StorageObject object,
+ long fileLength,
+ long targetStrideBytes,
+ int maxRecordBytes
+ ) throws IOException {
+ long minSegment = reader.minimumSegmentSize();
+ return FileSplitProvider.probeStridedBoundariesSerially(
+ reader.recordSplitter(maxRecordBytes),
+ object,
+ fileLength,
+ FileSplitProvider.macroSplitProbePositions(fileLength, targetStrideBytes, minSegment),
+ minSegment,
+ () -> false
+ );
}
private static StorageObject createInMemoryStorageObject(byte[] data, StoragePath 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..80d8ce8797809 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,14 @@ 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#probeStridedBoundary} through the same
+ * decorator chain used for uncompressed text files on object storage. A boundary probe deliberately does
+ * not abort: it opens a bounded window, then drains and closes it so the connection returns to the
+ * pool for the next probe (aborting would drop the connection and cost a handshake per probe). What the chain
+ * must preserve is the window's bound, so the total read stays a small fraction of the file rather than a
+ * range opened to end-of-file. A decorator that ignored the requested length would trip the upper bound here.
*/
- public void testMacroSplitDiscoveryAbortPropagatesThroughDecoratorChainWithoutDrain() throws IOException {
+ public void testMacroSplitDiscoveryDrainsBoundedWindowsThroughDecoratorChain() 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');
@@ -113,25 +115,30 @@ public void testMacroSplitDiscoveryAbortPropagatesThroughDecoratorChainWithoutDr
SegmentableFormatReader csvReader = (SegmentableFormatReader) new CsvFormatReader(blockFactory).withConfig(Map.of("mode", "plain"));
long stride = fileLength / 4;
- List starts = FileSplitProvider.computeRecordAlignedMacroSplitStarts(
- csvReader,
+ long minSegment = csvReader.minimumSegmentSize();
+ List starts = FileSplitProvider.probeStridedBoundariesSerially(
+ csvReader.recordSplitter(SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES),
chain,
fileLength,
- stride,
- SegmentableFormatReader.DEFAULT_MAX_RECORD_BYTES,
+ FileSplitProvider.macroSplitProbePositions(fileLength, stride, minSegment),
+ minSegment,
() -> false
);
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("a boundary probe pools its connection by draining, never aborts", 0, tracking.abortCalls.get());
+ assertTrue("probe streams must be closed through the chain", tracking.closed.get());
assertThat(
- "macro-split probes must not drain range streams; consumed "
+ "each probe drains its full window through the chain",
+ tracking.bytesConsumed.get(),
+ Matchers.greaterThanOrEqualTo((starts.size() - 1) * FileSplitProvider.MACRO_SPLIT_PROBE_WINDOW_BYTES)
+ );
+ assertThat(
+ "macro-split probes must read bounded windows, not ranges to end-of-file; consumed "
+ tracking.bytesConsumed.get()
+ " of "
+ fileLength
- + " bytes across "
- + tracking.abortCalls.get()
- + " probes",
+ + " bytes",
tracking.bytesConsumed.get(),
Matchers.lessThan(fileLength / 2)
);