From 87bd880b7fc73005c032bf0c11c482eec0add5d6 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Mon, 24 Aug 2026 08:31:07 +0000 Subject: [PATCH] Reuse a repeated aggregate sub-plan instead of evaluating it twice Signed-off-by: Lantao Jin --- .../analytics/AnalyticsSettings.java | 35 ++ .../analytics/exec/DefaultPlanExecutor.java | 16 +- .../analytics/exec/join/UnifiedDispatch.java | 10 +- .../analytics/planner/dag/DAGBuilder.java | 338 +++++++++++++++++- .../planner/dag/SharedSubplanReuse.java | 156 ++++++++ .../planner/dag/SharedSubplanReuseTests.java | 209 +++++++++++ .../analytics/qa/SharedSubplanReuseIT.java | 321 +++++++++++++++++ 7 files changed, 1064 insertions(+), 21 deletions(-) create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/SharedSubplanReuse.java create mode 100644 sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/SharedSubplanReuseTests.java create mode 100644 sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/SharedSubplanReuseIT.java 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 ad4f94be09dcc..bf5f826498b50 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 @@ -344,6 +344,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) @@ -433,7 +439,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 row) { + StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < row.size(); i++) { + if (i > 0) sb.append('|'); + sb.append(normalizeCell(row.get(i))); + } + return sb.append(']').toString(); + } + + private static String normalizeCell(Object cell) { + if (cell == null) return ""; + // Doubles that differ only in the last bits are the SAME answer — the defect is a missing row, not a + // differing tail — so compare numerics at a tolerance rather than bit-for-bit. + if (cell instanceof Number) return String.format(java.util.Locale.ROOT, "%.6f", ((Number) cell).doubleValue()); + return cell.toString(); + } +}