[Analytics Engine] CSE: evaluate a repeated aggregate sub-plan once instead of twice

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.
+     *
+     * <p>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.
+     *
+     * <p><b>Not an MPP setting</b>, 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.
+     *
+     * <p>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<Boolean> 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<Setting<?>> ALL_SETTINGS = List.of(
         MPP_ENABLED,
+        SUBPLAN_REUSE_ENABLED,
         BROADCAST_MAX_BYTES,
         MPP_SHUFFLE_PARTITIONS,
         MPP_SHUFFLE_RECV_TIMEOUT,

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

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<Integer, byte[]> captur
     private static void injectBroadcastsInPlace(Stage stage, Map<Integer, byte[]> capturedByBuildId) {
         if (stage.getFragment() != null) {
             List<OpenSearchBroadcastScan> scans = RelNodeUtils.findNodes(stage.getFragment(), OpenSearchBroadcastScan.class);
-            List<Map.Entry<Integer, byte[]>> 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<Integer, byte[]> 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<StagePlan> enriched = new ArrayList<>(stage.getPlanAlternatives().size());
                 for (StagePlan sp : stage.getPlanAlternatives()) {
                     List<InstructionNode> merged = new ArrayList<>(sp.instructions());
-                    for (Map.Entry<Integer, byte[]> e : toInject) {
+                    for (Map.Entry<Integer, byte[]> 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()));
                     }