diff --git a/docs/changelog/143209.yaml b/docs/changelog/143209.yaml new file mode 100644 index 0000000000000..ce5b1222fb28f --- /dev/null +++ b/docs/changelog/143209.yaml @@ -0,0 +1,5 @@ +area: ES|QL +issues: [] +pr: 143209 +summary: Add data node execution for external sources +type: enhancement diff --git a/server/src/main/resources/transport/definitions/referable/esql_external_splits_in_data_node_request.csv b/server/src/main/resources/transport/definitions/referable/esql_external_splits_in_data_node_request.csv new file mode 100644 index 0000000000000..c3c28f9ea7e44 --- /dev/null +++ b/server/src/main/resources/transport/definitions/referable/esql_external_splits_in_data_node_request.csv @@ -0,0 +1 @@ +9299000 diff --git a/server/src/main/resources/transport/upper_bounds/9.4.csv b/server/src/main/resources/transport/upper_bounds/9.4.csv index af8361f633d5b..ab70347c1a681 100644 --- a/server/src/main/resources/transport/upper_bounds/9.4.csv +++ b/server/src/main/resources/transport/upper_bounds/9.4.csv @@ -1 +1 @@ -cluster_stats_multiple_synonym_graph_filters,9298000 +esql_external_splits_in_data_node_request,9299000 diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java index 12f2c72f47b7c..02ad79974aee5 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java @@ -34,6 +34,7 @@ import org.elasticsearch.xpack.esql.core.type.DataType; import org.elasticsearch.xpack.esql.core.util.Holder; import org.elasticsearch.xpack.esql.core.util.Queries; +import org.elasticsearch.xpack.esql.datasources.FilterPushdownRegistry; import org.elasticsearch.xpack.esql.expression.predicate.Predicates; import org.elasticsearch.xpack.esql.io.stream.PlanStreamWrapperQueryBuilder; import org.elasticsearch.xpack.esql.optimizer.LocalLogicalOptimizerContext; @@ -219,6 +220,24 @@ public static PhysicalPlan localPlan( return localPlan(plan, logicalOptimizer, physicalOptimizer, planTimeProfile); } + public static PhysicalPlan localPlan( + PlannerSettings plannerSettings, + EsqlFlags flags, + Configuration configuration, + FoldContext foldCtx, + PhysicalPlan plan, + SearchStats searchStats, + FilterPushdownRegistry filterPushdownRegistry, + PlanTimeProfile planTimeProfile + ) { + final var logicalOptimizer = new LocalLogicalPlanOptimizer(new LocalLogicalOptimizerContext(configuration, foldCtx, searchStats)); + var physicalOptimizer = new LocalPhysicalPlanOptimizer( + new LocalPhysicalOptimizerContext(plannerSettings, flags, configuration, foldCtx, searchStats, filterPushdownRegistry) + ); + + return localPlan(plan, logicalOptimizer, physicalOptimizer, planTimeProfile); + } + public static PhysicalPlan integrateEsFilterIntoFragment(PhysicalPlan plan, @Nullable QueryBuilder esFilter) { return esFilter == null ? plan : plan.transformUp(FragmentExec.class, f -> { var fragmentFilter = f.esFilter(); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java index da7194703ca4c..35709705d8dc5 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java @@ -54,6 +54,7 @@ import org.elasticsearch.xpack.esql.core.expression.Attribute; import org.elasticsearch.xpack.esql.core.expression.FoldContext; import org.elasticsearch.xpack.esql.core.util.Holder; +import org.elasticsearch.xpack.esql.datasources.FilterPushdownRegistry; import org.elasticsearch.xpack.esql.datasources.OperatorFactoryRegistry; import org.elasticsearch.xpack.esql.datasources.SplitDiscoveryPhase; import org.elasticsearch.xpack.esql.datasources.spi.ExternalSplit; @@ -80,6 +81,7 @@ import org.elasticsearch.xpack.esql.session.Configuration; import org.elasticsearch.xpack.esql.session.EsqlCCSUtils; import org.elasticsearch.xpack.esql.session.Result; +import org.elasticsearch.xpack.esql.stats.SearchContextStats; import java.util.ArrayList; import java.util.Collections; @@ -161,6 +163,7 @@ public class ComputeService { private final ExchangeService exchangeService; private final PlannerSettings.Holder plannerSettings; private final OperatorFactoryRegistry operatorFactoryRegistry; + private final FilterPushdownRegistry filterPushdownRegistry; @SuppressWarnings("this-escape") public ComputeService( @@ -170,7 +173,8 @@ public ComputeService( ThreadPool threadPool, BigArrays bigArrays, BlockFactory blockFactory, - OperatorFactoryRegistry operatorFactoryRegistry + OperatorFactoryRegistry operatorFactoryRegistry, + FilterPushdownRegistry filterPushdownRegistry ) { this.searchService = transportActionServices.searchService(); this.transportService = transportActionServices.transportService(); @@ -202,6 +206,7 @@ public ComputeService( ); this.plannerSettings = transportActionServices.plannerSettings(); this.operatorFactoryRegistry = operatorFactoryRegistry; + this.filterPushdownRegistry = filterPushdownRegistry != null ? filterPushdownRegistry : FilterPushdownRegistry.empty(); } PlannerSettings.Holder plannerSettings() { @@ -228,10 +233,10 @@ static ExternalDistributionStrategy resolveExternalDistributionStrategy(QueryPra }; } - PhysicalPlan applyExternalDistributionStrategy(PhysicalPlan plan, Configuration configuration) { + ExternalDistributionResult applyExternalDistributionStrategy(PhysicalPlan plan, Configuration configuration) { List externalSplits = collectExternalSplits(plan); if (externalSplits.isEmpty()) { - return plan; + return new ExternalDistributionResult(plan, null); } ExternalDistributionStrategy strategy = resolveExternalDistributionStrategy(configuration.pragmas()); @@ -245,14 +250,21 @@ PhysicalPlan applyExternalDistributionStrategy(PhysicalPlan plan, Configuration ExternalDistributionPlan distributionPlan = strategy.planDistribution(context); if (distributionPlan.distributed()) { - // TODO: PR5 will insert ExchangeExec and break the plan for data node dispatch. - // Until then, external sources always execute on the coordinator. - LOGGER.debug("external distribution requested but not yet supported; falling back to coordinator-only"); + LOGGER.debug( + "external distribution: distributing {} splits across {} nodes", + externalSplits.size(), + distributionPlan.nodeAssignments().size() + ); + return new ExternalDistributionResult(plan, distributionPlan); } - // Safety net: remove any ExchangeExec wrapping ExternalSourceExec to prevent - // the plan from being incorrectly split between coordinator and data nodes. - return collapseExternalSourceExchanges(plan); + return new ExternalDistributionResult(collapseExternalSourceExchanges(plan), null); + } + + record ExternalDistributionResult(PhysicalPlan plan, ExternalDistributionPlan distributionPlan) { + boolean isDistributed() { + return distributionPlan != null && distributionPlan.distributed(); + } } private static List collectExternalSplits(PhysicalPlan plan) { @@ -420,7 +432,8 @@ public void executePlan( PlanTimeProfile planTimeProfile ) { final PhysicalPlan splitPlan = discoverSplits(physicalPlan); - final PhysicalPlan resolvedPlan = applyExternalDistributionStrategy(splitPlan, configuration); + final ExternalDistributionResult distributionResult = applyExternalDistributionStrategy(splitPlan, configuration); + final PhysicalPlan resolvedPlan = distributionResult.plan(); Tuple coordinatorAndDataNodePlan = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode( resolvedPlan, configuration @@ -444,8 +457,15 @@ public void executePlan( } Map clusterToConcreteIndices = getIndices(resolvedPlan, EsRelation::concreteIndices); Runnable cancelQueryOnFailure = cancelQueryOnFailure(rootTask); + boolean hasConcreteIndices = false; + for (OriginalIndices v : clusterToConcreteIndices.values()) { + if (v.indices().length > 0) { + hasConcreteIndices = true; + break; + } + } if (dataNodePlan == null) { - if (clusterToConcreteIndices.values().stream().allMatch(v -> v.indices().length == 0) == false) { + if (hasConcreteIndices) { String error = "expected no concrete indices without data node plan; got " + clusterToConcreteIndices; assert false : error; listener.onFailure(new IllegalStateException(error)); @@ -484,13 +504,34 @@ public void executePlan( ); return; } - } else { - if (clusterToConcreteIndices.values().stream().allMatch(v -> v.indices().length == 0)) { - var error = "expected concrete indices with data node plan but got empty; data node plan " + dataNodePlan; - assert false : error; - listener.onFailure(new IllegalStateException(error)); - return; - } + } + // External source distribution: data node plan exists but no ES indices + if (distributionResult.isDistributed() && hasConcreteIndices == false) { + executeExternalDistribution( + sessionId, + rootTask, + flags, + configuration, + foldContext, + (ExchangeSinkExec) dataNodePlan, + coordinatorPlan, + resolvedPlan, + distributionResult.distributionPlan(), + collectedPages, + execInfo, + profileQualifier, + cancelQueryOnFailure, + exchangeSinkSupplier, + planTimeProfile, + listener + ); + return; + } + if (hasConcreteIndices == false) { + var error = "expected concrete indices with data node plan but got empty; data node plan " + dataNodePlan; + assert false : error; + listener.onFailure(new IllegalStateException(error)); + return; } Map clusterToOriginalIndices = getIndices(resolvedPlan, EsRelation::originalIndices); var localOriginalIndices = clusterToOriginalIndices.remove(LOCAL_CLUSTER); @@ -659,6 +700,78 @@ public void executePlan( } } + private void executeExternalDistribution( + String sessionId, + CancellableTask rootTask, + EsqlFlags flags, + Configuration configuration, + FoldContext foldContext, + ExchangeSinkExec dataNodePlan, + PhysicalPlan coordinatorPlan, + PhysicalPlan resolvedPlan, + ExternalDistributionPlan distributionPlan, + List collectedPages, + EsqlExecutionInfo execInfo, + String profileQualifier, + Runnable cancelQueryOnFailure, + Supplier exchangeSinkSupplier, + PlanTimeProfile planTimeProfile, + ActionListener listener + ) { + List outputAttributes = resolvedPlan.output(); + var exchangeSource = new ExchangeSourceHandler( + configuration.pragmas().exchangeBufferSize(), + transportService.getThreadPool().executor(ThreadPool.Names.SEARCH) + ); + listener = ActionListener.runBefore(listener, () -> exchangeService.removeExchangeSourceHandler(sessionId)); + exchangeService.addExchangeSourceHandler(sessionId, exchangeSource); + try ( + var computeListener = new ComputeListener( + transportService.getThreadPool(), + cancelQueryOnFailure, + listener.delegateFailureAndWrap((l, completionInfo) -> { + execInfo.markEndQuery(); + l.onResponse(new Result(outputAttributes, collectedPages, configuration, completionInfo, execInfo)); + }) + ) + ) { + try (Releasable ignored = exchangeSource.addEmptySink()) { + // Run the coordinator plan + runCompute( + rootTask, + new ComputeContext( + sessionId, + profileDescription(profileQualifier, "final"), + LOCAL_CLUSTER, + flags, + EmptyIndexedByShardId.instance(), + configuration, + foldContext, + exchangeSource::createExchangeSource, + exchangeSinkSupplier + ), + coordinatorPlan, + plannerSettings.get(), + LocalPhysicalOptimization.ENABLED, + planTimeProfile, + computeListener.acquireCompute() + ); + // Dispatch to each data node with its assigned splits + dataNodeComputeHandler.startExternalComputeOnDataNodes( + sessionId, + rootTask, + flags, + configuration, + dataNodePlan, + distributionPlan, + exchangeSource, + cancelQueryOnFailure, + computeListener + ); + } + } + } + // For queries like: FROM logs* | LIMIT 0 (including cross-cluster LIMIT 0 queries) private static void updateShardCountForCoordinatorOnlyQuery(EsqlExecutionInfo execInfo) { if (execInfo.isCrossClusterSearch() || execInfo.includeExecutionMetadata() == ALWAYS) { @@ -700,8 +813,10 @@ private static void updateExecutionInfoAfterCoordinatorOnlyQuery(EsqlExecutionIn */ static void failIfAllShardsFailed(EsqlExecutionInfo execInfo, List finalResults) { // do not fail if any final result has results - if (finalResults.stream().anyMatch(p -> p.getPositionCount() > 0)) { - return; + for (Page p : finalResults) { + if (p.getPositionCount() > 0) { + return; + } } int totalFailedShards = 0; for (EsqlExecutionInfo.Cluster cluster : execInfo.clusterInfo.values()) { @@ -773,16 +888,28 @@ void runCompute( List localContexts = new ArrayList<>(); context.searchExecutionContexts().iterable().forEach(localContexts::add); + boolean hasExternalSource = plan.anyMatch(p -> p instanceof ExternalSourceExec); var localPlan = switch (localPhysicalOptimization) { - case ENABLED -> PlannerUtils.localPlan( - plannerSettings, - context.flags(), - localContexts, - context.configuration(), - context.foldCtx(), - plan, - planTimeProfile - ); + case ENABLED -> hasExternalSource + ? PlannerUtils.localPlan( + plannerSettings, + context.flags(), + context.configuration(), + context.foldCtx(), + plan, + SearchContextStats.from(localContexts), + filterPushdownRegistry, + planTimeProfile + ) + : PlannerUtils.localPlan( + plannerSettings, + context.flags(), + localContexts, + context.configuration(), + context.foldCtx(), + plan, + planTimeProfile + ); case DISABLED -> plan; }; if (LOGGER.isDebugEnabled()) { diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java index e96862099feb1..5af877ae71f6e 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java @@ -13,11 +13,13 @@ import org.elasticsearch.action.ActionRunnable; import org.elasticsearch.action.OriginalIndices; import org.elasticsearch.action.support.ChannelActionListener; +import org.elasticsearch.action.support.IndicesOptions; import org.elasticsearch.action.support.RefCountingRunnable; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.cluster.project.ProjectResolver; import org.elasticsearch.cluster.routing.SplitShardCountSummary; import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.compute.lucene.EmptyIndexedByShardId; import org.elasticsearch.compute.lucene.IndexedByShardId; import org.elasticsearch.compute.lucene.IndexedByShardIdFromSingleton; import org.elasticsearch.compute.operator.DriverCompletionInfo; @@ -49,7 +51,9 @@ import org.elasticsearch.transport.TransportService; import org.elasticsearch.xpack.esql.action.EsqlCapabilities; import org.elasticsearch.xpack.esql.core.expression.FoldContext; +import org.elasticsearch.xpack.esql.datasources.spi.ExternalSplit; import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec; +import org.elasticsearch.xpack.esql.plan.physical.ExternalSourceExec; import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan; import org.elasticsearch.xpack.esql.planner.PlanConcurrencyCalculator; import org.elasticsearch.xpack.esql.planner.PlannerSettings; @@ -240,6 +244,93 @@ protected void sendRequest( ); } + void startExternalComputeOnDataNodes( + String sessionId, + CancellableTask parentTask, + EsqlFlags flags, + Configuration configuration, + ExchangeSinkExec dataNodePlan, + ExternalDistributionPlan distributionPlan, + ExchangeSourceHandler exchangeSource, + Runnable runOnTaskFailure, + ComputeListener parentComputeListener + ) { + var queryPragmas = configuration.pragmas(); + boolean sentAny = false; + for (Map.Entry> entry : distributionPlan.nodeAssignments().entrySet()) { + String nodeId = entry.getKey(); + List nodeSplits = entry.getValue(); + if (nodeSplits.isEmpty()) { + continue; + } + + DiscoveryNode node = clusterService.state().nodes().get(nodeId); + if (node == null) { + parentComputeListener.acquireCompute() + .onFailure(new IllegalStateException("node [" + nodeId + "] not found in cluster state")); + return; + } + + final Transport.Connection connection; + try { + connection = transportService.getConnection(node); + } catch (Exception e) { + parentComputeListener.acquireCompute().onFailure(e); + return; + } + + sentAny = true; + var childSessionId = computeService.newChildSession(sessionId); + ExchangeService.openExchange( + transportService, + connection, + childSessionId, + queryPragmas.exchangeBufferSize(), + esqlExecutor, + parentComputeListener.acquireAvoid().delegateFailureAndWrap((l, unused) -> { + var computeListener = parentComputeListener; + var dataNodeRequest = new DataNodeRequest( + childSessionId, + configuration, + "", + List.of(), + Map.of(), + dataNodePlan, + new String[0], + IndicesOptions.STRICT_EXPAND_OPEN, + queryPragmas.nodeLevelReduction(), + false, + nodeSplits + ); + transportService.sendChildRequest( + connection, + ComputeService.DATA_ACTION_NAME, + dataNodeRequest, + parentTask, + TransportRequestOptions.EMPTY, + new ActionListenerResponseHandler<>( + computeListener.acquireCompute().map(r -> r.completionInfo()), + DataNodeComputeResponse::new, + esqlExecutor + ) + ); + var remoteSink = exchangeService.newRemoteSink(parentTask, childSessionId, transportService, connection); + exchangeSource.addRemoteSink( + remoteSink, + configuration.allowPartialResults() == false, + () -> {}, + queryPragmas.concurrentExchangeClients(), + computeListener.acquireAvoid() + ); + l.onResponse(null); + }) + ); + } + if (sentAny == false) { + parentComputeListener.acquireCompute().onResponse(DriverCompletionInfo.EMPTY); + } + } + private static final Logger LOGGER = LogManager.getLogger(DataNodeComputeHandler.class); private class DataNodeRequestExecutor { @@ -572,14 +663,18 @@ private void runComputeOnDataNode( @Override public void messageReceived(DataNodeRequest request, TransportChannel channel, Task task) { ActionListener listener = new ChannelActionListener<>(channel); - ReductionPlan reductionPlan; Configuration configuration = request.configuration(); - // We can avoid synchronization (for the most part) since the array elements are never modified, and the array is only added to, - // with its size being known before we start the computation. PlanTimeProfile planTimeProfile = null; if (configuration.profile()) { planTimeProfile = new PlanTimeProfile(); } + + if (request.externalSplits().isEmpty() == false && request.shards().isEmpty()) { + handleExternalSourceRequest(request, (CancellableTask) task, listener, planTimeProfile); + return; + } + + ReductionPlan reductionPlan; if (request.plan() instanceof ExchangeSinkExec plan) { reductionPlan = ComputeService.reductionPlan( computeService.plannerSettings().get(), @@ -606,7 +701,8 @@ public void messageReceived(DataNodeRequest request, TransportChannel channel, T request.indices(), request.indicesOptions(), request.runNodeLevelReduction(), - request.reductionLateMaterialization() + request.reductionLateMaterialization(), + request.externalSplits() ); // the sender doesn't support retry on shard failures, so we need to fail fast here. final boolean failFastOnShardFailures = supportShardLevelRetryFailure(channel.getVersion()) == false; @@ -625,6 +721,83 @@ public void messageReceived(DataNodeRequest request, TransportChannel channel, T ); } + private void handleExternalSourceRequest( + DataNodeRequest request, + CancellableTask task, + ActionListener listener, + PlanTimeProfile planTimeProfile + ) { + if (request.plan() instanceof ExchangeSinkExec == false) { + listener.onFailure(new IllegalStateException("expected exchange sink for external compute; got " + request.plan())); + return; + } + ExchangeSinkExec sinkExec = (ExchangeSinkExec) request.plan(); + Configuration configuration = request.configuration(); + final String sessionId = request.sessionId(); + + // Inject external splits into the ExternalSourceExec within the plan + PhysicalPlan planWithSplits = sinkExec.child() + .transformUp(ExternalSourceExec.class, exec -> exec.withSplits(request.externalSplits())); + ExchangeSinkExec updatedSinkExec = new ExchangeSinkExec( + sinkExec.source(), + sinkExec.output(), + sinkExec.isIntermediateAgg(), + planWithSplits + ); + + try ( + ComputeListener computeListener = new ComputeListener( + threadPool, + computeService.cancelQueryOnFailure(task), + listener.map(profiles -> new DataNodeComputeResponse(profiles, Map.of())) + ) + ) { + var parentListener = computeListener.acquireAvoid(); + try { + var externalSink = exchangeService.getSinkHandler(sessionId); + String internalSessionId = sessionId + "[n]"; + task.addListener( + () -> { exchangeService.finishSinkHandler(sessionId, new TaskCancelledException(task.getReasonCancelled())); } + ); + EsqlFlags flags = computeService.createFlags(); + + var computeContext = new ComputeContext( + internalSessionId, + ComputeService.DATA_DESCRIPTION, + request.clusterAlias(), + flags, + EmptyIndexedByShardId.instance(), + configuration, + configuration.newFoldContext(), + null, + () -> externalSink.createExchangeSink(() -> {}) + ); + computeService.runCompute( + task, + computeContext, + updatedSinkExec, + computeService.plannerSettings().get(), + LocalPhysicalOptimization.ENABLED, + planTimeProfile, + ActionListener.wrap(resp -> { + externalSink.addCompletionListener(ActionListener.running(() -> { + exchangeService.finishSinkHandler(sessionId, null); + computeListener.acquireCompute().onResponse(resp); + })); + }, e -> { + LOGGER.debug("Error in external source compute on data node", e); + exchangeService.finishSinkHandler(sessionId, e); + computeListener.acquireCompute().onFailure(e); + }) + ); + parentListener.onResponse(null); + } catch (Exception e) { + exchangeService.finishSinkHandler(sessionId, e); + parentListener.onFailure(e); + } + } + } + static boolean supportShardLevelRetryFailure(TransportVersion transportVersion) { return transportVersion.supports(ESQL_RETRY_ON_SHARD_LEVEL_FAILURE); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequest.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequest.java index f757c6a838858..ac283419c94d3 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequest.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequest.java @@ -29,6 +29,7 @@ import org.elasticsearch.tasks.Task; import org.elasticsearch.tasks.TaskId; import org.elasticsearch.transport.AbstractTransportRequest; +import org.elasticsearch.xpack.esql.datasources.spi.ExternalSplit; import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput; import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput; import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec; @@ -48,6 +49,9 @@ final class DataNodeRequest extends AbstractTransportRequest implements IndicesRequest.Replaceable { private static final TransportVersion REDUCE_LATE_MATERIALIZATION = TransportVersion.fromName("esql_reduce_late_materialization"); + private static final TransportVersion EXTERNAL_SPLITS_IN_DATA_NODE_REQUEST = TransportVersion.fromName( + "esql_external_splits_in_data_node_request" + ); private static final Logger logger = LogManager.getLogger(DataNodeRequest.class); @@ -61,6 +65,7 @@ final class DataNodeRequest extends AbstractTransportRequest implements IndicesR private final IndicesOptions indicesOptions; private final boolean runNodeLevelReduction; private final boolean reductionLateMaterialization; + private final List externalSplits; DataNodeRequest( String sessionId, @@ -73,6 +78,34 @@ final class DataNodeRequest extends AbstractTransportRequest implements IndicesR IndicesOptions indicesOptions, boolean runNodeLevelReduction, boolean reductionLateMaterialization + ) { + this( + sessionId, + configuration, + clusterAlias, + shards, + aliasFilters, + plan, + indices, + indicesOptions, + runNodeLevelReduction, + reductionLateMaterialization, + List.of() + ); + } + + DataNodeRequest( + String sessionId, + Configuration configuration, + String clusterAlias, + List shards, + Map aliasFilters, + PhysicalPlan plan, + String[] indices, + IndicesOptions indicesOptions, + boolean runNodeLevelReduction, + boolean reductionLateMaterialization, + List externalSplits ) { this.sessionId = sessionId; this.configuration = configuration; @@ -84,6 +117,7 @@ final class DataNodeRequest extends AbstractTransportRequest implements IndicesR this.indicesOptions = indicesOptions; this.runNodeLevelReduction = runNodeLevelReduction; this.reductionLateMaterialization = reductionLateMaterialization; + this.externalSplits = externalSplits != null ? List.copyOf(externalSplits) : List.of(); } DataNodeRequest(StreamInput in) throws IOException { @@ -118,6 +152,11 @@ final class DataNodeRequest extends AbstractTransportRequest implements IndicesR } else { this.reductionLateMaterialization = false; } + if (in.getTransportVersion().supports(EXTERNAL_SPLITS_IN_DATA_NODE_REQUEST)) { + this.externalSplits = in.readNamedWriteableCollectionAsList(ExternalSplit.class); + } else { + this.externalSplits = List.of(); + } } @Override @@ -139,6 +178,9 @@ public void writeTo(StreamOutput out) throws IOException { if (out.getTransportVersion().supports(REDUCE_LATE_MATERIALIZATION)) { out.writeBoolean(reductionLateMaterialization); } + if (out.getTransportVersion().supports(EXTERNAL_SPLITS_IN_DATA_NODE_REQUEST)) { + out.writeNamedWriteableCollection(externalSplits); + } } @Override @@ -219,9 +261,17 @@ boolean reductionLateMaterialization() { return reductionLateMaterialization; } + List externalSplits() { + return externalSplits; + } + @Override public String getDescription() { - return "shards=" + shards + " plan=" + plan; + String desc = "shards=" + shards + " plan=" + plan; + if (externalSplits.isEmpty() == false) { + desc += " externalSplits=" + externalSplits.size(); + } + return desc; } @Override @@ -243,7 +293,9 @@ public boolean equals(Object o) { && getParentTask().equals(request.getParentTask()) && Arrays.equals(indices, request.indices) && indicesOptions.equals(request.indicesOptions) - && runNodeLevelReduction == request.runNodeLevelReduction; + && runNodeLevelReduction == request.runNodeLevelReduction + && reductionLateMaterialization == request.reductionLateMaterialization + && externalSplits.equals(request.externalSplits); } @Override @@ -257,7 +309,9 @@ public int hashCode() { plan, Arrays.hashCode(indices), indicesOptions, - runNodeLevelReduction + runNodeLevelReduction, + reductionLateMaterialization, + externalSplits ); } @@ -272,7 +326,8 @@ public DataNodeRequest withPlan(ExchangeSinkExec newPlan) { indices, indicesOptions, runNodeLevelReduction, - reductionLateMaterialization + reductionLateMaterialization, + externalSplits ); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java index 15fb73daec89e..ad78377ce245a 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java @@ -52,6 +52,7 @@ import org.elasticsearch.xpack.esql.action.EsqlResponseListener; import org.elasticsearch.xpack.esql.analysis.AnalyzerSettings; import org.elasticsearch.xpack.esql.core.async.AsyncTaskManagementService; +import org.elasticsearch.xpack.esql.datasources.FilterPushdownRegistry; import org.elasticsearch.xpack.esql.datasources.OperatorFactoryRegistry; import org.elasticsearch.xpack.esql.enrich.AbstractLookupService; import org.elasticsearch.xpack.esql.enrich.EnrichLookupService; @@ -199,6 +200,7 @@ public TransportEsqlQueryAction( OperatorFactoryRegistry operatorFactoryRegistry = planExecutor.dataSourceModule() .createOperatorFactoryRegistry(threadPool.executor(ThreadPool.Names.SEARCH)); + FilterPushdownRegistry filterPushdownRegistry = planExecutor.dataSourceModule().filterPushdownRegistry(); this.computeService = new ComputeService( services, enrichLookupService, @@ -206,7 +208,8 @@ public TransportEsqlQueryAction( threadPool, bigArrays, blockFactoryProvider.blockFactory(), - operatorFactoryRegistry + operatorFactoryRegistry, + filterPushdownRegistry ); this.activityLogger = new ActivityLogger<>( diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/ExternalSourceDataNodeTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/ExternalSourceDataNodeTests.java new file mode 100644 index 0000000000000..82aa34472b544 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/ExternalSourceDataNodeTests.java @@ -0,0 +1,494 @@ +/* + * 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.plugin; + +import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.cluster.node.DiscoveryNodeUtils; +import org.elasticsearch.cluster.node.DiscoveryNodes; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.compute.aggregation.AggregatorMode; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.esql.core.expression.Attribute; +import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; +import org.elasticsearch.xpack.esql.core.expression.Literal; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.core.type.EsField; +import org.elasticsearch.xpack.esql.datasources.FileSplit; +import org.elasticsearch.xpack.esql.datasources.spi.ExternalSplit; +import org.elasticsearch.xpack.esql.datasources.spi.StoragePath; +import org.elasticsearch.xpack.esql.plan.physical.AggregateExec; +import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec; +import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec; +import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec; +import org.elasticsearch.xpack.esql.plan.physical.ExternalSourceExec; +import org.elasticsearch.xpack.esql.plan.physical.LimitExec; +import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan; +import org.elasticsearch.xpack.esql.planner.PlannerUtils; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.elasticsearch.cluster.node.DiscoveryNodeRole.DATA_HOT_NODE_ROLE; +import static org.hamcrest.Matchers.equalTo; + +/** + * Tests for data node execution of external sources. + * Covers distribution planning, plan structure, and request building. + */ +public class ExternalSourceDataNodeTests extends ESTestCase { + + private static final Source SRC = Source.EMPTY; + + // --- Distribution result tests --- + + public void testDistributionResultNotDistributedWhenNoSplits() { + ExternalSourceExec source = createExternalSourceExec(); + LimitExec limit = new LimitExec(SRC, source, new Literal(SRC, 10, DataType.INTEGER), null); + + List splits = collectExternalSplits(limit); + assertTrue("No splits expected on source without splits", splits.isEmpty()); + } + + public void testDistributionResultDistributedWithSplitsAndAggregation() { + List splits = createSplits(5); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + AggregateExec agg = new AggregateExec(SRC, source, List.of(), List.of(), AggregatorMode.INITIAL, List.of(), null); + ExchangeExec exchange = new ExchangeExec(SRC, agg); + + AdaptiveStrategy strategy = new AdaptiveStrategy(); + DiscoveryNodes nodes = createNodes(3); + ExternalDistributionContext context = new ExternalDistributionContext(exchange, splits, nodes, QueryPragmas.EMPTY); + ExternalDistributionPlan plan = strategy.planDistribution(context); + + assertTrue("Should distribute with aggregation and multiple splits", plan.distributed()); + assertFalse("Node assignments should not be empty", plan.nodeAssignments().isEmpty()); + } + + public void testDistributionResultLocalWithSingleSplit() { + List splits = createSplits(1); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + + AdaptiveStrategy strategy = new AdaptiveStrategy(); + DiscoveryNodes nodes = createNodes(3); + ExternalDistributionContext context = new ExternalDistributionContext(source, splits, nodes, QueryPragmas.EMPTY); + ExternalDistributionPlan plan = strategy.planDistribution(context); + + assertFalse("Single split should stay local", plan.distributed()); + } + + public void testCoordinatorOnlyPragmaOverride() { + List splits = createSplits(10); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + AggregateExec agg = new AggregateExec(SRC, source, List.of(), List.of(), AggregatorMode.INITIAL, List.of(), null); + + Settings settings = Settings.builder().put("external_distribution", "coordinator_only").build(); + QueryPragmas pragmas = new QueryPragmas(settings); + ExternalDistributionStrategy strategy = ComputeService.resolveExternalDistributionStrategy(pragmas); + assertTrue("Should be coordinator-only strategy", strategy instanceof CoordinatorOnlyStrategy); + + DiscoveryNodes nodes = createNodes(3); + ExternalDistributionContext context = new ExternalDistributionContext(agg, splits, nodes, pragmas); + ExternalDistributionPlan plan = strategy.planDistribution(context); + + assertFalse("coordinator_only pragma should prevent distribution", plan.distributed()); + } + + public void testRoundRobinPragmaDistribution() { + List splits = createSplits(6); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + + Settings settings = Settings.builder().put("external_distribution", "round_robin").build(); + QueryPragmas pragmas = new QueryPragmas(settings); + ExternalDistributionStrategy strategy = ComputeService.resolveExternalDistributionStrategy(pragmas); + assertTrue("Should be round-robin strategy", strategy instanceof RoundRobinStrategy); + + DiscoveryNodes nodes = createNodes(3); + ExternalDistributionContext context = new ExternalDistributionContext(source, splits, nodes, pragmas); + ExternalDistributionPlan plan = strategy.planDistribution(context); + + assertTrue("round_robin should distribute", plan.distributed()); + assertEquals("Should have 3 node assignments", 3, plan.nodeAssignments().size()); + for (List nodeSplits : plan.nodeAssignments().values()) { + assertEquals("Each node should get 2 splits", 2, nodeSplits.size()); + } + } + + // --- Plan structure tests --- + + public void testExternalSourcePlanBreaksAtExchange() { + List splits = createSplits(3); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + ExchangeExec exchange = new ExchangeExec(SRC, source); + LimitExec limit = new LimitExec(SRC, exchange, new Literal(SRC, 10, DataType.INTEGER), null); + + var result = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(limit, null); + + PhysicalPlan coordinatorPlan = result.v1(); + PhysicalPlan dataNodePlan = result.v2(); + + assertNotNull("Data node plan should exist", dataNodePlan); + assertTrue("Data node plan should be ExchangeSinkExec", dataNodePlan instanceof ExchangeSinkExec); + + ExchangeSinkExec sink = (ExchangeSinkExec) dataNodePlan; + assertTrue("Sink child should be ExternalSourceExec", sink.child() instanceof ExternalSourceExec); + + ExternalSourceExec sinkSource = (ExternalSourceExec) sink.child(); + assertEquals("Splits should be preserved", 3, sinkSource.splits().size()); + + assertTrue("Coordinator plan should have ExchangeSourceExec", coordinatorPlan instanceof LimitExec); + LimitExec coordLimit = (LimitExec) coordinatorPlan; + assertTrue("Coordinator child should be ExchangeSourceExec", coordLimit.child() instanceof ExchangeSourceExec); + } + + public void testExternalSourceExecWithSplitsInjection() { + ExternalSourceExec source = createExternalSourceExec(); + assertTrue("Initially no splits", source.splits().isEmpty()); + + List splits = createSplits(4); + ExternalSourceExec withSplits = source.withSplits(splits); + + assertEquals("Should have 4 splits", 4, withSplits.splits().size()); + assertEquals("Source path preserved", source.sourcePath(), withSplits.sourcePath()); + assertEquals("Source type preserved", source.sourceType(), withSplits.sourceType()); + assertEquals("Attributes preserved", source.output(), withSplits.output()); + } + + // --- Split injection into plan tests --- + + public void testSplitInjectionIntoExchangeSinkPlan() { + ExternalSourceExec source = createExternalSourceExec(); + ExchangeSinkExec sink = new ExchangeSinkExec(SRC, source.output(), false, source); + + List splits = createSplits(4); + PhysicalPlan planWithSplits = sink.child().transformUp(ExternalSourceExec.class, exec -> exec.withSplits(splits)); + + assertTrue("Transformed plan should be ExternalSourceExec", planWithSplits instanceof ExternalSourceExec); + ExternalSourceExec injected = (ExternalSourceExec) planWithSplits; + assertEquals("Should have 4 splits after injection", 4, injected.splits().size()); + } + + public void testSplitInjectionPreservesOtherFields() { + ExternalSourceExec source = createExternalSourceExec(); + List splits = createSplits(2); + ExternalSourceExec withSplits = source.withSplits(splits); + + assertEquals(source.sourcePath(), withSplits.sourcePath()); + assertEquals(source.sourceType(), withSplits.sourceType()); + assertEquals(source.output(), withSplits.output()); + assertEquals(source.config(), withSplits.config()); + assertEquals(source.sourceMetadata(), withSplits.sourceMetadata()); + assertEquals(2, withSplits.splits().size()); + } + + // --- Node assignment tests --- + + public void testRoundRobinDistributionEvenSplits() { + List splits = createSplits(9); + List nodes = new ArrayList<>(); + DiscoveryNodes.Builder builder = DiscoveryNodes.builder(); + for (int i = 0; i < 3; i++) { + DiscoveryNode node = DiscoveryNodeUtils.builder("node-" + i).roles(Set.of(DATA_HOT_NODE_ROLE)).build(); + nodes.add(node); + builder.add(node); + } + + ExternalDistributionPlan plan = RoundRobinStrategy.assignRoundRobin(splits, nodes); + + assertTrue(plan.distributed()); + assertEquals(3, plan.nodeAssignments().size()); + for (List nodeSplits : plan.nodeAssignments().values()) { + assertEquals("Each of 3 nodes should get 3 of 9 splits", 3, nodeSplits.size()); + } + } + + public void testRoundRobinDistributionUnevenSplits() { + List splits = createSplits(7); + List nodes = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + nodes.add(DiscoveryNodeUtils.builder("node-" + i).roles(Set.of(DATA_HOT_NODE_ROLE)).build()); + } + + ExternalDistributionPlan plan = RoundRobinStrategy.assignRoundRobin(splits, nodes); + + assertTrue(plan.distributed()); + int totalAssigned = 0; + for (List nodeSplits : plan.nodeAssignments().values()) { + totalAssigned += nodeSplits.size(); + } + assertEquals("All splits should be assigned", 7, totalAssigned); + } + + // --- ExternalDistributionResult tests --- + + public void testDistributionResultIsDistributedWithRealPlan() { + List splits = createSplits(4); + Map> assignments = Map.of("node-0", splits.subList(0, 2), "node-1", splits.subList(2, 4)); + ExternalDistributionPlan distributionPlan = new ExternalDistributionPlan(assignments, true); + + var result = new ComputeService.ExternalDistributionResult(createExternalSourceExec(), distributionPlan); + assertTrue("Should be distributed when plan says distributed", result.isDistributed()); + assertThat(result.distributionPlan().nodeAssignments().size(), equalTo(2)); + } + + public void testDistributionResultNotDistributedWithNullPlan() { + var result = new ComputeService.ExternalDistributionResult(createExternalSourceExec(), null); + assertFalse("Should not be distributed with null plan", result.isDistributed()); + } + + public void testDistributionResultNotDistributedWithLocalPlan() { + var result = new ComputeService.ExternalDistributionResult(createExternalSourceExec(), ExternalDistributionPlan.LOCAL); + assertFalse("Should not be distributed with LOCAL plan", result.isDistributed()); + } + + // --- Dispatch condition tests --- + + public void testExternalDispatchConditionWithSplitsAndNoShards() { + List splits = createSplits(3); + boolean hasExternalSplits = splits.isEmpty() == false; + boolean hasShardsEmpty = List.of().isEmpty(); + assertTrue("Should dispatch to external handler", hasExternalSplits && hasShardsEmpty); + } + + public void testExternalDispatchConditionWithSplitsAndShards() { + List splits = createSplits(3); + boolean hasExternalSplits = splits.isEmpty() == false; + boolean hasShardsEmpty = false; + assertFalse("Should NOT dispatch to external handler when shards present", hasExternalSplits && hasShardsEmpty); + } + + public void testExternalDispatchConditionWithNoSplits() { + List splits = List.of(); + boolean hasExternalSplits = splits.isEmpty() == false; + assertFalse("Should NOT dispatch to external handler when no splits", hasExternalSplits); + } + + // --- handleExternalSourceRequest failure path tests --- + + public void testHandleExternalSourceRequiresExchangeSinkExec() { + ExternalSourceExec source = createExternalSourceExec().withSplits(createSplits(3)); + PhysicalPlan plan = source; + assertFalse("A non-ExchangeSinkExec plan should be rejected by handleExternalSourceRequest", plan instanceof ExchangeSinkExec); + } + + public void testHandleExternalSourcePlanTransformationWithSinkExec() { + ExternalSourceExec source = createExternalSourceExec(); + ExchangeSinkExec sink = new ExchangeSinkExec(SRC, source.output(), false, source); + + List splits = createSplits(5); + PhysicalPlan planWithSplits = sink.child().transformUp(ExternalSourceExec.class, exec -> exec.withSplits(splits)); + ExchangeSinkExec updatedSink = new ExchangeSinkExec(sink.source(), sink.output(), sink.isIntermediateAgg(), planWithSplits); + + assertTrue("Updated sink child should be ExternalSourceExec", updatedSink.child() instanceof ExternalSourceExec); + ExternalSourceExec injected = (ExternalSourceExec) updatedSink.child(); + assertThat("Should have 5 splits after injection", injected.splits().size(), equalTo(5)); + assertThat("Output should be preserved", updatedSink.output(), equalTo(sink.output())); + } + + public void testHandleExternalSourcePlanTransformationWithNestedPlan() { + ExternalSourceExec source = createExternalSourceExec(); + LimitExec limit = new LimitExec(SRC, source, new Literal(SRC, 100, DataType.INTEGER), null); + ExchangeSinkExec sink = new ExchangeSinkExec(SRC, source.output(), false, limit); + + List splits = createSplits(3); + PhysicalPlan planWithSplits = sink.child().transformUp(ExternalSourceExec.class, exec -> exec.withSplits(splits)); + ExchangeSinkExec updatedSink = new ExchangeSinkExec(sink.source(), sink.output(), sink.isIntermediateAgg(), planWithSplits); + + assertTrue("Updated sink child should be LimitExec", updatedSink.child() instanceof LimitExec); + LimitExec updatedLimit = (LimitExec) updatedSink.child(); + assertTrue("Limit child should be ExternalSourceExec", updatedLimit.child() instanceof ExternalSourceExec); + ExternalSourceExec injected = (ExternalSourceExec) updatedLimit.child(); + assertThat("Should have 3 splits after injection through nested plan", injected.splits().size(), equalTo(3)); + } + + // --- startExternalComputeOnDataNodes failure path tests --- + + public void testStartExternalComputeSkipsEmptySplitAssignments() { + List splits = List.of(); + Map> assignments = Map.of("node-0", splits); + + boolean sentAny = false; + for (Map.Entry> entry : assignments.entrySet()) { + if (entry.getValue().isEmpty()) { + continue; + } + sentAny = true; + } + assertFalse("Should not send when all assignments are empty", sentAny); + } + + public void testStartExternalComputeHandlesMissingNode() { + DiscoveryNodes nodes = createNodes(2); + String missingNodeId = "node-missing"; + assertNull("Missing node should not be found", nodes.get(missingNodeId)); + } + + public void testStartExternalComputeNodeLookupSuccess() { + DiscoveryNodes nodes = createNodes(3); + assertNotNull("node-0 should exist", nodes.get("node-0")); + assertNotNull("node-1 should exist", nodes.get("node-1")); + assertNotNull("node-2 should exist", nodes.get("node-2")); + } + + // --- Full distribution flow tests --- + + public void testExternalDistributionBuildsSplitAssignments() { + List splits = createSplits(10); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + AggregateExec agg = new AggregateExec(SRC, source, List.of(), List.of(), AggregatorMode.INITIAL, List.of(), null); + ExchangeExec exchange = new ExchangeExec(SRC, agg); + + DiscoveryNodes nodes = createNodes(3); + Settings settings = Settings.builder().put("external_distribution", "round_robin").build(); + QueryPragmas pragmas = new QueryPragmas(settings); + ExternalDistributionStrategy strategy = ComputeService.resolveExternalDistributionStrategy(pragmas); + ExternalDistributionContext context = new ExternalDistributionContext(exchange, splits, nodes, pragmas); + ExternalDistributionPlan plan = strategy.planDistribution(context); + + assertTrue("Should distribute", plan.distributed()); + assertThat("Should have 3 node assignments", plan.nodeAssignments().size(), equalTo(3)); + + int totalAssigned = 0; + for (List nodeSplits : plan.nodeAssignments().values()) { + assertTrue("Each node should have at least 1 split", nodeSplits.size() >= 1); + totalAssigned += nodeSplits.size(); + } + assertThat("All 10 splits should be assigned", totalAssigned, equalTo(10)); + } + + public void testExternalDistributionPlanBreaksCorrectlyForDataNodes() { + List splits = createSplits(6); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + ExchangeExec exchange = new ExchangeExec(SRC, source); + LimitExec limit = new LimitExec(SRC, exchange, new Literal(SRC, 50, DataType.INTEGER), null); + + var result = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(limit, null); + PhysicalPlan coordinatorPlan = result.v1(); + PhysicalPlan dataNodePlan = result.v2(); + + assertNotNull("Should have data node plan", dataNodePlan); + assertTrue("Data node plan should be ExchangeSinkExec", dataNodePlan instanceof ExchangeSinkExec); + + ExchangeSinkExec sinkExec = (ExchangeSinkExec) dataNodePlan; + assertTrue("Sink child should be ExternalSourceExec", sinkExec.child() instanceof ExternalSourceExec); + ExternalSourceExec sinkSource = (ExternalSourceExec) sinkExec.child(); + assertThat("Splits should be preserved in data node plan", sinkSource.splits().size(), equalTo(6)); + + // Verify coordinator plan structure + assertTrue("Coordinator plan should be LimitExec", coordinatorPlan instanceof LimitExec); + LimitExec coordLimit = (LimitExec) coordinatorPlan; + assertTrue("Coordinator child should be ExchangeSourceExec", coordLimit.child() instanceof ExchangeSourceExec); + + // Verify the data node plan can be used to inject per-node splits + List nodeSplits = splits.subList(0, 2); + PhysicalPlan injected = sinkExec.child().transformUp(ExternalSourceExec.class, exec -> exec.withSplits(nodeSplits)); + ExchangeSinkExec updatedSink = new ExchangeSinkExec(sinkExec.source(), sinkExec.output(), sinkExec.isIntermediateAgg(), injected); + + ExternalSourceExec updatedSource = (ExternalSourceExec) updatedSink.child(); + assertThat("Injected plan should have 2 splits", updatedSource.splits().size(), equalTo(2)); + } + + public void testExternalDistributionWithAggregationPlanBreaks() { + List splits = createSplits(4); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + List output = source.output(); + AggregateExec agg = new AggregateExec(SRC, source, List.of(), List.of(), AggregatorMode.INITIAL, output, null); + ExchangeExec exchange = new ExchangeExec(SRC, agg); + + var result = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(exchange, null); + PhysicalPlan dataNodePlan = result.v2(); + + assertNotNull("Should have data node plan for aggregation", dataNodePlan); + assertTrue("Data node plan should be ExchangeSinkExec", dataNodePlan instanceof ExchangeSinkExec); + + ExchangeSinkExec sinkExec = (ExchangeSinkExec) dataNodePlan; + assertTrue("Sink child should be AggregateExec", sinkExec.child() instanceof AggregateExec); + AggregateExec sinkAgg = (AggregateExec) sinkExec.child(); + assertTrue("Agg child should be ExternalSourceExec", sinkAgg.child() instanceof ExternalSourceExec); + + ExternalSourceExec sinkSource = (ExternalSourceExec) sinkAgg.child(); + assertThat("Splits preserved through aggregation", sinkSource.splits().size(), equalTo(4)); + + // Verify split injection works through the aggregation + List nodeSplits = splits.subList(0, 1); + PhysicalPlan injected = sinkExec.child().transformUp(ExternalSourceExec.class, exec -> exec.withSplits(nodeSplits)); + ExchangeSinkExec updatedSink = new ExchangeSinkExec(sinkExec.source(), sinkExec.output(), sinkExec.isIntermediateAgg(), injected); + + AggregateExec updatedAgg = (AggregateExec) updatedSink.child(); + ExternalSourceExec updatedSource = (ExternalSourceExec) updatedAgg.child(); + assertThat("Injected through aggregation should have 1 split", updatedSource.splits().size(), equalTo(1)); + } + + public void testCollapseExternalSourceExchangesPreservesSplits() { + List splits = createSplits(3); + ExternalSourceExec source = createExternalSourceExec().withSplits(splits); + ExchangeExec exchange = new ExchangeExec(SRC, source); + LimitExec limit = new LimitExec(SRC, exchange, new Literal(SRC, 10, DataType.INTEGER), null); + + PhysicalPlan collapsed = ComputeService.collapseExternalSourceExchanges(limit); + + assertTrue("Collapsed plan should be LimitExec", collapsed instanceof LimitExec); + LimitExec collapsedLimit = (LimitExec) collapsed; + assertTrue("Exchange should be removed, child should be ExternalSourceExec", collapsedLimit.child() instanceof ExternalSourceExec); + + ExternalSourceExec collapsedSource = (ExternalSourceExec) collapsedLimit.child(); + assertThat("Splits should be preserved after collapse", collapsedSource.splits().size(), equalTo(3)); + } + + public void testCollapseExternalSourceExchangesLeavesNonExternalExchange() { + ExternalSourceExec source = createExternalSourceExec(); + LimitExec limit = new LimitExec(SRC, source, new Literal(SRC, 10, DataType.INTEGER), null); + ExchangeExec exchange = new ExchangeExec(SRC, limit); + + PhysicalPlan collapsed = ComputeService.collapseExternalSourceExchanges(exchange); + + assertTrue("Non-external-child exchange should be preserved", collapsed instanceof ExchangeExec); + } + + // --- Helpers --- + + private static ExternalSourceExec createExternalSourceExec() { + return new ExternalSourceExec( + SRC, + "s3://bucket/data/*.parquet", + "parquet", + List.of( + new FieldAttribute(SRC, "name", new EsField("name", DataType.KEYWORD, Map.of(), false, EsField.TimeSeriesFieldType.NONE)) + ), + Map.of(), + Map.of(), + null + ); + } + + private static List createSplits(int count) { + List splits = new ArrayList<>(); + for (int i = 0; i < count; i++) { + splits.add( + new FileSplit("parquet", StoragePath.of("s3://bucket/file" + i + ".parquet"), 0, 1024, "parquet", Map.of(), Map.of()) + ); + } + return splits; + } + + private static DiscoveryNodes createNodes(int count) { + DiscoveryNodes.Builder builder = DiscoveryNodes.builder(); + for (int i = 0; i < count; i++) { + builder.add(DiscoveryNodeUtils.builder("node-" + i).roles(Set.of(DATA_HOT_NODE_ROLE)).build()); + } + return builder.build(); + } + + private static List collectExternalSplits(PhysicalPlan plan) { + List splits = new ArrayList<>(); + plan.forEachDown(ExternalSourceExec.class, exec -> splits.addAll(exec.splits())); + return splits; + } +}