[analytics-datafusion] Apply aggregate mode in indexed executor, remove dc() gate (opensearch-project#22102)

The indexed executor now applies aggregate_mode stripping (Mode::Partial)
after building the physical plan, producing Binary HLL state instead of
Int64. This fixes the dc() schema mismatch without needing to suppress
filter delegation at marking time.

Changes:
- indexed_executor.rs: apply_aggregate_mode after physical plan build
- substrait_to_tree.rs: delegation_possible UDF returns first arg
  (transparent pass-through) instead of panicking
- Remove engine-native-merge gate from OpenSearchFilterRule
- Remove dead code: PlannerContext.hasEngineNativeMergeAggregate,
  PlannerImpl.containsEngineNativeMergeAggregate
- Expand FilterDelegationExtendedIT with stddev_pop/samp, percentile,
  additional dc() shapes (29 exact-value + 6 double-tolerance cases)

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

Author: sandeshkr419

sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs

@@ -319,7 +319,23 @@ mod tests {
             !rules.iter().any(|r| r.name() == combine_name),
             "CombinePartialFinalAggregate should be filtered out"
         );
-        // Verify we still have other rules
         assert!(!rules.is_empty(), "Should have other optimizer rules");
     }
+
+    /// Verifies apply_aggregate_mode(Partial) strips the Final aggregate and keeps
+    /// only the Partial subtree — the core behavior the indexed executor relies on
+    /// for engine-native-merge (dc/HLL) queries.
+    #[tokio::test]
+    async fn test_apply_partial_strips_final() {
+        let plan = make_agg_plan().await;
+        let display_before = plan_string(&plan);
+        assert!(display_before.contains("AggregateExec: mode=Final"), "expected Final in plan");
+        assert!(display_before.contains("AggregateExec: mode=Partial"), "expected Partial in plan");
+
+        let stripped = apply_aggregate_mode(plan, Mode::Partial).unwrap();
+        let display_after = plan_string(&stripped);
+        assert!(!display_after.contains("mode=Final"), "Final should be stripped");
+        assert!(display_after.contains("mode=Partial"), "Partial should remain");
+    }
+
 }

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs

@@ -500,6 +500,7 @@ async unsafe fn execute_indexed_with_context_inner(
 
     let query_config = Arc::new(handle.query_config);
     let num_partitions = query_config.target_partitions.max(1);
+    let aggregate_mode = handle.aggregate_mode;
     let ctx = handle.ctx;
     let table_name = handle.table_name;
     let table_path = handle.table_path;
@@ -880,6 +881,13 @@ async unsafe fn execute_indexed_with_context_inner(
     // types. The target is schema_coerce::coerce_inferred_schema(physical_schema) — same
     // narrowing the partition-stream registration uses, so consumer-side StreamingTable
     // and producer-side batches agree by construction (see crate::relabel_exec).
+    // Apply aggregate mode stripping when prepare_partial_plan was called (engine-native-merge).
+    // This makes the indexed executor produce Binary HLL state (Partial) instead of Int64 (Final).
+    let physical_plan = if aggregate_mode != crate::agg_mode::Mode::Default {
+        crate::agg_mode::apply_aggregate_mode(physical_plan, aggregate_mode)?
+    } else {
+        physical_plan
+    };
     let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema());
     let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?;
     log_debug!("DataFusion physical plan:\n{}", displayable(physical_plan.as_ref()).indent(true));

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/substrait_to_tree.rs

@@ -460,13 +460,24 @@ impl ScalarUDFImpl for DelegationPossibleUdf {
     }
     fn invoke_with_args(
         &self,
-        _args: ScalarFunctionArgs,
+        args: ScalarFunctionArgs,
     ) -> datafusion::common::Result<ColumnarValue> {
-        Err(datafusion::common::DataFusionError::Internal(format!(
-            "{} UDF body invoked — expr_to_bool_tree did not unwrap the marker; \
-                 treat as a serious correctness bug",
-            DELEGATION_POSSIBLE_FUNCTION_NAME
-        )))
+        // Pass-through: evaluate the original predicate (first arg) directly.
+        // Expected on the prepare_partial_plan path where FilterExec retains the
+        // delegation_possible wrapper. On the indexed path, expr_to_bool_tree should
+        // unwrap the marker before execution — log a warning if we reach here
+        // unexpectedly so correctness issues surface in logs without crashing.
+        log::warn!(
+            "delegation_possible UDF evaluated at runtime — expected only on \
+             prepare_partial_plan path; if this appears on a data-node indexed query, \
+             expr_to_bool_tree may have missed the marker"
+        );
+        args.args.into_iter().next().ok_or_else(|| {
+            datafusion::common::DataFusionError::Internal(format!(
+                "{} UDF invoked with no arguments",
+                DELEGATION_POSSIBLE_FUNCTION_NAME
+            ))
+        })
     }
 }
 

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java

@@ -33,7 +33,6 @@ public class PlannerContext {
     private final boolean preferMetadataDriver;
     private int annotationIdCounter;
     private RuleProfilingListener.PlannerProfile lastProfile;
-    private boolean hasEngineNativeMergeAggregate;
     // Cluster settings the planner consults at planning time (oversampling factor + delegation
     // block-list). Defaults to planner defaults; DefaultPlanExecutor injects the live, settings-backed
     // instance via setPlannerSettings before planning.
@@ -128,14 +127,6 @@ public OpenSearchDistributionTraitDef getDistributionTraitDef() {
         return distributionTraitDef;
     }
 
-    public boolean hasEngineNativeMergeAggregate() {
-        return hasEngineNativeMergeAggregate;
-    }
-
-    public void setHasEngineNativeMergeAggregate(boolean value) {
-        this.hasEngineNativeMergeAggregate = value;
-    }
-
     /**
      * Mirrors the {@code analytics.planner.prefer_metadata_driver} cluster setting at planning
      * time. When {@code false}, {@code OpenSearchTableScanRule} skips the permissive

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java

@@ -18,8 +18,6 @@
 import org.apache.calcite.rel.RelHomogeneousShuttle;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelShuttle;
-import org.apache.calcite.rel.core.Aggregate;
-import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.Sort;
@@ -54,7 +52,6 @@
 import org.opensearch.analytics.planner.rules.OpenSearchUnionRule;
 import org.opensearch.analytics.planner.rules.OpenSearchUnionSplitRule;
 import org.opensearch.analytics.planner.rules.OpenSearchValuesRule;
-import org.opensearch.analytics.spi.AggregateFunction;
 
 import java.util.List;
 import java.util.Optional;
@@ -102,7 +99,6 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con
         modifiedRelNode = reduceExpressions(modifiedRelNode, listener);
         modifiedRelNode = pushdownRules(modifiedRelNode, listener);
         modifiedRelNode = decomposeAggregates(modifiedRelNode, listener);
-        context.setHasEngineNativeMergeAggregate(containsEngineNativeMergeAggregate(modifiedRelNode));
         modifiedRelNode = mark(modifiedRelNode, context, listener);
         LOGGER.debug("After marking:\n{}", RelOptUtil.toString(modifiedRelNode));
         // TODO(combine-delegated-predicates): a post-marking HEP rule should fuse same-backend
@@ -341,25 +337,6 @@ private static RelNode decomposeAggregates(RelNode input, RuleProfilingListener
      * <p>TODO: add SortPushdown rule here — pushes Sort below Exchange to data nodes for top-K
      * optimization.
      */
-    /**
-     * True if the plan contains an aggregate with engine-native-merge semantics (intermediate
-     * wire type differs from output type, e.g. APPROX_COUNT_DISTINCT emits Binary HLL state).
-     * When present, filter delegation must be suppressed to avoid plan-shape divergence between
-     * derive_schema_from_partial_plan and the data-node execution.
-     */
-    private static boolean containsEngineNativeMergeAggregate(RelNode node) {
-        if (node instanceof Aggregate agg) {
-            for (AggregateCall call : agg.getAggCallList()) {
-                if (AggregateFunction.isEngineNativeMerge(call)) {
-                    return true;
-                }
-            }
-        }
-        for (RelNode child : node.getInputs()) {
-            if (containsEngineNativeMergeAggregate(child)) return true;
-        }
-        return false;
-    }
 
     private static RelNode mark(RelNode input, PlannerContext context, RuleProfilingListener listener) {
         return HepPhase.named("marking")

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java

@@ -254,19 +254,6 @@ private List<String> resolveViableBackends(
             viableSet.retainAll(registry.scalarBackendsAnyFormat(scalarFunc, returnType));
         }
 
-        // TODO: Temporary workaround — suppress dual-viable filters for dc() (engine-native-merge)
-        // queries. When delegation_possible UDF wraps a predicate in an APPROX_COUNT_DISTINCT plan,
-        // the shard execution path produces Int64 instead of the expected Binary HLL intermediate,
-        // causing a schema mismatch at the reduce sink. Proper fix requires changes to the DataFusion
-        // local execution path to honour the prepared partial plan when delegation_possible is present.
-        if (context.hasEngineNativeMergeAggregate() && viableSet.size() > 1) {
-            List<String> delegationAcceptors = registry.delegationAcceptors(DelegationType.FILTER);
-            boolean drivingBackendSurvives = viableSet.stream().anyMatch(b -> !delegationAcceptors.contains(b));
-            if (drivingBackendSurvives) {
-                viableSet.removeAll(delegationAcceptors);
-            }
-        }
-
         // Per-backend delegation block-list. Drop backends that block this predicate so it is never
         // marked viable for them — but only when a non-blocked backend survives: blocking is a
         // delegation knob and must never make a predicate unexecutable.

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/FilterDelegationExtendedIT.java

@@ -84,7 +84,30 @@ private record Case(String pplSuffix, long expected, boolean isRowCount) {}
 
         // ===== IN =====
         new Case("where str0 in ('apple', 'cherry') | stats count() as c", 5, false),
-        new Case("where str0 in ('apple', 'cherry') | stats sum(num0) as c", 210, false)
+        new Case("where str0 in ('apple', 'cherry') | stats sum(num0) as c", 210, false),
+
+        // ===== AVG — additional exact-integer cases =====
+        new Case("where str0 != 'cherry' | stats avg(num0) as c", 50, false),
+
+        // ===== dc() across more predicate shapes =====
+        new Case("where str0 != 'banana' | stats dc(str0) as c", 3, false),
+        new Case("where isnull(str1) | stats dc(num0) as c", 4, false),
+        new Case("where isnotnull(str1) | stats dc(num0) as c", 6, false),
+
+        // ===== PERCENTILE =====
+        new Case("where str0 != 'apple' | stats percentile(num0, 50) as c", 70, false),
+        new Case("where isnull(str1) | stats percentile(num0, 50) as c", 70, false)
+    );
+
+    private record StatisticalCase(String pplSuffix, double expected) {}
+
+    private static final List<StatisticalCase> STATISTICAL_CASES = List.of(
+        new StatisticalCase("where str0 != 'apple' | stats stddev_pop(num0) as c", 20.0),
+        new StatisticalCase("where str0 != 'apple' | stats stddev_samp(num0) as c", 21.602469),
+        new StatisticalCase("where isnull(str1) | stats stddev_pop(num0) as c", 25.860201),
+        new StatisticalCase("where isnull(str1) | stats stddev_samp(num0) as c", 29.860788),
+        new StatisticalCase("where isnotnull(str1) | stats stddev_pop(num0) as c", 27.487371),
+        new StatisticalCase("where isnotnull(str1) | stats stddev_samp(num0) as c", 30.110906)
     );
 
     private static boolean dataProvisioned = false;
@@ -116,6 +139,28 @@ public void testFilterDelegationExtended() throws Exception {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    public void testStatisticalAggregatesWithDelegation() throws Exception {
+        List<String> failures = new java.util.ArrayList<>();
+        for (StatisticalCase dc : STATISTICAL_CASES) {
+            try {
+                String ppl = "source = " + DATASET.indexName + " | " + dc.pplSuffix;
+                Map<String, Object> result = executePpl(ppl);
+                List<List<Object>> rows = (List<List<Object>>) result.get("datarows");
+                assertNotNull("datarows null for [" + dc.pplSuffix + "]", rows);
+                assertEquals("expected 1 agg row for [" + dc.pplSuffix + "]", 1, rows.size());
+                double actual = ((Number) rows.get(0).get(0)).doubleValue();
+                assertEquals("[" + dc.pplSuffix + "]", dc.expected, actual, 0.1);
+            } catch (AssertionError | Exception e) {
+                failures.add(dc.pplSuffix + " → " + e.getMessage());
+            }
+        }
+        if (!failures.isEmpty()) {
+            fail("Statistical aggregate failures (" + failures.size() + " of " + STATISTICAL_CASES.size() + "):\n"
+                + String.join("\n", failures));
+        }
+    }
+
     @SuppressWarnings("unchecked")
     private void assertAggCount(String pplSuffix, long expected) throws Exception {
         String ppl = "source = " + DATASET.indexName + " | " + pplSuffix;