[BugFix] Restore dedup pushdown when combined with WHERE clause (#5482)

[ryan-gh-bot]

Description

PPL dedup is initially compiled into a ROW_NUMBER() OVER (PARTITION BY field) window pattern with an IS_NOT_NULL(field) bucket-non-null filter directly above the scan. PPLSimplifyDedupRule is responsible for folding that pattern into a LogicalDedup operator, which DedupPushdownRule then pushes to the shard as a composite aggregation + top_hits. The HEP program registered both FilterMergeRule.DEFAULT and PPLSimplifyDedupRule.DEDUP_SIMPLIFY_RULE via addRuleCollection. When a user where precedes dedup, the user's filter sat adjacent to the bucket-non-null filter; FilterMergeRule fired first and merged them into AND(IS_NOT_NULL(field), <user predicate>), which no longer satisfied PPLSimplifyDedupRule's mayBeFilterFromBucketNonNull operand predicate. The simplify rule never fired, no LogicalDedup was produced, and DedupPushdownRule had nothing to match — dedup fell through to the in-memory ROW_NUMBER window form on the coordinator, which exceeds timeouts on large indices.

The fix replaces the single addRuleCollection with two sequential addRuleInstance phases so each rule runs to a fixed point in a defined order: PPLSimplifyDedupRule first (folds the dedup pattern into LogicalDedup while the bucket-non-null filter is still pure), then FilterMergeRule (consolidates any remaining adjacent filters). Per Calcite's documentation, addRuleCollection is non-deterministic in rule order, so addRuleInstance is the correct primitive to express this ordering invariant.

Verification

Before (current source, prior to this PR):

source=EMP | where SAL > 1000 | dedup DEPTNO against the existing HEP_PROGRAM ordering. The added regression test testDedupAfterWhereWithFilterMergeFirstFailsToSimplify exercises the exact rule sequence that production uses today (FilterMergeRule first, PPLSimplifyDedupRule second) on the un-merged plan PPL emits. The optimized plan retains the ROW_NUMBER window form and never produces a LogicalDedup, so DedupPushdownRule cannot match:

LogicalProject(EMPNO=[$0], ..., DEPTNO=[$7])
  LogicalFilter(condition=[<=($8, 1)])
    LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $7)])
      LogicalFilter(condition=[AND(IS NOT NULL($7), >($5, 1000))])
        LogicalTableScan(table=[[scott, EMP]])

After (this branch):

$ JAVA_HOME=/usr/lib/jvm/java-21-amazon-corretto ./gradlew :ppl:test \
    --tests org.opensearch.sql.ppl.calcite.CalcitePPLDedupTest

CalcitePPLDedupTest > testDedupAfterWhereProducesLogicalDedup PASSED
CalcitePPLDedupTest > testDedupAfterWhereWithFilterMergeFirstFailsToSimplify PASSED
CalcitePPLDedupTest > testDedup1 PASSED
CalcitePPLDedupTest > testDedup2 PASSED
CalcitePPLDedupTest > testDedupKeepEmpty1 PASSED
CalcitePPLDedupTest > testDedupKeepEmpty2 PASSED
CalcitePPLDedupTest > testDedupExpr PASSED
CalcitePPLDedupTest > testRenameDedup PASSED
CalcitePPLDedupTest > testSortThenDedup PASSED
CalcitePPLDedupTest > testSortThenDedupWithEval PASSED
CalcitePPLDedupTest > testSortFieldProjectedAwayBeforeDedup PASSED

BUILD SUCCESSFUL

The new testDedupAfterWhereProducesLogicalDedup test runs the same rule sequence used by CalciteToolsHelper#HEP_PROGRAM after this PR (PPLSimplifyDedupRule first, then FilterMergeRule) and asserts that the optimized plan contains LogicalDedup and no longer contains ROW_NUMBER. The optimized plan now folds the dedup pattern into a LogicalDedup operator with the user where predicate preserved as a separate filter below it, ready for DedupPushdownRule to match and rewrite into a composite + top_hits aggregation:

LogicalProject(EMPNO=[$0], ..., DEPTNO=[$7])
  LogicalDedup(dedup_fields=[$7], allowed_dedup=1, keepEmpty=false, consecutive=false)
    LogicalProject(EMPNO=[$0], ..., DEPTNO=[$7])
      LogicalFilter(condition=[>($5, 1000)])
        LogicalTableScan(table=[[scott, EMP]])

Dedup pushdown is restored when a where clause precedes dedup. All 11 existing CalcitePPLDedupTest cases continue to pass.

Related Issues

Resolves #5482

Check List

• New functionality includes testing.
• New functionality has been documented.
• New PPL command checklist all confirmed.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff or -s.
• Public documentation issue/PR created.

---

This pull request was authored by @ryan-gh-bot in response to the maintainer command /triage on #5482 (invoked by @RyanL1997). This is a maintainer-triggered automated contribution. Please review carefully before merging.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 0bf8f7c)

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 The isNotNullOnRef method does not check if getOperands() is empty before calling get(0). If a malformed IS_NOT_NULL call with zero operands reaches this code, it will throw an IndexOutOfBoundsException. While the operand predicate mayBeFilterFromBucketNonNull adds an isEmpty() guard at line 826, the standalone isNotNullOnRef at line 823 is also called from PPLSimplifyDedupRule line 175 without that guard, creating a potential crash path if a zero-operand IS_NOT_NULL appears in the partition key check loop. static boolean isNotNullOnRef(RexNode rex) { return rex instanceof RexCall rexCall && rexCall.isA(SqlKind.IS_NOT_NULL) && !rexCall.getOperands().isEmpty() && rexCall.getOperands().get(0) instanceof RexInputRef;

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 0bf8f7c

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Remove redundant operand check The !rexCall.getOperands().isEmpty() check is redundant because IS_NOT_NULL always has exactly one operand by definition. Accessing getOperands().get(0) without this check is safe and simplifies the logic. core/src/main/java/org/opensearch/sql/calcite/plan/rule/PPLSimplifyDedupRule.java [174-180] private static boolean isNotNullOnPartitionKey(RexNode rex, Set<Integer> partitionKeyIndices) { return rex instanceof RexCall rexCall && rexCall.isA(SqlKind.IS_NOT_NULL) - && !rexCall.getOperands().isEmpty() && rexCall.getOperands().get(0) instanceof RexInputRef ref && partitionKeyIndices.contains(ref.getIndex()); } Suggestion importance[1-10]: 4 __ Why: While IS_NOT_NULL typically has exactly one operand, the check !rexCall.getOperands().isEmpty() provides defensive programming against malformed expressions. Removing it offers minimal benefit and slightly reduces code robustness. The same pattern is used consistently in PlanUtils.isNotNullOnRef at line 826, suggesting it's an intentional defensive practice.	Low
	Handle non-reference partition keys Handle non-RexInputRef partition keys gracefully. If dedupColumns contains complex expressions that aren't simple column references, they won't be added to partitionKeyIndices, potentially causing valid bucket-non-null filters to be incorrectly rejected. Consider logging a warning or adding a comment explaining this limitation. core/src/main/java/org/opensearch/sql/calcite/plan/rule/PPLSimplifyDedupRule.java [130-134] for (RexNode key : dedupColumns) { if (key instanceof RexInputRef ref) { partitionKeyIndices.add(ref.getIndex()); } + // Note: Complex expressions in dedupColumns are not tracked as partition keys. + // This may cause the rule to bail out if bucket-non-null filters reference them. } Suggestion importance[1-10]: 3 __ Why: The suggestion correctly identifies that non-RexInputRef partition keys won't be tracked, but this appears to be intentional design rather than a bug. Adding a comment could improve code clarity, but the impact is minimal since the defensive check at line 148 already handles this case by bailing out when no matching filters are found.	Low

---

Previous suggestions

Suggestions up to commit 19f506d

Category	Suggestion	Impact
General	Make filter assertion more specific The assertion checking for "1000" is too generic and may match unintended occurrences in the plan. Consider checking for a more specific pattern like ">($7, 1000)" or "SAL > 1000" to ensure the user's WHERE clause filter is actually present in the expected location. ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLDedupTest.java [380-386] String rawExplain = raw.explain(); Assert.assertTrue( "Raw plan should contain the bucket-non-null filter:\n" + rawExplain, rawExplain.contains("IS NOT NULL")); Assert.assertTrue( "Raw plan should contain the user where filter:\n" + rawExplain, - rawExplain.contains("1000")); + rawExplain.contains(">") && rawExplain.contains("1000")); Suggestion importance[1-10]: 4 __ Why: The suggestion correctly identifies that checking for "1000" alone could match unintended occurrences. However, the improved code rawExplain.contains(">") && rawExplain.contains("1000") is still too generic and doesn't guarantee the filter is in the expected location. A more specific pattern would be better, but the suggestion does improve robustness slightly.	Low

Suggestions up to commit 8bfb649

Category	Suggestion	Impact
General	Make filter assertion more specific The assertion checking for "1000" is too generic and may match unintended occurrences in the plan. Consider checking for a more specific pattern like ">($7, 1000)" or "SAL > 1000" to ensure the user's WHERE clause filter is actually present in the expected location. ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLDedupTest.java [383-388] String rawExplain = raw.explain(); Assert.assertTrue( "Raw plan should contain the bucket-non-null filter:\n" + rawExplain, rawExplain.contains("IS NOT NULL")); Assert.assertTrue( "Raw plan should contain the user where filter:\n" + rawExplain, - rawExplain.contains("1000")); + rawExplain.contains(">") && rawExplain.contains("1000")); Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that checking for "1000" alone could match unintended occurrences. The improved code adds a check for ">" to make the assertion more specific, though it still doesn't guarantee the exact filter pattern. This is a reasonable improvement for test robustness but has moderate impact.	Low

[RyanL1997]

Hi @ryan-gh-bot , thanks for the change, left some reviews.

[RyanL1997]

@ryan-gh-bot can you explain the reason for removing this?

[ryan-gh-bot]

The hepRuleList was only ever used as the argument to addRuleCollection(...). The fix here depends on PPLSimplifyDedupRule running strictly before FilterMergeRule, but per Calcite's docs addRuleCollection is non-deterministic in rule firing order — passing the rules in a List does not guarantee they fire in list order. So I switched to two sequential addRuleInstance(...) phases, which is the documented primitive for expressing an ordering invariant (each phase runs to a fixed point before the next begins).

Once the rules were inlined into the builder chain, the List<RelOptRule> field became dead code, so I dropped it (and the now-unused java.util.List / org.apache.calcite.plan.RelOptRule imports). The rule set itself is unchanged — same two rules, just expressed in a way that pins the order.

[RyanL1997]

@ryan-gh-bot it is unnecessary for mentioning the specific issue link here. You can just remove this.

[ryan-gh-bot]

Done — pushed 19f506ddb.

[github-actions]

Persistent review updated to latest commit 19f506d

[RyanL1997]

@ryan-gh-bot the rule-reordering fix works but treats a symptom rather than the root cause. The actual brittleness is in PlanUtils.mayBeFilterFromBucketNonNull — its allMatch(isNotNullOnRef) requirement means any future code path that produces a merged conjunction (a new rule, a different optimization order, a SQL-side caller) will silently re-disable dedup pushdown the same way. The HEP ordering is now a load-bearing invariant with nothing enforcing it; if anyone later adds a third rule to HEP_PROGRAM or refactors it, the bug comes back silently. The new testDedupAfterWhereWithFilterMergeFirstFailsToSimplify test pins the implementation strategy rather than the user-visible behavior.

Please revise to take option 1 from the issue instead:

1. Extend mayBeFilterFromBucketNonNull (core/.../utils/PlanUtils.java) to accept AND(IS_NOT_NULL(partition_col), <rest>) — i.e., recognize an AND whose operands contain an IS_NOT_NULL on the partition column rather than requiring every conjunct to be one.
2. In PPLSimplifyDedupRule.apply, split the non-IS_NOT_NULL conjuncts off and emit them as a separate filter below the new LogicalDedup, so the user predicate is preserved and pushdown still works.
3. Revert the addRuleCollection → addRuleInstance change in CalciteToolsHelper.HEP_PROGRAM — once the predicate is order-independent, rule order no longer matters.
4. Replace testDedupAfterWhereWithFilterMergeFirstFailsToSimplify (which pins buggy behavior) with a test that asserts the user-visible contract: with where upstream, dedup still produces a LogicalDedup regardless of rule order. Keep testDedupAfterWhereProducesLogicalDedup.

Please sanity-check that the loosened operand can't accidentally match user-written IS_NOT_NULL(x) AND ... filters outside the dedup pattern — the four-level operand chain (project → filter with row_number <= N → project containing ROW_NUMBER-dedup → filter) should be specific enough, but worth confirming.

[ryan-gh-bot]

Working on /@triage (invoked by @RyanL1997) — posting result when done.

[ryan-gh-bot]

Done — pushed 0bf8f7c73.

[github-actions]

Persistent review updated to latest commit 0bf8f7c

[dai-chen]

@ryan-gh-bot the rule-reordering fix works but treats a symptom rather than the root cause. The actual brittleness is in PlanUtils.mayBeFilterFromBucketNonNull — its allMatch(isNotNullOnRef) requirement means any future code path that produces a merged conjunction (a new rule, a different optimization order, a SQL-side caller) will silently re-disable dedup pushdown the same way. The HEP ordering is now a load-bearing invariant with nothing enforcing it; if anyone later adds a third rule to HEP_PROGRAM or refactors it, the bug comes back silently. The new testDedupAfterWhereWithFilterMergeFirstFailsToSimplify test pins the implementation strategy rather than the user-visible behavior.

Please revise to take option 1 from the issue instead:

1. Extend mayBeFilterFromBucketNonNull (core/.../utils/PlanUtils.java) to accept AND(IS_NOT_NULL(partition_col), <rest>) — i.e., recognize an AND whose operands contain an IS_NOT_NULL on the partition column rather than requiring every conjunct to be one.
2. In PPLSimplifyDedupRule.apply, split the non-IS_NOT_NULL conjuncts off and emit them as a separate filter below the new LogicalDedup, so the user predicate is preserved and pushdown still works.
3. Revert the addRuleCollection → addRuleInstance change in CalciteToolsHelper.HEP_PROGRAM — once the predicate is order-independent, rule order no longer matters.
4. Replace testDedupAfterWhereWithFilterMergeFirstFailsToSimplify (which pins buggy behavior) with a test that asserts the user-visible contract: with where upstream, dedup still produces a LogicalDedup regardless of rule order. Keep testDedupAfterWhereProducesLogicalDedup.

Please sanity-check that the loosened operand can't accidentally match user-written IS_NOT_NULL(x) AND ... filters outside the dedup pattern — the four-level operand chain (project → filter with row_number <= N → project containing ROW_NUMBER-dedup → filter) should be specific enough, but worth confirming.

I'm weighing the pros and cons of these two options:

• Option 1 — make PPLSimplifyDedupRule run first: If this rule runs before FilterMergeRule, it matches the original unmerged pattern and fixes the problem directly. The concern is that some other rule could run ahead of it and reintroduce the mismatch.
• Option 2 - current PR approach: Does this cover all the plan shapes visitDedupe can produce, or does it only fix the plain where + dedup case? e.g. dedup N, multiple dedup fields, and the sort/collation cases. Or these cases are not supported by DSL query anyway?

[penghuo]

1. Add end-to-end test asserts dedup is actually pushed down when where precedes dedup.
2. Refactor PR description to match the implementation.

[penghuo]

could u simply AI comments.

[penghuo]

nit: not necessary change, previous code is clean.

[penghuo]

nit: The added !rexCall.getOperands().isEmpty() before .get(0) can never be false — IS NOT NULL is always unary in Calcite.

[penghuo]

Reuse isNotNullOnRef