[BugFix] Make dedup simplify operand order-independent (#5488)

Address review feedback on #5488: extend mayBeFilterFromBucketNonNull to accept the merged conjunction shape FilterMergeRule produces, so PPLSimplifyDedupRule fires regardless of whether FilterMergeRule has already merged the user where clause into the bucket-non-null filter. PPLSimplifyDedupRule.apply now splits the bottom filter into IS_NOT_NULL conjuncts on partition keys (absorbed into LogicalDedup semantics) and any remaining conjuncts (preserved as a separate filter below the new LogicalDedup), so a user predicate that was folded in is no longer dropped.

With the operand predicate order-independent, the HEP rule order is no longer a load-bearing invariant. Revert the addRuleCollection -> addRuleInstance change in CalciteToolsHelper.HEP_PROGRAM that the previous patch introduced. Replace the regression test that pinned the buggy rule order with one that asserts the user-visible contract: with where preceding dedup, a LogicalDedup is produced and the user predicate is preserved regardless of which order FilterMergeRule and PPLSimplifyDedupRule fire.

Signed-off-by: ryan-gh-bot <ryan-gh-bot@users.noreply.github.com>

Author: ryan-gh-bot

core/src/main/java/org/opensearch/sql/calcite/plan/rule/PPLSimplifyDedupRule.java

@@ -6,11 +6,14 @@
 package org.opensearch.sql.calcite.plan.rule;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.plan.RelRule;
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollations;
@@ -22,6 +25,7 @@
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.rex.RexWindow;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.type.SqlTypeName;
@@ -115,8 +119,41 @@ protected void apply(
 
     RelCollation inputCollation = extractCollationFromWindow(windows.get(0));
 
+    // Split the bucket-non-null filter into two parts:
+    //   1) IS_NOT_NULL conjuncts on a partition key — these are the bucket-non-null guards PPL
+    //      emits as part of the dedup pattern; LogicalDedup absorbs their semantics.
+    //   2) Everything else — for example, a user `where` predicate that FilterMergeRule may
+    //      have folded into the same conjunction, or a user IS_NOT_NULL filter on a non-partition
+    //      column. These must be preserved as a separate filter below the new LogicalDedup so
+    //      user-visible behavior is unchanged regardless of whether FilterMergeRule fired.
+    Set<Integer> partitionKeyIndices = new HashSet<>();
+    for (RexNode key : dedupColumns) {
+      if (key instanceof RexInputRef ref) {
+        partitionKeyIndices.add(ref.getIndex());
+      }
+    }
+    List<RexNode> bucketNonNullConjuncts = new ArrayList<>();
+    List<RexNode> remainingConjuncts = new ArrayList<>();
+    for (RexNode conjunct : RelOptUtil.conjunctions(bucketNonNullFilter.getCondition())) {
+      if (isNotNullOnPartitionKey(conjunct, partitionKeyIndices)) {
+        bucketNonNullConjuncts.add(conjunct);
+      } else {
+        remainingConjuncts.add(conjunct);
+      }
+    }
+    // Defensive: if no IS_NOT_NULL conjunct on a partition key is present, this filter is not
+    // actually a bucket-non-null filter — bail out without transforming. The loose operand
+    // predicate may have matched on an unrelated AND that happens to contain an IS_NOT_NULL on
+    // some other ref.
+    if (bucketNonNullConjuncts.isEmpty()) {
+      return;
+    }
+
     RelBuilder relBuilder = call.builder();
     relBuilder.push(bucketNonNullFilter.getInput());
+    if (!remainingConjuncts.isEmpty()) {
+      relBuilder.filter(RexUtil.composeConjunction(relBuilder.getRexBuilder(), remainingConjuncts));
+    }
     List<Pair<RexNode, String>> targetProjections =
         projectWithWindow.getNamedProjects().stream()
             .filter(p -> !p.getKey().isA(SqlKind.ROW_NUMBER))
@@ -134,6 +171,14 @@ protected void apply(
     call.transformTo(relBuilder.build());
   }
 
+  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());
+  }
+
   private static @Nullable RelCollation extractCollationFromWindow(RexWindow window) {
     if (window.orderKeys.isEmpty()) {
       return null;

core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java

@@ -36,6 +36,7 @@
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.time.Instant;
+import java.util.List;
 import java.util.Properties;
 import java.util.function.Consumer;
 import org.apache.calcite.adapter.enumerable.EnumerableConvention;
@@ -58,6 +59,7 @@
 import org.apache.calcite.plan.Convention;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptSchema;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelOptTable.ViewExpander;
@@ -549,25 +551,12 @@ public RelNode visit(TableScan scan) {
     }
   }
 
-  /**
-   * Try to optimize the plan by using HepPlanner.
-   *
-   * <p>Rules are added as separate {@code addRuleInstance} phases so each rule is run to a fixed
-   * point before the next phase begins. {@code PPLSimplifyDedupRule} must run before {@code
-   * FilterMergeRule}: PPL emits the dedup pattern with an {@code IS NOT NULL} bucket-non-null
-   * filter directly above the scan, and the simplify rule's operand requires that filter to remain
-   * a pure {@code IS NOT NULL} (or {@code AND} of {@code IS NOT NULL}s). When a user {@code where}
-   * clause precedes the dedup, the user's filter sits adjacent to the bucket-non-null filter; if
-   * {@code FilterMergeRule} fires first it folds the two into a single conjunction, the simplify
-   * rule's operand match fails, no {@link org.opensearch.sql.calcite.plan.rel.LogicalDedup} is
-   * produced, and dedup pushdown is silently disabled. See
-   * https://github.com/opensearch-project/sql/issues/5482.
-   */
+  /** Try to optimize the plan by using HepPlanner */
+  private static final List<RelOptRule> hepRuleList =
+      List.of(FilterMergeRule.Config.DEFAULT.toRule(), PPLSimplifyDedupRule.DEDUP_SIMPLIFY_RULE);
+
   private static final HepProgram HEP_PROGRAM =
-      new HepProgramBuilder()
-          .addRuleInstance(PPLSimplifyDedupRule.DEDUP_SIMPLIFY_RULE)
-          .addRuleInstance(FilterMergeRule.Config.DEFAULT.toRule())
-          .build();
+      new HepProgramBuilder().addRuleCollection(hepRuleList).build();
 
   public static RelNode optimize(RelNode plan, CalcitePlanContext context) {
     Util.discard(context);

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -785,17 +785,45 @@ static Mapping mapping(List<RexNode> rexNodes, RelDataType schema) {
     return Mappings.target(getSelectColumns(rexNodes), schema.getFieldCount());
   }
 
+  /**
+   * Loose structural check for the bucket-non-null filter PPL emits as part of the {@code dedup}
+   * pattern. Returns true for either:
+   *
+   * <ul>
+   *   <li>a pure {@code IS NOT NULL($ref)}, or
+   *   <li>any {@code AND} whose conjuncts contain at least one {@code IS NOT NULL($ref)}.
+   * </ul>
+   *
+   * <p>This intentionally accepts the un-merged shape ({@code IS NOT NULL} alone) and the merged
+   * shape that {@link org.apache.calcite.rel.rules.FilterMergeRule} produces when a user {@code
+   * where} clause precedes the dedup ({@code AND(IS NOT NULL($pk), <user predicate>)}). The
+   * concrete partition-key match — and the split-out of any user predicate that was folded in —
+   * happens inside {@code PPLSimplifyDedupRule#apply}, which has access to the partition keys
+   * declared by the {@code ROW_NUMBER} window above this filter. Keeping the operand predicate
+   * order-independent means the simplify rule fires regardless of whether {@code FilterMergeRule}
+   * has already run, so dedup pushdown is no longer load-bearing on HEP rule ordering.
+   *
+   * <p>Misclassification is bounded by the surrounding four-level operand chain: a project that
+   * does not contain the {@code _row_number_dedup_} column, above a filter whose condition is
+   * {@code _row_number_ <= N} and which does contain {@code _row_number_dedup_}, above a project
+   * that does contain {@code _row_number_dedup_} (with a {@code ROW_NUMBER} window function), above
+   * this filter. That signature is unique to the PPL-emitted dedup pattern, so a user-written
+   * filter such as {@code IS NOT NULL(x) AND ...} cannot accidentally match.
+   */
   static boolean mayBeFilterFromBucketNonNull(LogicalFilter filter) {
     RexNode condition = filter.getCondition();
-    return isNotNullOnRef(condition)
-        || (condition instanceof RexCall rexCall
-            && rexCall.getOperator().equals(SqlStdOperatorTable.AND)
-            && rexCall.getOperands().stream().allMatch(PlanUtils::isNotNullOnRef));
+    if (isNotNullOnRef(condition)) {
+      return true;
+    }
+    return condition instanceof RexCall rexCall
+        && rexCall.getOperator().equals(SqlStdOperatorTable.AND)
+        && rexCall.getOperands().stream().anyMatch(PlanUtils::isNotNullOnRef);
   }
 
-  private static boolean isNotNullOnRef(RexNode rex) {
+  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;
   }
 

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLDedupTest.java

@@ -362,41 +362,37 @@ public void testSortFieldProjectedAwayBeforeDedup() {
 
   /**
    * When a user {@code where} precedes {@code dedup}, the user filter sits adjacent to the
-   * bucket-non-null filter that PPL emits for the dedup pattern. The HEP optimizer must run {@link
-   * PPLSimplifyDedupRule} before {@link FilterMergeRule}; otherwise the merged condition breaks the
-   * simplify-rule's operand match and dedup falls through to the in-memory {@code ROW_NUMBER}
-   * window form, defeating dedup pushdown to the shard.
-   *
-   * <p>This test mirrors the production HEP rule sequence in {@code CalciteToolsHelper#HEP_PROGRAM}
-   * and asserts that a {@code LogicalDedup} is produced in the presence of a user {@code where}.
+   * bucket-non-null filter that PPL emits for the dedup pattern. {@link PPLSimplifyDedupRule} must
+   * fold the dedup pattern into a {@link org.opensearch.sql.calcite.plan.rel.LogicalDedup} so that
+   * {@code DedupPushdownRule} can match it; otherwise dedup falls through to the in-memory {@code
+   * ROW_NUMBER} window form, defeating dedup pushdown to the shard. The user predicate must be
+   * preserved as a separate filter below the new {@code LogicalDedup}.
    */
   @Test
   public void testDedupAfterWhereProducesLogicalDedup() {
     String ppl = "source=EMP | where SAL > 1000 | dedup DEPTNO";
     RelNode raw = getRelNodeRaw(ppl);
 
     // Sanity: the un-merged plan has both the bucket-non-null filter and the user where filter
-    // adjacent to each other above the scan — exactly the shape that triggers the bug.
+    // adjacent to each other above the scan — exactly the shape that triggered the original bug.
     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(">($5, 1000)"));
     Assert.assertTrue(
         "Raw plan should contain ROW_NUMBER prior to simplification:\n" + rawExplain,
         rawExplain.contains("ROW_NUMBER"));
 
-    // Apply the production HEP rule sequence: PPLSimplifyDedupRule first, then FilterMergeRule.
+    // Apply rules in the order PPLSimplifyDedupRule -> FilterMergeRule.
     HepProgram program =
         new HepProgramBuilder()
             .addRuleInstance(PPLSimplifyDedupRule.DEDUP_SIMPLIFY_RULE)
             .addRuleInstance(FilterMergeRule.Config.DEFAULT.toRule())
             .build();
-    HepPlanner planner = new HepPlanner(program);
-    planner.setRoot(raw);
-    RelNode optimized = planner.findBestExp();
+    RelNode optimized = runHepPlanner(raw, program);
 
     String optimizedExplain = optimized.explain();
     Assert.assertTrue(
@@ -406,40 +402,56 @@ public void testDedupAfterWhereProducesLogicalDedup() {
     Assert.assertFalse(
         "Optimized plan should not retain ROW_NUMBER after simplification:\n" + optimizedExplain,
         optimizedExplain.contains("ROW_NUMBER"));
+    Assert.assertTrue(
+        "User where predicate must be preserved as a filter below LogicalDedup:\n"
+            + optimizedExplain,
+        optimizedExplain.contains(">($5, 1000)"));
   }
 
   /**
-   * Companion to {@link #testDedupAfterWhereProducesLogicalDedup} that pins the buggy behavior:
-   * when {@link FilterMergeRule} runs before {@link PPLSimplifyDedupRule}, the merged conjunction
-   * defeats the simplify-rule's operand match and no {@code LogicalDedup} is produced. This guards
-   * against accidentally swapping the rule order in {@code CalciteToolsHelper#HEP_PROGRAM} back to
-   * the buggy ordering.
+   * Companion to {@link #testDedupAfterWhereProducesLogicalDedup} that pins the user-visible
+   * contract: with a {@code where} preceding {@code dedup}, a {@code LogicalDedup} must be produced
+   * regardless of the order in which {@link FilterMergeRule} and {@link PPLSimplifyDedupRule} fire.
+   *
+   * <p>The simplify rule's bucket-non-null operand predicate is order-independent — it accepts both
+   * a pure {@code IS NOT NULL} and an {@code AND} that contains an {@code IS NOT NULL} on a
+   * partition key — so {@code FilterMergeRule} firing first no longer disables dedup pushdown. This
+   * guards against re-introducing the original regression by reordering, removing, or adding rules
+   * to {@code CalciteToolsHelper#HEP_PROGRAM}.
    */
   @Test
-  public void testDedupAfterWhereWithFilterMergeFirstFailsToSimplify() {
+  public void testDedupAfterWhereProducesLogicalDedupRegardlessOfRuleOrder() {
     String ppl = "source=EMP | where SAL > 1000 | dedup DEPTNO";
     RelNode raw = getRelNodeRaw(ppl);
 
-    // Buggy order: FilterMergeRule first, then PPLSimplifyDedupRule.
+    // Swapped order: FilterMergeRule first, then PPLSimplifyDedupRule. This is the order that
+    // produced the original bug (#5482) before the fix.
     HepProgram program =
         new HepProgramBuilder()
             .addRuleInstance(FilterMergeRule.Config.DEFAULT.toRule())
             .addRuleInstance(PPLSimplifyDedupRule.DEDUP_SIMPLIFY_RULE)
             .build();
-    HepPlanner planner = new HepPlanner(program);
-    planner.setRoot(raw);
-    RelNode optimized = planner.findBestExp();
+    RelNode optimized = runHepPlanner(raw, program);
 
     String optimizedExplain = optimized.explain();
-    Assert.assertFalse(
-        "Buggy order: PPLSimplifyDedupRule should not match after FilterMergeRule consolidates"
-            + " the bucket-non-null filter with the user where filter, so no LogicalDedup is"
-            + " produced:\n"
+    Assert.assertTrue(
+        "Even with FilterMergeRule firing first, PPLSimplifyDedupRule must still produce"
+            + " LogicalDedup so DedupPushdownRule can match it:\n"
             + optimizedExplain,
         optimizedExplain.contains("LogicalDedup"));
+    Assert.assertFalse(
+        "ROW_NUMBER window form should be removed after simplification:\n" + optimizedExplain,
+        optimizedExplain.contains("ROW_NUMBER"));
     Assert.assertTrue(
-        "Buggy order: ROW_NUMBER window form should remain since simplification was skipped:\n"
+        "User where predicate must be preserved as a filter below LogicalDedup, even when"
+            + " FilterMergeRule fired first and folded it into the bucket-non-null filter:\n"
             + optimizedExplain,
-        optimizedExplain.contains("ROW_NUMBER"));
+        optimizedExplain.contains(">($5, 1000)"));
+  }
+
+  private static RelNode runHepPlanner(RelNode root, HepProgram program) {
+    HepPlanner planner = new HepPlanner(program);
+    planner.setRoot(root);
+    return planner.findBestExp();
   }
 }