[Analytics-Engine] Decorrelate RexSubQuery before marking in PlannerImpl

[RyanL1997]

Description

PPL queries that contain EXISTS / NOT EXISTS / IN / SOME / ANY subqueries — and the subsearch shapes that lower to them — currently crash on the analytics-engine route with:

HTTP/1.1 500 Internal Server Error
{
  "error": {
    "reason": "There was internal problem at backend",
    "details": "Unrecognized filter operator [EXISTS / EXISTS]",
    "type": "IllegalStateException"
  }
}

The whole class CalcitePPLExistsSubqueryIT (19 tests) fails this way against an analytics-engine-routed cluster.

Root cause

The full failure path:

PPL `... | where [exists ...]`
  → SQL plugin's CalciteRelNodeVisitor.analyze(...)
  → LogicalFilter(condition=[EXISTS({...subquery...})])           ← RexSubQuery of SqlKind.EXISTS
  → AnalyticsExecutionEngine → analytics-engine PlannerImpl
  → marking phase → OpenSearchFilterRule.resolveViableBackends(predicate)
  → ScalarFunction.fromSqlOperatorWithFallback(EXISTS) returns null
  → throw IllegalStateException("Unrecognized filter operator [EXISTS / EXISTS]")

Where:

• Error site: OpenSearchFilterRule.java:155-160 — resolves every leaf predicate through ScalarFunction.fromSqlOperatorWithFallback, which doesn't (and shouldn't) cover RexSubQuery operators. RexSubQuery extends RexCall, so the upstream instanceof RexCall check waves it through.
• Why the SQL legacy engine doesn't see this: the legacy path goes through RelRunner.prepareStatement(rel), which internally runs SubQueryRemoveRule + decorrelation inside Calcite's Volcano-driven physical conversion. The analytics-engine path goes AnalyticsExecutionEngine → QueryPlanExecutor.execute → PlannerImpl.runAllOptimizations, bypassing that pipeline.
• Where the rewrite belongs: the analytics-engine's PlannerImpl already runs four HEP phases (reduce-expressions → pushdown-rules → aggregate-decompose → marking) plus a Volcano CBO pass. Subquery removal is a logical-plan rewrite — it belongs in this pipeline, before marking lowers the surrounding Logical* nodes into OpenSearch* shapes that the capability rules expect.

Fix

Add a new subquery-remove phase at the top of PlannerImpl.runAllOptimizations:

private static RelNode removeSubQueries(RelNode input, RuleProfilingListener listener) {
    HepProgramBuilder builder = new HepProgramBuilder();
    builder.addRuleCollection(
        List.of(
            CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
            CoreRules.PROJECT_SUB_QUERY_TO_CORRELATE,
            CoreRules.JOIN_SUB_QUERY_TO_CORRELATE
        )
    );
    HepPlanner planner = new HepPlanner(builder.build());
    // ... profiling-listener wiring identical to other phases ...
    planner.setRoot(input);
    RelNode withCorrelates = planner.findBestExp();
    return RelDecorrelator.decorrelateQuery(
        withCorrelates,
        RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)
    );
}

Wired in as:

RelNode modifiedRelNode = rawRelNode;
modifiedRelNode = removeSubQueries(modifiedRelNode, listener);    // NEW: Phase 0
modifiedRelNode = reduceExpressions(modifiedRelNode, listener);
modifiedRelNode = pushdownRules(modifiedRelNode, listener);
modifiedRelNode = decomposeAggregates(modifiedRelNode, listener);
modifiedRelNode = mark(modifiedRelNode, context, listener);
modifiedRelNode = cbo(modifiedRelNode, rawRelNode, context, listener);

The three *_SUB_QUERY_TO_CORRELATE rules rewrite RexSubQuerys into LogicalCorrelate. RelDecorrelator.decorrelateQuery then converts the correlates into straight joins. Every later phase observes a subquery-free tree.

Placement discussion

Per the Component Responsibilities table in opensearch-project/sql#5246, logical-plan rewrites are an SQL-plugin concern in the long-term design. The unified-query path on the SQL side, however, doesn't have a logical-optimizer stage yet — CalciteToolsHelper.HEP_PROGRAM (FilterMerge + PPLSimplifyDedup) runs only on the legacy engine path, and UnifiedQueryPlanner.plan() returns the RelNode after postAnalysisRules() shuttles. AE already owns RelNode → physical-execution.

Placing the rewrite here:

1. Every AE client (PPL today, ANSI SQL / Spark / future frontends) gets correct subquery handling without each frontend having to remember to run SubQueryRemoveRule + RelDecorrelator before handoff.
2. The fix is local to the consumer that actually had the limitation (analytics-engine), not spread across N callers.
3. When a UnifiedQueryOptimizer lands SQL-side per Revert changes in AbstractPointGeometryFieldMapper #5246, this phase can migrate upward — same code, different host.

After-fix verification

Same CalcitePPLExistsSubqueryIT class, same cluster, same routing (-Dtests.analytics.parquet_indices=true), with a vanilla-main SQL plugin (no SQL-side rewrite — isolates the AE change):

Run	Pass	Fail	Failure signature
Before	0 / 19	19 / 19	All IllegalStateException: Unrecognized filter operator [EXISTS / EXISTS]
After (this PR)	6 / 19	13 / 19	All AssertionError (result-row mismatches) — zero remaining IllegalStateExceptions

The "Unrecognized filter operator [EXISTS / EXISTS]" signature is gone from every test — the RCA is fully addressed.

Per-test breakdown after the fix:

• PASS (6): testEmptyExistsSubquery, testExistsSubqueryWithConjunction, testNotExistsSubquery, testNotExistsSubqueryInFilter, testSubsearchMaxOut2, testSubsearchMaxOut4
• FAIL — result mismatch only (13): testSimpleExistsSubquery, testSimpleExistsSubqueryInFilter, testUncorrelatedExistsSubquery, testUncorrelatedExistsSubqueryCheckTheReturnContentOfInnerTableIsEmptyOrNot, testNestedExistsSubquery, testExistsSubqueryAndAggregation, testIssue3566, testSubsearchMaxOut1, testSubsearchMaxOut3, testSubsearchMaxOutNegativeMeansUnlimited, testSubsearchMaxOutUncorrelated, testCorrelatedSubsearchMaxOutZeroMeansUnlimited, testUncorrelatedSubsearchMaxOutZeroMeansUnlimited

The decorrelated plan for the failing tests looks correct on inspection — e.g. testSimpleExistsSubquery produces:

LogicalJoin(condition=[=(worker.id, $f0)], joinType=[inner])
  LogicalTableScan(worker)
  LogicalProject(uid=[$0], $f1=[true])
    LogicalAggregate(group=[{0}])         ← dedups inner side by uid
      LogicalProject(uid=[$3])
        LogicalFilter(condition=[IS NOT NULL($3)])
          LogicalTableScan(work_information)

— textbook EXISTS semantics (INNER join with the inner side group-by-deduplicated). The 2× row multiplication that's observed suggests the bare-group-by LogicalAggregate(group=[{0}]) isn't preserving its dedup semantic through the analytics-engine's join + aggregate execution path. That's a separate downstream issue from the subquery-removal RCA and will be tracked under its own change.

Out of scope

Anything beyond the EXISTS / IN / SOME / ANY subquery-removal RCA:

• Bare-group-by Aggregate dedup semantics inside a decorrelated join — needs analytics-engine planner / executor changes. Tracked separately.
• RelDecorrelator will still raise for correlated subqueries Calcite can't decorrelate (e.g. correlated nondeterministic functions). Those queries used to crash with "Unrecognized filter operator"; with this PR they'll crash with a Calcite decorrelation message — net same outcome (failure → failure), different error string.

Tests

• New unit-test class SubQueryPlanShapeTests — 3 tests covering uncorrelated EXISTS, correlated EXISTS, and IN-subquery. Each parses real SQL through SqlPlannerTestFixture, runs the full planner, and asserts no RexSubQuery survives.
• RuleProfilingListenerTests — EXPECTED_PHASES extended with subquery-remove so the profile-listener phase-list assertion still matches the actual phase ordering. All 5 tests in that class still green.
• :sandbox:plugins:analytics-engine:test — 311 tests, 0 failures, 0 errors.
• End-to-end via CalcitePPLExistsSubqueryIT against an analytics-engine-routed cluster — 6/19 pass, 13/19 fail with downstream result-mismatch (see verification table above).

Companion / supersedes

This PR is the AE-side alternative to opensearch-project/sql#5460, which applied the same rewrite in the SQL plugin's AnalyticsExecutionEngine. After alignment in the #5460 thread, the rewrite is being moved here per the Component Responsibilities placement in opensearch-project/sql#5246 (subquery removal is logical optimization → SQL-plugin column long-term, but AE today is the actual consumer with the limitation, and every AE client benefits from a single fix here). #5460 will close as superseded once this lands.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit a8b1151)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Possible Issue RelDecorrelator.decorrelateQuery is called with a RelBuilder created from Contexts.empty(). If decorrelation requires access to schema metadata, type factories, or other context-dependent services that would normally be present in a properly configured context, this may cause failures. The empty context might lack necessary configuration that RelDecorrelator expects when rewriting complex correlated subqueries. return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null));

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to a8b1151

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Handle decorrelation exceptions properly The decorrelation step could fail with an exception, but the finally block will still execute endPhase. If decorrelation throws, the phase timing will be recorded as complete even though the operation failed. Consider wrapping the decorrelation in proper error handling or moving endPhase before the return to ensure accurate profiling. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [142-150] try { planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); // RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a // straight join shape that the marking + capability rules already handle. - return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); -} finally { + RelNode result = RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); if (listener != null) listener.endPhase("subquery-remove"); + return result; +} catch (Exception e) { + if (listener != null) listener.endPhase("subquery-remove"); + throw e; } Suggestion importance[1-10]: 4 __ Why: While the concern about exception handling is valid, the suggestion changes the profiling semantics by calling endPhase twice in different paths. The current finally block ensures endPhase is always called, which is a common pattern. The improvement is marginal and the suggested approach is debatable.	Low

---

Previous suggestions

Suggestions up to commit fdf67ea

Category	Suggestion	Impact
Possible issue	Validate planner result before decorrelation If planner.findBestExp() returns null or the decorrelation fails, the method will throw an unchecked exception. Add validation to ensure withCorrelates is non-null before decorrelation and handle potential decorrelation failures gracefully. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [143-147] planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); +if (withCorrelates == null) { + throw new IllegalStateException("Subquery removal phase failed to produce a valid plan"); +} // RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a // straight join shape that the marking + capability rules already handle. return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); Suggestion importance[1-10]: 5 __ Why: Adding a null check for withCorrelates improves robustness by catching potential planner failures early. However, HepPlanner.findBestExp() typically doesn't return null in normal operation, making this a defensive check rather than addressing a critical issue.	Low

Suggestions up to commit 05c262a

Category	Suggestion	Impact
General	Add error handling for decorrelation If decorrelateQuery throws an exception, the listener's phase will end but the exception from decorrelation won't be properly logged or contextualized. Consider wrapping the decorrelation call in a try-catch to provide better error context before re-throwing. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [137-150] HepPlanner planner = new HepPlanner(builder.build()); if (listener != null) { planner.addListener(listener); listener.beginPhase("subquery-remove"); } try { planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); - ... - return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); + try { + return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); + } catch (Exception e) { + LOGGER.error("Failed to decorrelate query after subquery removal", e); + throw e; + } } finally { if (listener != null) listener.endPhase("subquery-remove"); } Suggestion importance[1-10]: 4 __ Why: While adding error logging could marginally improve debugging, the finally block already ensures proper cleanup. The suggestion adds minimal value since exceptions will propagate naturally, and the additional try-catch only adds logging without changing behavior. The impact is minor for maintainability.	Low

Suggestions up to commit 9f49af5

Category	Suggestion	Impact
Possible issue	Validate correlate conversion succeeded The decorrelation step may fail or produce incorrect results if withCorrelates still contains unresolved correlations or if the planner didn't successfully convert all subqueries. Add validation to check that withCorrelates is in the expected state before decorrelation to prevent silent failures. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [143-147] planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); +if (withCorrelates == null) { + throw new IllegalStateException("Subquery-to-correlate conversion failed"); +} // RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a // straight join shape that the marking + capability rules already handle. return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); Suggestion importance[1-10]: 7 __ Why: Adding a null check for withCorrelates is a reasonable defensive programming practice that could prevent silent failures if the planner fails to convert subqueries. This improves error handling and makes debugging easier, though findBestExp() likely doesn't return null in normal operation.	Medium
	Pass proper context to RelBuilder The RelBuilder is created with a null schema parameter, which may cause issues when the decorrelator needs to resolve table references or perform schema-dependent operations. Consider passing the appropriate schema context from the input node's cluster to ensure proper decorrelation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [147] -return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); +RelBuilder relBuilder = RelBuilder.proto(input.getCluster().getPlanner().getContext()).create(input.getCluster(), null); +return RelDecorrelator.decorrelateQuery(withCorrelates, relBuilder); Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern about passing Contexts.empty() to RelBuilder.proto(), which might cause issues during decorrelation. However, the improved code assumes input.getCluster().getPlanner().getContext() exists and is the correct context, which may not be accurate without deeper knowledge of the codebase. The suggestion is reasonable but speculative.	Low

Suggestions up to commit 2632a5d

Category	Suggestion	Impact
Possible issue	Add null check for optimization result If planner.findBestExp() returns null or fails to find a valid plan, the subsequent decorrelateQuery call will fail with a NullPointerException. Add a null check or validation to handle cases where optimization doesn't produce a valid result. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [143-147] planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); +if (withCorrelates == null) { + throw new IllegalStateException("Subquery removal phase failed to produce a valid plan"); +} // RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a // straight join shape that the marking + capability rules already handle. return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); Suggestion importance[1-10]: 6 __ Why: Valid defensive programming suggestion. Adding a null check for withCorrelates before calling decorrelateQuery would prevent potential NullPointerException and provide clearer error messaging if the optimization phase fails unexpectedly.	Low

Suggestions up to commit b37e0d2

Category	Suggestion	Impact
General	Validate subquery removal before decorrelation The decorrelation step could fail or produce unexpected results if withCorrelates contains unresolved correlations. Add validation to check if the plan after findBestExp() still contains RexSubQuery nodes before attempting decorrelation, and log a warning or throw an exception if subqueries remain. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [143-147] planner.setRoot(input); RelNode withCorrelates = planner.findBestExp(); +// Validate that all RexSubQuery nodes were converted to LogicalCorrelate +if (containsRexSubQuery(withCorrelates)) { + LOGGER.warn("RexSubQuery nodes remain after correlation phase"); +} // RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a // straight join shape that the marking + capability rules already handle. return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); Suggestion importance[1-10]: 5 __ Why: Adding validation to check for remaining RexSubQuery nodes before decorrelation is a reasonable defensive programming practice. However, the suggestion references a containsRexSubQuery() method that doesn't exist in the codebase, and the test file SubQueryPlanShapeTests.java already provides comprehensive validation that subqueries are removed, making this additional check somewhat redundant.	Low
	Pass proper context to RelBuilder The RelBuilder is created with a null schema parameter, which may cause issues when decorrelating queries that reference specific schemas. Consider passing the appropriate schema from the input node's cluster to ensure proper schema resolution during decorrelation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java [147] -return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)); +RelBuilder relBuilder = RelBuilder.proto(input.getCluster().getPlanner().getContext()).create(input.getCluster(), null); +return RelDecorrelator.decorrelateQuery(withCorrelates, relBuilder); Suggestion importance[1-10]: 3 __ Why: The suggestion proposes using input.getCluster().getPlanner().getContext() instead of Contexts.empty(), but this assumes the planner context is more appropriate without evidence from the PR that the current implementation causes issues. The change is speculative and may not provide meaningful improvement.	Low

[github-actions]

Persistent review updated to latest commit d1140a5

[github-actions]

Persistent review updated to latest commit b37e0d2

[github-actions]

❌ Gradle check result for b37e0d2: null

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[songkant-aws]

Thanks for landing the subquery decorrelation work, @ryanbogan! This is a solid foundation — RelDecorrelator running before marking unblocks a lot of correlated-subquery shapes that we couldn't plan before.

I've opened a draft PR #21804 that builds on top of this one with two follow-up marking-phase fixes (LITERAL_AGG lowering for NOT IN shapes, and relaxing OpenSearchJoinRule to accept all join condition shapes so RelDecorrelator's pure non-equi join output reaches DataFusion). #21804 will stay draft until #21795 merges, then rebase onto main so only the layered diff shows. Looking forward to your review on this one.

[github-actions]

Persistent review updated to latest commit 2632a5d

[github-actions]

✅ Gradle check result for 2632a5d: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.36%. Comparing base (0c60e50) to head (a8b1151).

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21795      +/-   ##
============================================
+ Coverage     73.34%   73.36%   +0.02%     
- Complexity    75417    75429      +12     
============================================
  Files          6032     6032              
  Lines        342404   342404              
  Branches      49235    49235              
============================================
+ Hits         251142   251218      +76     
+ Misses        71272    71266       -6     
+ Partials      19990    19920      -70     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 9f49af5

[github-actions]

❕ Gradle check result for 9f49af5: UNSTABLE

Please review all flaky tests that succeeded after retry and create an issue if one does not already exist to track the flaky failure.

[github-actions]

Persistent review updated to latest commit 05c262a

[github-actions]

❌ Gradle check result for 05c262a: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit fdf67ea

[github-actions]

✅ Gradle check result for fdf67ea: SUCCESS

[github-actions]

Persistent review updated to latest commit a8b1151

[github-actions]

✅ Gradle check result for a8b1151: SUCCESS