diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsSettings.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsSettings.java
index 99fca18527586..9cd746c8dba5b 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsSettings.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsSettings.java
@@ -224,6 +224,40 @@ private AnalyticsSettings() {}
Setting.Property.Dynamic
);
+ /**
+ * Compute a sub-plan that the query evaluates MORE THAN ONCE only once, feeding every consumer from that
+ * one result.
+ *
+ *
This is a CORRECTNESS fix before it is an optimization. A query that inlines the same aggregate
+ * subquery twice — TPC-H q15 joins {@code revenue0} and then filters
+ * {@code where total_revenue = [ … max(total_revenue) ]} over the same {@code revenue0}, because the
+ * spec's VIEW has no PPL equivalent — aggregates each copy independently. {@code SUM(double)} is not
+ * associative, so the copies' partial sums merge in different orders, disagree in the last bits, and the
+ * exact {@code =} matches nothing: q15 then returns 1 row or 0 rows at random (measured 11/20 correct
+ * without this, 20/20 with it). Sharing one evaluation makes both consumers read identical rows, so the
+ * comparison holds whatever order the sum ran in — and halves the work.
+ *
+ *
Not an MPP setting, despite living alongside them historically: sharing is done by
+ * {@code DAGBuilder} for every analytics query and is deliberately NOT gated on {@link #MPP_ENABLED} — the
+ * wrong answer it prevents happens coordinator-centric too. In fact it applies MORE often with distribution
+ * off, because a distributed plan can put the two references in different fragments, where sharing does not
+ * currently reach.
+ *
+ *
Default {@code true}, and it is a KILL SWITCH rather than an opt-in feature flag: the same posture
+ * Spark takes for the equivalent transform ({@code spark.sql.exchange.reuse}, internal, default true since
+ * 2.0.0). {@code SharedSubplanReuse} keeps sharing narrow — only a COMPLETE aggregate subtree with no
+ * shuffle/broadcast/late-materialization boundary — and {@code DAGBuilder} rebuilds without sub-plan reuse when the
+ * consumer would not buffer the shared input. What no internal fallback can catch is a WRONG digest match
+ * (two subtrees that normalize equal without being equivalent), which would be a silent wrong answer; set
+ * this to {@code false} to revert that class of incident without a rollback.
+ */
+ public static final Setting SUBPLAN_REUSE_ENABLED = Setting.boolSetting(
+ "analytics.planner.subplan_reuse.enabled",
+ true,
+ Setting.Property.NodeScope,
+ Setting.Property.Dynamic
+ );
+
/**
* Master switch for hash-shuffle disk spill. When {@code true}, a query whose per-query shuffle
* footprint would exceed the on-heap budget spills its oldest buffered Arrow-IPC chunks to disk
@@ -363,6 +397,7 @@ private AnalyticsSettings() {}
/** All engine-level settings registered by {@code AnalyticsPlugin.getSettings()}. */
public static final List> ALL_SETTINGS = List.of(
MPP_ENABLED,
+ SUBPLAN_REUSE_ENABLED,
BROADCAST_MAX_BYTES,
MPP_SHUFFLE_PARTITIONS,
MPP_SHUFFLE_RECV_TIMEOUT,
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java
index 6861772cd55d4..6bd80a1cb93a9 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java
@@ -347,6 +347,12 @@ private void executeInternal(
final Settings perQuerySettings = Settings.builder()
.put(clusterService.getSettings())
.put(AnalyticsSettings.MPP_ENABLED.getKey(), clusterService.getClusterSettings().get(AnalyticsSettings.MPP_ENABLED))
+ // Sub-plan reuse changes how DAGBuilder cuts the plan, so a static node-bootstrap read would make a dynamic
+ // enable/disable a silent no-op — which is the whole point of keeping it as a kill switch.
+ .put(
+ AnalyticsSettings.SUBPLAN_REUSE_ENABLED.getKey(),
+ clusterService.getClusterSettings().get(AnalyticsSettings.SUBPLAN_REUSE_ENABLED)
+ )
.put(
AnalyticsSettings.MPP_SHUFFLE_AGGREGATE_ENABLED.getKey(),
clusterService.getClusterSettings().get(AnalyticsSettings.MPP_SHUFFLE_AGGREGATE_ENABLED)
@@ -436,7 +442,15 @@ private void executeInternal(
);
}
final String fullPlan = profile ? RelOptUtil.toString(plan) : null;
- QueryDAG dag = DAGBuilder.build(plan, capabilityRegistry, clusterService, indexNameExpressionResolver);
+ // Sub-plan reuse is NOT gated on MPP_ENABLED: a plan that computes the same aggregate twice returns the wrong
+ // answer coordinator-centric too (TPC-H q15), so the sharing has to apply either way.
+ QueryDAG dag = DAGBuilder.build(
+ plan,
+ capabilityRegistry,
+ clusterService,
+ indexNameExpressionResolver,
+ AnalyticsSettings.SUBPLAN_REUSE_ENABLED.get(perQuerySettings)
+ );
// Dispatch resolution under the GENERAL post-CBO scheduler. The enforcement pass placed every
// exchange (shuffle/broadcast) + pre-split any distributed aggregate; DAGBuilder cut at those and
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/join/UnifiedDispatch.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/join/UnifiedDispatch.java
index 0b8da80c87065..9feef947aea19 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/join/UnifiedDispatch.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/join/UnifiedDispatch.java
@@ -361,18 +361,22 @@ private static Stage stripBuildChildren(Stage stage, Map captur
private static void injectBroadcastsInPlace(Stage stage, Map capturedByBuildId) {
if (stage.getFragment() != null) {
List scans = RelNodeUtils.findNodes(stage.getFragment(), OpenSearchBroadcastScan.class);
- List> toInject = new ArrayList<>();
+ // Keyed by build id, so SEVERAL scans on the SAME build inject once. That happens under sub-plan reuse
+ // (SharedSubplanReuse), where every consumer of a shared sub-plan points at one build id and must
+ // resolve the one registered table — duplicating the instruction would re-register the same
+ // namedInputId and ship the payload twice.
+ Map toInject = new LinkedHashMap<>();
for (OpenSearchBroadcastScan scan : scans) {
byte[] ipc = capturedByBuildId.get(scan.getBuildStageId());
if (ipc != null) {
- toInject.add(Map.entry(scan.getBuildStageId(), ipc));
+ toInject.putIfAbsent(scan.getBuildStageId(), ipc);
}
}
if (!toInject.isEmpty()) {
List enriched = new ArrayList<>(stage.getPlanAlternatives().size());
for (StagePlan sp : stage.getPlanAlternatives()) {
List merged = new ArrayList<>(sp.instructions());
- for (Map.Entry e : toInject) {
+ for (Map.Entry e : toInject.entrySet()) {
// buildSideIndex 0: informational only — the NamedScan resolves by name on the data node.
merged.add(new BroadcastInjectionInstructionNode("broadcast-" + e.getKey(), 0, e.getValue()));
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java
index 1ae407ba33813..c4598c8d760b8 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java
@@ -9,6 +9,8 @@
package org.opensearch.analytics.planner.dag;
import org.apache.calcite.rel.RelNode;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
import org.opensearch.analytics.exec.OrdinalAppendingSink;
import org.opensearch.analytics.planner.CapabilityRegistry;
import org.opensearch.analytics.planner.CapabilityResolutionUtils;
@@ -53,6 +55,8 @@
*/
public class DAGBuilder {
+ private static final Logger LOGGER = LogManager.getLogger(DAGBuilder.class);
+
private DAGBuilder() {}
public static QueryDAG build(
@@ -60,6 +64,68 @@ public static QueryDAG build(
CapabilityRegistry registry,
ClusterService clusterService,
IndexNameExpressionResolver indexNameExpressionResolver
+ ) {
+ return build(cboOutput, registry, clusterService, indexNameExpressionResolver, /* subplanReuseEnabled */ false);
+ }
+
+ /**
+ * {@code subplanReuseEnabled} shares a sub-plan this plan computes more than once instead of computing it twice —
+ * see {@link SharedSubplanReuse}, which explains why that is a correctness fix (TPC-H q15) and not only a
+ * saving. Applies to the plan reached through {@code sever}; a root that is itself an exchange or a
+ * late-materialization wrapper is cut before {@code sever} runs and is left alone.
+ */
+ public static QueryDAG build(
+ RelNode cboOutput,
+ CapabilityRegistry registry,
+ ClusterService clusterService,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
+ ) {
+ if (subplanReuseEnabled) {
+ QueryDAG shared = buildInternal(cboOutput, registry, clusterService, indexNameExpressionResolver, true);
+ // Sharing is only sound while the consumer BUFFERS its inputs. The memtable reduce sink is selected
+ // on inputCount > 1, so a stage whose ONLY child is the shared one read twice would get the streaming
+ // sink and see an empty second read. Rather than ship that footgun, detect the shape and rebuild
+ // without sub-plan reuse — a missed reuse is slower, a once-consumable double read is wrong.
+ if (sharedInputsAreBuffered(shared.rootStage())) {
+ return shared;
+ }
+ LOGGER.debug("Rebuilding the DAG without sub-plan sharing: the consumer would not buffer the shared input");
+ }
+ return buildInternal(cboOutput, registry, clusterService, indexNameExpressionResolver, false);
+ }
+
+ /**
+ * True when no stage reads a single child stage more than once, which is the only shape the streaming
+ * (once-consumable) reduce sink cannot serve. See {@link #cutShared}.
+ */
+ private static boolean sharedInputsAreBuffered(Stage stage) {
+ if (stage.getFragment() != null && stage.getChildStages().size() == 1) {
+ int childId = stage.getChildStages().getFirst().getStageId();
+ int refs = 0;
+ for (OpenSearchStageInputScan scan : RelNodeUtils.findNodes(stage.getFragment(), OpenSearchStageInputScan.class)) {
+ if (scan.getChildStageId() == childId) {
+ refs++;
+ }
+ }
+ if (refs > 1) {
+ return false;
+ }
+ }
+ for (Stage child : stage.getChildStages()) {
+ if (!sharedInputsAreBuffered(child)) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private static QueryDAG buildInternal(
+ RelNode cboOutput,
+ CapabilityRegistry registry,
+ ClusterService clusterService,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
) {
int[] counter = { 0 };
List childStages = new ArrayList<>();
@@ -69,16 +135,38 @@ public static QueryDAG build(
// Root IS an ExchangeReducer — pure gather (no compute above the exchange).
// Cut directly: child stage is the subtree below, root fragment is
// ExchangeReducer → StageInputScan.
- rootFragment = cutAtExchange(reducer, counter, childStages, registry, clusterService, indexNameExpressionResolver);
+ rootFragment = cutAtExchange(
+ reducer,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
} else if (cboOutput instanceof OpenSearchLateMaterialization lm) {
// LM at root, no above-ops (e.g. `source = t | where ... | sort col | head N`):
// promote the LM stage to rootStage and skip the synthetic post-LM stage that would
// wrap a bare StageInputScan placeholder.
- cutAtLateMaterialization(lm, counter, childStages, registry, clusterService, indexNameExpressionResolver);
+ cutAtLateMaterialization(lm, counter, childStages, registry, clusterService, indexNameExpressionResolver, subplanReuseEnabled);
assert childStages.size() == 1 : "cutAtLateMaterialization must add exactly one child (the LM stage)";
return new QueryDAG(newQueryId(), childStages.getFirst());
} else {
- rootFragment = sever(cboOutput, counter, childStages, registry, clusterService, indexNameExpressionResolver);
+ SharedSubplanReuse reuse = null;
+ if (subplanReuseEnabled) {
+ SharedSubplanReuse detected = SharedSubplanReuse.detect(cboOutput);
+ reuse = detected.isEmpty() ? null : detected;
+ }
+ rootFragment = sever(
+ cboOutput,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ reuse,
+ subplanReuseEnabled
+ );
}
// Sink provider is needed whenever the root stage runs a backend plan locally —
@@ -117,6 +205,31 @@ private static String newQueryId() {
return UUID.randomUUID().toString();
}
+ /**
+ * Severs a fragment root, giving it its OWN sub-plan-reuse scope when reuse is on.
+ *
+ *
Scoping per FRAGMENT is a correctness requirement, not tidiness: {@link #cutShared} makes the shared
+ * sub-plan a child stage of the fragment being severed, and the consumer resolves it by the named table
+ * {@code input-} that only ITS stage registers. Sharing across two fragments would leave one
+ * of them scanning an input that is not among its child stages — {@code No table named 'input-N'}.
+ */
+ private static RelNode severFragment(
+ RelNode fragmentRoot,
+ int[] counter,
+ List childStages,
+ CapabilityRegistry registry,
+ ClusterService clusterService,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
+ ) {
+ SharedSubplanReuse reuse = null;
+ if (subplanReuseEnabled) {
+ SharedSubplanReuse detected = SharedSubplanReuse.detect(fragmentRoot);
+ reuse = detected.isEmpty() ? null : detected;
+ }
+ return sever(fragmentRoot, counter, childStages, registry, clusterService, indexNameExpressionResolver, reuse, subplanReuseEnabled);
+ }
+
private static RelNode sever(
RelNode node,
int[] counter,
@@ -124,23 +237,93 @@ private static RelNode sever(
CapabilityRegistry registry,
ClusterService clusterService,
IndexNameExpressionResolver indexNameExpressionResolver
+ ) {
+ return sever(node, counter, childStages, registry, clusterService, indexNameExpressionResolver, null, false);
+ }
+
+ /**
+ * {@code reuse} non-null enables sub-plan reuse: a sub-plan that this plan computes more than
+ * once is cut ONCE into a build stage, and later occurrences become an {@code OpenSearchBroadcastScan} on
+ * that same build id so every consumer reads one materialized copy. See {@link SharedSubplanReuse} for why
+ * that is a correctness fix, not just an optimization. Only threaded through {@code sever}'s own recursion —
+ * the {@code cut*} helpers below sever their children WITHOUT it, so a repeat hidden under a shuffle or
+ * broadcast boundary is simply not shared (a missed reuse, which is the pre-existing behaviour).
+ */
+ private static RelNode sever(
+ RelNode node,
+ int[] counter,
+ List childStages,
+ CapabilityRegistry registry,
+ ClusterService clusterService,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ SharedSubplanReuse reuse,
+ boolean subplanReuseEnabled
) {
List newInputs = new ArrayList<>();
List rawInputs = node.getInputs();
for (int inputIndex = 0; inputIndex < rawInputs.size(); inputIndex++) {
RelNode input = rawInputs.get(inputIndex);
- if (input instanceof OpenSearchExchangeReducer reducer) {
- newInputs.add(cutAtExchange(reducer, counter, childStages, registry, clusterService, indexNameExpressionResolver));
+ String sharedDigest = reuse == null ? null : reuse.sharedDigestOf(input);
+ if (sharedDigest != null) {
+ newInputs.add(
+ cutShared(
+ input,
+ sharedDigest,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ reuse,
+ subplanReuseEnabled
+ )
+ );
+ } else if (input instanceof OpenSearchExchangeReducer reducer) {
+ newInputs.add(
+ cutAtExchange(reducer, counter, childStages, registry, clusterService, indexNameExpressionResolver, subplanReuseEnabled)
+ );
} else if (input instanceof OpenSearchShuffleExchange shuffle) {
newInputs.add(
- cutShuffle(shuffle, counter, childStages, registry, clusterService, node, inputIndex, indexNameExpressionResolver)
+ cutShuffle(
+ shuffle,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ node,
+ inputIndex,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ )
);
} else if (input instanceof OpenSearchBroadcastExchange broadcast) {
- newInputs.add(cutBroadcast(broadcast, counter, childStages, registry, clusterService, indexNameExpressionResolver));
+ newInputs.add(
+ cutBroadcast(
+ broadcast,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ )
+ );
} else if (input instanceof OpenSearchLateMaterialization lm) {
- newInputs.add(cutAtLateMaterialization(lm, counter, childStages, registry, clusterService, indexNameExpressionResolver));
+ newInputs.add(
+ cutAtLateMaterialization(
+ lm,
+ counter,
+ childStages,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ )
+ );
} else {
- newInputs.add(sever(input, counter, childStages, registry, clusterService, indexNameExpressionResolver));
+ newInputs.add(
+ sever(input, counter, childStages, registry, clusterService, indexNameExpressionResolver, reuse, subplanReuseEnabled)
+ );
}
}
if (node.getInputs().isEmpty()) return node;
@@ -185,11 +368,20 @@ private static RelNode cutAtLateMaterialization(
List parentChildStages,
CapabilityRegistry registry,
ClusterService clusterService,
- IndexNameExpressionResolver indexNameExpressionResolver
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
) {
// 1. Reduce child — Sort+Limit reduce above shard scans. Multi-shard QTF only.
List reduceChildren = new ArrayList<>();
- RelNode reduceFragment = sever(lm.getInput(), counter, reduceChildren, registry, clusterService, indexNameExpressionResolver);
+ RelNode reduceFragment = severFragment(
+ lm.getInput(),
+ counter,
+ reduceChildren,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
if (reduceChildren.isEmpty()) {
throw new IllegalStateException(
"QTF rewriter fired but the wrapper's input has no ExchangeReducer below it — "
@@ -270,13 +462,22 @@ private static RelNode cutAtExchange(
List parentChildStages,
CapabilityRegistry registry,
ClusterService clusterService,
- IndexNameExpressionResolver indexNameExpressionResolver
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
) {
// Recurse into the child fragment with full sever() so any nested ExchangeReducers
// (e.g. a Join below a top-level gather Reducer) are also cut into their own child
// stages rather than being left intact inside the shard-local fragment.
List grandchildren = new ArrayList<>();
- RelNode childFragment = sever(reducer.getInput(), counter, grandchildren, registry, clusterService, indexNameExpressionResolver);
+ RelNode childFragment = severFragment(
+ reducer.getInput(),
+ counter,
+ grandchildren,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
int childStageId = counter[0]++;
// Stage execution location is decided by the fragment's contents, not the grandchild
@@ -355,14 +556,23 @@ private static RelNode cutShuffle(
ClusterService clusterService,
RelNode parent,
int parentInputIndex,
- IndexNameExpressionResolver indexNameExpressionResolver
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
) {
// Recurse into the shuffle's input with full sever() so any nested exchanges below the
// shuffle (e.g. a partial-aggregate that itself reduces) are also cut into their own
// stages. M2 today only composes shuffle over a shard scan, but the recursion makes the
// cutter robust to future plan shapes.
List grandchildren = new ArrayList<>();
- RelNode childFragment = sever(shuffle.getInput(), counter, grandchildren, registry, clusterService, indexNameExpressionResolver);
+ RelNode childFragment = severFragment(
+ shuffle.getInput(),
+ counter,
+ grandchildren,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
int childStageId = counter[0]++;
// Decide the producer's locality by whether its fragment has a shard scan — NOT by child-stage
@@ -457,16 +667,110 @@ private static boolean containsAnyInput(RelNode root, Class t
* the empty-grandchildren branch would need to gain a sink provider just like the other
* cutters. For now the build always reduces to a leaf shard fragment.
*/
+ /**
+ * Cuts a SHARED sub-plan (see {@link SharedSubplanReuse}) into ONE ordinary child stage the first time it is
+ * seen, and returns an {@link OpenSearchStageInputScan} on it. Later occurrences scan that SAME child stage
+ * id, so the consumer's fragment references the one named table {@code input-} more than once —
+ * which is what makes two exact-equality consumers read identical rows.
+ *
+ *
Why an ordinary child stage and not a broadcast build. The consumer registers each child input as
+ * a re-readable {@code MemTable}: a multi-input coordinator stage is routed to
+ * {@code DatafusionMemtableReduceSink}, which buffers each input and hands it across in one
+ * {@code registerMemtable} call, so two reads of one input resolve to the same materialized batches with no
+ * extra machinery. Routing through {@code BROADCAST_BUILD} instead would need the broadcast-injection
+ * handler, which requires a DataFusion session created by a prior SHARD-SCAN handler — a coordinator reduce
+ * stage has none, and it fails with {@code BroadcastInjectionHandler: expected DataFusionSessionState … got
+ * null} (measured).
+ *
+ *
Depends on the consumer buffering its inputs. The streaming reduce sink registers each input as a
+ * bounded {@code StreamingTable}, consumable ONCE, so a second read would come back empty. The memtable sink
+ * is selected on {@code inputCount > 1}, so sharing is only safe while the consumer keeps another input
+ * besides the shared one.
+ */
+ private static RelNode cutShared(
+ RelNode shared,
+ String digest,
+ int[] counter,
+ List parentChildStages,
+ CapabilityRegistry registry,
+ ClusterService clusterService,
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ SharedSubplanReuse reuse,
+ boolean subplanReuseEnabled
+ ) {
+ List viableBackends = ((OpenSearchRelNode) shared).getViableBackends();
+ Integer existing = reuse.alreadyCutStageId(digest);
+ if (existing == null) {
+ List grandchildren = new ArrayList<>();
+ RelNode childFragment = severFragment(
+ shared,
+ counter,
+ grandchildren,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
+
+ int childStageId = counter[0]++;
+ // Same locality rule as cutAtExchange: a fragment containing a TableScan runs on shards even when it
+ // also has grandchildren; a coordinator-side fragment consumes its grandchildren through a sink.
+ boolean fragmentHasShardScan = containsAnyInput(childFragment, OpenSearchTableScan.class);
+ TargetResolver targetResolver = fragmentHasShardScan
+ ? new ShardTargetResolver(childFragment, clusterService, indexNameExpressionResolver)
+ : null;
+ ExchangeSinkProvider childSinkProvider = null;
+ if (!grandchildren.isEmpty() && !fragmentHasShardScan) {
+ List reduceViable = CapabilityResolutionUtils.filterByReduceCapability(registry, viableBackends);
+ childSinkProvider = registry.getBackend(reduceViable.getFirst()).getExchangeSinkProvider();
+ }
+ Stage sharedStage = new Stage(
+ childStageId,
+ childFragment,
+ grandchildren,
+ ExchangeInfo.singleton(),
+ childSinkProvider,
+ targetResolver
+ );
+ parentChildStages.add(sharedStage);
+ reuse.recordCut(digest, childStageId);
+ existing = childStageId;
+ LOGGER.debug("Sharing a repeated sub-plan: cut into child stage {}", childStageId);
+ } else {
+ LOGGER.debug("Sharing a repeated sub-plan: reusing child stage {}", existing);
+ }
+ // Per-column storage comes from the node this scan stands in for — a leaf that reports a short
+ // getOutputFieldStorage() truncates the storage union and fails conversion with
+ // "RexInputRef[N] has no matching FieldStorageInfo entry".
+ return new OpenSearchStageInputScan(
+ shared.getCluster(),
+ shared.getTraitSet(),
+ existing,
+ shared.getRowType(),
+ viableBackends,
+ ((OpenSearchRelNode) shared).getOutputFieldStorage()
+ );
+ }
+
private static RelNode cutBroadcast(
OpenSearchBroadcastExchange broadcast,
int[] counter,
List parentChildStages,
CapabilityRegistry registry,
ClusterService clusterService,
- IndexNameExpressionResolver indexNameExpressionResolver
+ IndexNameExpressionResolver indexNameExpressionResolver,
+ boolean subplanReuseEnabled
) {
List grandchildren = new ArrayList<>();
- RelNode childFragment = sever(broadcast.getInput(), counter, grandchildren, registry, clusterService, indexNameExpressionResolver);
+ RelNode childFragment = severFragment(
+ broadcast.getInput(),
+ counter,
+ grandchildren,
+ registry,
+ clusterService,
+ indexNameExpressionResolver,
+ subplanReuseEnabled
+ );
int childStageId = counter[0]++;
TargetResolver targetResolver = grandchildren.isEmpty()
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/SharedSubplanReuse.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/SharedSubplanReuse.java
new file mode 100644
index 0000000000000..a0bc0b526098d
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/SharedSubplanReuse.java
@@ -0,0 +1,156 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.planner.dag;
+
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.rel.RelNode;
+import org.opensearch.analytics.planner.rel.AggregateMode;
+import org.opensearch.analytics.planner.rel.OpenSearchAggregate;
+import org.opensearch.analytics.planner.rel.OpenSearchBroadcastExchange;
+import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization;
+import org.opensearch.analytics.planner.rel.OpenSearchShuffleExchange;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+/**
+ * Sub-plan reuse detection for a plan that computes the SAME complete aggregate more than once.
+ *
+ *
The correctness problem this solves. A query may inline one subquery twice — TPC-H q15 is
+ * {@code supplier ⋈ revenue0} plus {@code where total_revenue = [ … max(total_revenue) ]} over the same
+ * {@code revenue0}, because the spec's {@code revenue0} VIEW has no PPL equivalent. Each copy is aggregated
+ * independently, and {@code SUM(double)} is not associative, so the two copies' partial sums are merged in
+ * different orders and disagree in the last bits. The exact {@code =} then matches nothing and the row
+ * vanishes: q15 returns 1 row or 0 rows at random (measured ~9/20 correct, in every distribution
+ * configuration — coordinator-centric included, so this is not an MPP artifact).
+ *
+ *
Making float summation order-independent would be one cure; computing the shared relation ONCE is the
+ * other, and it is correct by construction rather than by numeric luck — both consumers then read the very
+ * same rows, so the equality holds whatever order the sum ran in. It also halves the work.
+ *
+ *
How. {@link DAGBuilder} asks this class, as it severs the plan, whether a node it is about to
+ * descend into is a shared sub-plan. The first occurrence is cut into ONE ordinary child stage; every
+ * occurrence then becomes an {@code OpenSearchStageInputScan} on that same child stage id, so the consumer's
+ * fragment scans the single named table {@code input-} more than once. That works with no new
+ * transport because a multi-input coordinator stage is served by {@code DatafusionMemtableReduceSink}, which
+ * buffers each child input into a re-readable {@code MemTable}.
+ *
+ *
Deliberately narrow. Only a COMPLETE aggregate ({@link AggregateMode#FINAL} or
+ * {@link AggregateMode#SINGLE}) is a candidate, and only when its subtree contains no shuffle, broadcast or
+ * late-materialization boundary. That bounds what gets buffered (an aggregate's output, not a raw scan's) and
+ * keeps the shared stage a plain gather. Anything else is left alone — a missed reuse costs performance, a wrong
+ * one costs correctness.
+ *
+ *
Sharing is scoped per FRAGMENT and per buffered consumer, both enforced in {@link DAGBuilder}: a
+ * shared stage must be a direct child of the stage scanning it (else {@code No table named 'input-N'}), and the
+ * consumer must buffer its inputs (the streaming reduce sink's inputs are once-consumable, so a second read
+ * returns nothing). {@code DAGBuilder} rebuilds without sub-plan reuse rather than emit either shape.
+ *
+ * @opensearch.internal
+ */
+final class SharedSubplanReuse {
+
+ /**
+ * Annotation ids are a per-query sequential counter ({@code ANNOTATED_PREDICATE(id=0, …)}), so two
+ * semantically identical subtrees carry DIFFERENT ids — in q15 one copy has {@code id=0,1,4,5} and the
+ * other {@code id=2,3,6,7}. They must not defeat the match, and dropping them is safe: the surviving copy
+ * keeps its own annotations, and the eliminated copy's are simply no longer referenced.
+ */
+ private static final Pattern ANNOTATION_ID = Pattern.compile("id=\\d+, ");
+
+ private final Set sharedDigests;
+ private final Map cutStageIdByDigest = new HashMap<>();
+
+ private SharedSubplanReuse(Set sharedDigests) {
+ this.sharedDigests = sharedDigests;
+ }
+
+ /** Digests every candidate in {@code root} and retains those occurring more than once. */
+ static SharedSubplanReuse detect(RelNode root) {
+ Map counts = new HashMap<>();
+ Deque queue = new ArrayDeque<>();
+ queue.push(root);
+ while (!queue.isEmpty()) {
+ RelNode node = queue.pop();
+ if (isCandidate(node)) {
+ counts.merge(digestOf(node), 1, Integer::sum);
+ }
+ for (RelNode input : node.getInputs()) {
+ queue.push(input);
+ }
+ }
+ Set shared = new HashSet<>();
+ for (Map.Entry e : counts.entrySet()) {
+ if (e.getValue() > 1) {
+ shared.add(e.getKey());
+ }
+ }
+ return new SharedSubplanReuse(shared);
+ }
+
+ /** True when no candidate repeats, so {@link DAGBuilder} can skip every reuse check. */
+ boolean isEmpty() {
+ return sharedDigests.isEmpty();
+ }
+
+ /** The digest of {@code node} if it is a shared sub-plan, else {@code null}. */
+ String sharedDigestOf(RelNode node) {
+ if (!isCandidate(node)) {
+ return null;
+ }
+ String digest = digestOf(node);
+ return sharedDigests.contains(digest) ? digest : null;
+ }
+
+ /** The child-stage id already cut for {@code digest}, or {@code null} on first encounter. */
+ Integer alreadyCutStageId(String digest) {
+ return cutStageIdByDigest.get(digest);
+ }
+
+ void recordCut(String digest, int childStageId) {
+ cutStageIdByDigest.put(digest, childStageId);
+ }
+
+ private static boolean isCandidate(RelNode node) {
+ if (!(node instanceof OpenSearchAggregate aggregate)) {
+ return false;
+ }
+ if (aggregate.getMode() != AggregateMode.FINAL && aggregate.getMode() != AggregateMode.SINGLE) {
+ return false;
+ }
+ return !containsUnsupportedBoundary(node);
+ }
+
+ /**
+ * A shared subtree is cut as a plain gather stage. A shuffle / broadcast / late-materialization boundary
+ * inside would make that stage need producer or injection wiring too, so those subtrees are not shared.
+ */
+ private static boolean containsUnsupportedBoundary(RelNode node) {
+ if (node instanceof OpenSearchShuffleExchange
+ || node instanceof OpenSearchBroadcastExchange
+ || node instanceof OpenSearchLateMaterialization) {
+ return true;
+ }
+ for (RelNode input : node.getInputs()) {
+ if (containsUnsupportedBoundary(input)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private static String digestOf(RelNode node) {
+ return ANNOTATION_ID.matcher(RelOptUtil.toString(node)).replaceAll("");
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/SharedSubplanReuseTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/SharedSubplanReuseTests.java
new file mode 100644
index 0000000000000..551160d52458a
--- /dev/null
+++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/SharedSubplanReuseTests.java
@@ -0,0 +1,209 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.planner.dag;
+
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.core.CorrelationId;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.opensearch.analytics.planner.BasePlannerRulesTests;
+import org.opensearch.analytics.planner.PlannerContext;
+import org.opensearch.analytics.planner.RelNodeUtils;
+import org.opensearch.analytics.planner.rel.AggregateMode;
+import org.opensearch.analytics.planner.rel.OpenSearchAggregate;
+import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Sub-plan reuse in {@link DAGBuilder} ({@code analytics.planner.subplan_reuse.enabled}).
+ *
+ *
Why this matters beyond saving work: a plan that computes the same aggregate twice returns the WRONG
+ * ANSWER when the two copies are compared for exact equality, because {@code SUM(double)} is not associative
+ * and the copies' partial sums merge in different orders. That is TPC-H q15, which returns 1 row or 0 rows at
+ * random. Sharing one evaluation makes the comparison hold by construction. See {@link SharedSubplanReuse}.
+ */
+public class SharedSubplanReuseTests extends BasePlannerRulesTests {
+
+ /** Off (the default): the duplicated aggregate is cut twice, and nothing references a shared build. */
+ public void testReuseDisabled_duplicateAggregateIsComputedTwice() {
+ PlannerContext context = buildContext("parquet", 2, intFields());
+ RelNode cbo = runPlanner(joinOfTwoIdenticalAggregates(), context);
+
+ QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER);
+
+ assertEquals("both copies of the aggregate are cut as their own stages with reuse off", 2, stageInputScans(dag).size());
+ assertEquals("both copies of the aggregate are still computed", 2, completeAggregateCount(dag));
+ }
+
+ /**
+ * The consumer would have the shared aggregate as its ONLY input, so it would be served by the streaming
+ * (once-consumable) reduce sink and the second read would come back empty. {@code DAGBuilder} must detect
+ * that and rebuild without sub-plan reuse — a missed reuse is slower, a double read of a once-consumable input is wrong.
+ */
+ public void testReuseFallsBackWhenTheSharedInputWouldNotBeBuffered() {
+ PlannerContext context = buildContext("parquet", 2, intFields());
+ RelNode cbo = runPlanner(joinOfTwoIdenticalAggregates(), context);
+
+ QueryDAG dag = DAGBuilder.build(
+ cbo,
+ context.getCapabilityRegistry(),
+ mockClusterService(),
+ TEST_RESOLVER,
+ /* subplanReuseEnabled */ true
+ );
+
+ List scans = stageInputScans(dag);
+ assertEquals("the two copies stay on separate stages", 2, scans.size());
+ assertNotEquals(
+ "sharing must NOT happen when the consumer would not buffer the shared input",
+ scans.get(0).getChildStageId(),
+ scans.get(1).getChildStageId()
+ );
+ assertEquals("both copies are still computed (the fallback)", 2, completeAggregateCount(dag));
+ }
+
+ /**
+ * On, with the consumer keeping another input: the aggregate is cut ONCE and BOTH consumers scan that one
+ * child stage, so the buffered memtable input is read twice and every consumer sees identical rows.
+ */
+ public void testReuseEnabled_duplicateAggregateIsSharedByChildStageId() {
+ PlannerContext context = buildContextPerIndex("parquet", Map.of("test_index", 2, "other_idx", 2));
+ RelNode cbo = runPlanner(joinKeepingAnotherInputBesideTheSharedAggregate(), context);
+
+ QueryDAG dag = DAGBuilder.build(
+ cbo,
+ context.getCapabilityRegistry(),
+ mockClusterService(),
+ TEST_RESOLVER,
+ /* subplanReuseEnabled */ true
+ );
+
+ // THE point of the feature: one evaluation, so both consumers read identical rows and an exact-equality
+ // comparison between them cannot fall foul of float accumulation order.
+ assertEquals("the duplicated aggregate is computed exactly ONCE", 1, completeAggregateCount(dag));
+
+ Map refsByChildStage = stageInputScans(dag).stream()
+ .collect(
+ java.util.stream.Collectors.groupingBy(OpenSearchStageInputScan::getChildStageId, java.util.stream.Collectors.counting())
+ );
+ assertTrue(
+ "exactly one child stage must be scanned twice (the shared aggregate), got " + refsByChildStage,
+ refsByChildStage.values().stream().filter(c -> c == 2L).count() == 1
+ );
+ }
+
+ // ─── helpers ───────────────────────────────────────────────────────────────
+
+ /**
+ * {@code Join(Aggregate(scan), Aggregate(scan))} where the two aggregates are IDENTICAL — the shape q15
+ * produces by inlining one subquery twice.
+ */
+ private RelNode joinOfTwoIdenticalAggregates() {
+ RelNode left = identicalAggregate();
+ RelNode right = identicalAggregate();
+ RexNode condition = rexBuilder.makeCall(
+ SqlStdOperatorTable.EQUALS,
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 0),
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 2)
+ );
+ return LogicalJoin.create(left, right, List.of(), condition, Set.of(), JoinRelType.INNER);
+ }
+
+ /**
+ * {@code Join(Join(otherScan, sharedAgg), sharedAgg)} — the q15 skeleton. The consumer keeps a second input
+ * (the other scan's gather) besides the shared aggregate, so it buffers its inputs and sharing is sound.
+ */
+ private RelNode joinKeepingAnotherInputBesideTheSharedAggregate() {
+ RelNode other = stubScan(mockTable("other_idx", "status", "size"));
+ RelNode inner = LogicalJoin.create(
+ other,
+ identicalAggregate(),
+ List.of(),
+ rexBuilder.makeCall(
+ SqlStdOperatorTable.EQUALS,
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 0),
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 2)
+ ),
+ Set.of(),
+ JoinRelType.INNER
+ );
+ return LogicalJoin.create(
+ inner,
+ identicalAggregate(),
+ List.of(),
+ rexBuilder.makeCall(
+ SqlStdOperatorTable.EQUALS,
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 0),
+ rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 4)
+ ),
+ Set.of(),
+ JoinRelType.INNER
+ );
+ }
+
+ private RelNode identicalAggregate() {
+ RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+ AggregateCall sum = AggregateCall.create(
+ SqlStdOperatorTable.SUM,
+ false,
+ List.of(1),
+ -1,
+ scan,
+ typeFactory.createSqlType(SqlTypeName.INTEGER),
+ "total"
+ );
+ return LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(sum));
+ }
+
+ private static List stageInputScans(QueryDAG dag) {
+ List found = new ArrayList<>();
+ collectScans(dag.rootStage(), found);
+ return found;
+ }
+
+ private static void collectScans(Stage stage, List found) {
+ if (stage.getFragment() != null) {
+ found.addAll(RelNodeUtils.findNodes(stage.getFragment(), OpenSearchStageInputScan.class));
+ }
+ for (Stage child : stage.getChildStages()) {
+ collectScans(child, found);
+ }
+ }
+
+ /** Counts FINAL/SINGLE aggregates across every stage — one per surviving evaluation. */
+ private static int completeAggregateCount(QueryDAG dag) {
+ return countCompleteAggregates(dag.rootStage());
+ }
+
+ private static int countCompleteAggregates(Stage stage) {
+ int count = 0;
+ if (stage.getFragment() != null) {
+ for (OpenSearchAggregate aggregate : RelNodeUtils.findNodes(stage.getFragment(), OpenSearchAggregate.class)) {
+ if (aggregate.getMode() == AggregateMode.FINAL || aggregate.getMode() == AggregateMode.SINGLE) {
+ count++;
+ }
+ }
+ }
+ for (Stage child : stage.getChildStages()) {
+ count += countCompleteAggregates(child);
+ }
+ return count;
+ }
+
+}
diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SharedSubplanReuseIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SharedSubplanReuseIT.java
new file mode 100644
index 0000000000000..28a04c50841df
--- /dev/null
+++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SharedSubplanReuseIT.java
@@ -0,0 +1,321 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.qa;
+
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * End-to-end tests for sub-plan reuse ({@code analytics.planner.subplan_reuse.enabled}).
+ *
+ *
What breaks without it. A query that inlines the same aggregate subquery TWICE and then compares the
+ * two results for exact equality is nondeterministic: {@code SUM(double)} is not associative, the two copies'
+ * partial sums merge in different orders across shards and intra-shard slices, they disagree in the last bits,
+ * and the {@code =} matches nothing — so the query returns a row on some runs and nothing on others. TPC-H q15
+ * is this shape ({@code revenue0} is a VIEW in the spec, with no PPL equivalent, so it gets inlined twice) and
+ * measured 11/20 correct. Sharing one evaluation makes both consumers read identical rows, so the comparison
+ * holds regardless of accumulation order.
+ *
+ *
Why the assertion is REPEATED runs, not a single one. The bug is a coin flip, so one passing run
+ * proves nothing — the sweep passed q15 by luck more than once while the bug was live. Each test here runs the
+ * query {@link #RUNS} times and requires EVERY run to be correct; that is the only assertion that distinguishes
+ * "fixed" from "got lucky".
+ *
+ *
Why the data looks like this. The amounts are decimals with no exact binary representation, spread
+ * over {@link #SHARDS} shards with many rows per group, so the partial sums genuinely differ by ULPs depending
+ * on merge order. With round numbers they would agree exactly every time and these tests would pass whether or
+ * not anything was shared — proving nothing. Data adequacy was therefore checked once, out of band, by forcing
+ * the toggle off and observing this query drop to 0 rows within a few runs.
+ *
+ *
Note what is NOT asserted. Turning sharing off is not expected to make anything fail — it is an
+ * optimization, so both settings must produce the same ANSWER; only the PLAN differs. The flakiness that
+ * motivated this work is a pre-existing defect of exact float equality, not a property of the toggle. So the
+ * tests here assert (a) the plan really does share, via the per-stage profile, and (b) results are correct and
+ * unchanged. {@link #testSharedAggregateIsComputedOnceInThePlan} is the one that proves the optimization
+ * applied at all.
+ */
+public class SharedSubplanReuseIT extends AnalyticsRestTestCase {
+
+ private static final String FACTS = "cse_facts";
+ private static final String DIM = "cse_dim";
+ private static final int SHARDS = 3;
+ private static final int GROUPS = 40;
+ private static final int ROWS_PER_GROUP = 60;
+ /** Enough repeats that a ~50% coin flip cannot survive by chance (0.5^10 is under 0.1%). */
+ private static final int RUNS = 10;
+
+ private static boolean dataProvisioned = false;
+
+ @Override
+ public void tearDown() throws Exception {
+ resetSetting("analytics.planner.subplan_reuse.enabled");
+ resetSetting("analytics.mpp.enabled");
+ super.tearDown();
+ }
+
+ /**
+ * The q15 shape: join against an aggregate subquery, then filter on exact equality with a {@code max()} over
+ * the SAME subquery. With sharing on, every run must return the one matching group.
+ */
+ public void testExactEqualityOverSharedAggregate_isDeterministic() throws Exception {
+ ensureDataProvisioned();
+ enableSubplanReuse();
+
+ for (int run = 0; run < RUNS; run++) {
+ List> rows = executePplRows(sharedAggregateEqualityQuery());
+ assertEquals(
+ "run " + run + ": exact equality against a shared aggregate must match the top group on EVERY run. "
+ + "A run returning 0 rows means the equality did not match on that run — the cause is NOT "
+ + "asserted here (float accumulation order is the known one, but a plan or resource problem "
+ + "would also land here; note a resource failure would instead throw before this assertion).",
+ 1,
+ rows.size()
+ );
+ }
+ }
+
+ /**
+ * The direct evidence that the optimization applied: with sharing off the aggregate subquery is scanned by
+ * two separate stages, with it on by one. Read from the per-stage {@code fragment} the query profile
+ * publishes, so it asserts the executed DAG rather than the pre-DAG plan text.
+ */
+ public void testSharedAggregateIsComputedOnceInThePlan() throws Exception {
+ ensureDataProvisioned();
+ applySetting("analytics.mpp.enabled", "false");
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "false");
+ int stagesWithout = countStagesScanningFacts(sharedAggregateEqualityQuery());
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "true");
+ int stagesWith = countStagesScanningFacts(sharedAggregateEqualityQuery());
+
+ assertEquals("without sharing, the duplicated subquery is scanned by two stages", 2, stagesWithout);
+ assertEquals("with sharing, one stage scans it and both consumers read that stage", 1, stagesWith);
+ }
+
+ /**
+ * Sharing must not change the ANSWER, only how many times it is computed. Compared against the same query
+ * with sharing off — and on an INTEGER sum, which is associative, so the off arm is deterministic too and
+ * this comparison is stable rather than a coin flip.
+ */
+ public void testSharingDoesNotChangeResults_integerSumIsStable() throws Exception {
+ ensureDataProvisioned();
+ String ppl = "source = " + FACTS + " | stats sum(qty) as total_qty by group_id | sort group_id";
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "false");
+ List> without = executePplRows(ppl);
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "true");
+ List> with = executePplRows(ppl);
+
+ assertFalse("baseline must return rows (otherwise the comparison is vacuous)", without.isEmpty());
+ assertEquals("every group is returned", GROUPS, without.size());
+ assertRowMultisetEquals("sharing must not change results", without, with);
+ }
+
+ /**
+ * A query with NO duplicated sub-plan must be untouched — the detector has to be inert rather than
+ * rearranging plans it has no reason to touch.
+ */
+ public void testQueryWithoutDuplicateSubplan_isUnaffected() throws Exception {
+ ensureDataProvisioned();
+ String ppl = "source = " + FACTS + " | where group_id < 5 | stats count() as c, sum(qty) as q by group_id | sort group_id";
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "false");
+ List> without = executePplRows(ppl);
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "true");
+ List> with = executePplRows(ppl);
+
+ assertFalse("baseline must return rows", without.isEmpty());
+ assertRowMultisetEquals("a plan with nothing to share must be unchanged", without, with);
+ }
+
+ /** Sharing must survive the MPP path being on as well — no crash, no lost rows. */
+ public void testSharedAggregateUnderMpp_matchesNonMppRows() throws Exception {
+ ensureDataProvisioned();
+ String ppl = sharedAggregateEqualityQuery();
+
+ applySetting("analytics.planner.subplan_reuse.enabled", "true");
+ applySetting("analytics.mpp.enabled", "false");
+ List> nonMpp = executePplRows(ppl);
+ assertEquals("the non-MPP arm must return the single top group", 1, nonMpp.size());
+
+ applySetting("analytics.mpp.enabled", "true");
+ applySetting("analytics.mpp.distribute.min_rows", "1");
+ List> mpp = executePplRows(ppl);
+ // NOTE: under MPP the two copies can land in different fragments, where sharing does not apply — so this
+ // asserts only that enabling MPP alongside sub-plan reuse stays correct, NOT that sharing fired.
+ assertEquals("MPP arm must return one row too", 1, mpp.size());
+ assertRowMultisetEquals("MPP must not change the shared-aggregate answer", nonMpp, mpp);
+ }
+
+ // ─── query ─────────────────────────────────────────────────────────────────
+
+ /**
+ * {@code dim ⋈ (sum by group) } then {@code where total = [ max(total) over the SAME subquery ]} — the
+ * aggregate subquery text appears twice, which is what gives the planner something to share.
+ */
+ private String sharedAggregateEqualityQuery() {
+ String subquery = "source = " + FACTS + " | stats sum(amount) as total by group_id";
+ return "source = " + DIM + " | join right = rev ON dim_id = group_id [ " + subquery + " ] "
+ + "| where total = [ source = [ " + subquery + " ] | stats max(total) ] "
+ + "| fields dim_id, total";
+ }
+
+ // ─── provisioning ──────────────────────────────────────────────────────────
+
+ private void ensureDataProvisioned() throws IOException {
+ if (dataProvisioned) {
+ return;
+ }
+ createIndex(DIM, "{\"dim_id\":{\"type\":\"integer\"}}");
+ StringBuilder dim = new StringBuilder();
+ for (int g = 0; g < GROUPS; g++) {
+ dim.append("{\"index\":{}}\n");
+ dim.append("{\"dim_id\":").append(g).append("}\n");
+ }
+ bulkAndRefresh(DIM, dim.toString());
+
+ // Amounts are decimals with no exact binary representation and each group gets many of them, so the
+ // per-shard / per-slice partial sums differ by ULPs depending on the order they are merged in. Group
+ // totals increase with group_id so the max is unique and the expected row count is exactly 1.
+ createIndex(FACTS, "{\"group_id\":{\"type\":\"integer\"},\"amount\":{\"type\":\"double\"},\"qty\":{\"type\":\"integer\"}}");
+ StringBuilder facts = new StringBuilder();
+ for (int g = 0; g < GROUPS; g++) {
+ for (int r = 0; r < ROWS_PER_GROUP; r++) {
+ double amount = 0.1 + (g * 0.07) + (r * 0.013);
+ facts.append("{\"index\":{}}\n");
+ facts.append("{\"group_id\":").append(g).append(",\"amount\":").append(amount).append(",\"qty\":").append(r).append("}\n");
+ }
+ }
+ bulkAndRefresh(FACTS, facts.toString());
+ dataProvisioned = true;
+ }
+
+ /** Parquet-primary composite index — the analytics engine only plans over these. */
+ private void createIndex(String indexName, String mappingProperties) throws IOException {
+ try {
+ client().performRequest(new Request("DELETE", "/" + indexName));
+ } catch (Exception ignored) {
+ // first run — nothing to delete
+ }
+ Request request = new Request("PUT", "/" + indexName);
+ request.setJsonEntity(
+ "{"
+ + "\"settings\": {"
+ + " \"number_of_shards\": " + SHARDS + ","
+ + " \"number_of_replicas\": 0,"
+ + " \"index.pluggable.dataformat.enabled\": true,"
+ + " \"index.pluggable.dataformat\": \"composite\","
+ + " \"index.composite.primary_data_format\": \"parquet\","
+ + " \"index.composite.secondary_data_formats\": [\"lucene\"]"
+ + "},"
+ + "\"mappings\": { \"properties\": " + mappingProperties + " }"
+ + "}"
+ );
+ Map response = assertOkAndParse(client().performRequest(request), "Create index " + indexName);
+ assertEquals("index creation must be acknowledged", true, response.get("acknowledged"));
+
+ Request health = new Request("GET", "/_cluster/health/" + indexName);
+ health.addParameter("wait_for_status", "yellow");
+ health.addParameter("timeout", "60s");
+ client().performRequest(health);
+ }
+
+ private void bulkAndRefresh(String indexName, String bulkBody) throws IOException {
+ Request bulkRequest = new Request("POST", "/" + indexName + "/_bulk");
+ bulkRequest.setJsonEntity(bulkBody);
+ bulkRequest.addParameter("refresh", "true");
+ bulkRequest.setOptions(bulkRequest.getOptions().toBuilder().addHeader("Content-Type", "application/x-ndjson").build());
+ client().performRequest(bulkRequest);
+ client().performRequest(new Request("POST", "/" + indexName + "/_flush?force=true"));
+ }
+
+ // ─── helpers ───────────────────────────────────────────────────────────────
+
+ /**
+ * Sharing applies where both copies sit in ONE fragment, which is the coordinator-centric plan — the QA
+ * cluster turns MPP on globally, so these tests turn it off unless they are specifically exercising MPP.
+ */
+ private void enableSubplanReuse() throws IOException {
+ applySetting("analytics.planner.subplan_reuse.enabled", "true");
+ applySetting("analytics.mpp.enabled", "false");
+ }
+
+ /** Stages whose fragment scans {@link #FACTS} — one per surviving evaluation of the shared subquery. */
+ @SuppressWarnings("unchecked")
+ private int countStagesScanningFacts(String ppl) throws IOException {
+ Request request = new Request("POST", "/_analytics/ppl");
+ request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\", \"profile\": true}");
+ Map body = assertOkAndParse(client().performRequest(request), "PPL(profile): " + ppl);
+ Map profile = (Map) body.get("profile");
+ assertNotNull("profile must be present when profile=true", profile);
+ List> stages = (List>) profile.get("stages");
+ assertNotNull("profile.stages must be present", stages);
+ int count = 0;
+ for (Map stage : stages) {
+ List fragment = (List) stage.get("fragment");
+ if (fragment != null && String.join("\n", fragment).contains(FACTS)) {
+ count++;
+ }
+ }
+ return count;
+ }
+
+ private List> executePplRows(String ppl) throws IOException {
+ Request request = new Request("POST", "/_analytics/ppl");
+ request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}");
+ Response response = client().performRequest(request);
+ Map body = assertOkAndParse(response, "PPL: " + ppl);
+ @SuppressWarnings("unchecked")
+ List> rows = (List>) body.get("rows");
+ assertNotNull("Response missing 'rows' field for query: " + ppl, rows);
+ return rows;
+ }
+
+ private void applySetting(String key, String value) throws IOException {
+ Request request = new Request("PUT", "/_cluster/settings");
+ request.setJsonEntity("{\"transient\": {\"" + key + "\": " + value + "}}");
+ client().performRequest(request);
+ }
+
+ private void resetSetting(String key) throws IOException {
+ Request request = new Request("PUT", "/_cluster/settings");
+ request.setJsonEntity("{\"transient\": {\"" + key + "\": null}}");
+ client().performRequest(request);
+ }
+
+ private static void assertRowMultisetEquals(String message, List> expected, List> actual) {
+ List expectedNorm = expected.stream().map(SharedSubplanReuseIT::normalizeRow).sorted().toList();
+ List actualNorm = actual.stream().map(SharedSubplanReuseIT::normalizeRow).sorted().toList();
+ assertEquals(message, expectedNorm, actualNorm);
+ }
+
+ private static String normalizeRow(List