Allow more than 2 alternative plans for dedup

yuancu · yuancu · commit 35ac017e5526 · 2026-01-13T18:41:25.000+08:00
Signed-off-by: Yuanchun Shen &lt;yuanchu@amazon.com&gt;
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
@@ -22,6 +22,7 @@
 import static org.opensearch.sql.util.MatcherUtils.verifyErrorMessageContains;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Locale;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -2061,11 +2062,11 @@ public void testDedupExpr() throws IOException {
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | eval new_gender = lower(gender) | eval"
                 + " new_state = lower(state) | dedup 2 new_gender, new_state"));
-    expected = loadExpectedPlan("explain_dedup_expr4.yaml");
-    alternative = loadExpectedPlan("explain_dedup_expr4_alternative.yaml");
     assertYamlEqualsIgnoreId(
-        expected,
-        alternative,
+        List.of(
+            loadExpectedPlan("explain_dedup_expr4.yaml"),
+            loadExpectedPlan("explain_dedup_expr4_alternative.yaml"),
+            loadExpectedPlan("explain_dedup_expr4_alternative1.yaml")),
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval new_gender = lower(gender) | eval new_state = lower(state) | sort gender,"
@@ -2101,11 +2102,11 @@ public void testDedupRename() throws IOException {
             "source=opensearch-sql_test_index_account | eval tmp_gender = lower(gender) | eval"
                 + " tmp_state = lower(state) | rename tmp_gender as new_gender | rename tmp_state"
                 + " as new_state | dedup 2 new_gender, new_state"));
-    expected = loadExpectedPlan("explain_dedup_expr4.yaml");
-    alternative = loadExpectedPlan("explain_dedup_expr4_alternative.yaml");
     assertYamlEqualsIgnoreId(
-        expected,
-        alternative,
+        List.of(
+            loadExpectedPlan("explain_dedup_expr4.yaml"),
+            loadExpectedPlan("explain_dedup_expr4_alternative.yaml"),
+            loadExpectedPlan("explain_dedup_expr4_alternative1.yaml")),
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval tmp_gender = lower(gender) | eval tmp_state = lower(state) | rename"
@@ -2164,11 +2165,12 @@ public void testDedupWithExpr() throws IOException {
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | eval new_gender = lower(gender) | eval"
                 + " new_state = lower(state) | dedup 2 age, account_number"));
-    expected = loadExpectedPlan("explain_dedup_with_expr4.yaml");
-    alternative = loadExpectedPlan("explain_dedup_with_expr4_alternative.yaml");
+
     assertYamlEqualsIgnoreId(
-        expected,
-        alternative,
+        List.of(
+            loadExpectedPlan("explain_dedup_with_expr4.yaml"),
+            loadExpectedPlan("explain_dedup_with_expr4_alternative.yaml"),
+            loadExpectedPlan("explain_dedup_with_expr4_alternative2.yaml")),
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval new_gender = lower(gender) | eval new_state = lower(state) | sort gender,"
diff --git a/integ-test/src/test/java/org/opensearch/sql/util/MatcherUtils.java b/integ-test/src/test/java/org/opensearch/sql/util/MatcherUtils.java
@@ -466,6 +466,39 @@ public static void assertYamlEqualsIgnoreId(
     }
   }
 
+  /**
+   * Compares actual YAML with a list of expected YAML strings, trying each in order until one
+   * matches. This is useful when the DSL implementation can produce multiple valid plan variants.
+   *
+   * <p>The method iterates through each expected YAML string and attempts a comparison with the
+   * actual YAML. If any expected YAML matches, the assertion succeeds. If none match, the method
+   * throws an {@link AssertionError} from the last failed comparison.
+   *
+   * <p>Both expected and actual YAML strings are cleaned up (ignoring RelNode IDs, timestamps,
+   * PIDs, etc.) before comparison.
+   *
+   * @param expectedYamls list of expected YAML strings to try (must not be empty)
+   * @param actualYaml the actual YAML string to compare
+   * @throws AssertionError if none of the expected YAMLs match the actual YAML (reports the last
+   *     failure)
+   * @throws IllegalArgumentException if expectedYamls is empty
+   */
+  public static void assertYamlEqualsIgnoreId(List<String> expectedYamls, String actualYaml) {
+    if (expectedYamls.isEmpty()) {
+      throw new IllegalArgumentException("expectedYamls list must not be empty");
+    }
+    AssertionError lastError = null;
+    for (String expectedYaml : expectedYamls) {
+      try {
+        assertYamlEquals(cleanUpYaml(expectedYaml), cleanUpYaml(actualYaml));
+        return;
+      } catch (AssertionError e) {
+        lastError = e;
+      }
+    }
+    throw lastError;
+  }
+
   public static void assertYamlEquals(String expected, String actual) {
     String normalizedExpected = normalizeLineBreaks(expected).trim();
     String normalizedActual = normalizeLineBreaks(actual).trim();
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_expr4_alternative1.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_expr4_alternative1.yaml
@@ -0,0 +1,12 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(sort0=[$1], sort1=[$3], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last], fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], gender=[$1], age=[$2], state=[$3], new_gender=[$4], new_state=[$5])
+        LogicalFilter(condition=[<=($6, 2)])
+          LogicalProject(account_number=[$0], gender=[$1], age=[$2], state=[$3], new_gender=[$4], new_state=[$5], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $4, $5)])
+            LogicalFilter(condition=[AND(IS NOT NULL($4), IS NOT NULL($5))])
+              LogicalSort(sort0=[$1], sort1=[$3], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last])
+                LogicalProject(account_number=[$0], gender=[$4], age=[$8], state=[$7], new_gender=[LOWER($4)], new_state=[LOWER($7)])
+                  CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0, 1},agg#0=LITERAL_AGG(2)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"new_gender":{"terms":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["gender.keyword"]}},"missing_bucket":false,"order":"asc"}}},{"new_state":{"terms":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["state.keyword"]}},"missing_bucket":false,"order":"asc"}}}]},"aggregations":{"$f2":{"top_hits":{"from":0,"size":2,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["account_number","gender","age","state"],"excludes":[]},"script_fields":{"new_gender":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["gender.keyword"]}},"ignore_failure":false},"new_state":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["state.keyword"]}},"ignore_failure":false}}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_with_expr4_alternative2.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_with_expr4_alternative2.yaml
@@ -0,0 +1,12 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(sort0=[$1], sort1=[$3], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last], fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], gender=[$1], age=[$2], state=[$3], new_gender=[$4], new_state=[$5])
+        LogicalFilter(condition=[<=($6, 2)])
+          LogicalProject(account_number=[$0], gender=[$1], age=[$2], state=[$3], new_gender=[$4], new_state=[$5], _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY $1, $3)])
+            LogicalFilter(condition=[AND(IS NOT NULL($1), IS NOT NULL($3))])
+              LogicalSort(sort0=[$1], sort1=[$3], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last])
+                LogicalProject(account_number=[$0], gender=[$4], age=[$8], state=[$7], new_gender=[LOWER($4)], new_state=[LOWER($7)])
+                  CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0, 1},agg#0=LITERAL_AGG(2)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":false,"order":"asc"}}},{"state":{"terms":{"field":"state.keyword","missing_bucket":false,"order":"asc"}}}]},"aggregations":{"$f2":{"top_hits":{"from":0,"size":2,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["gender","state","account_number","age"],"excludes":[]},"script_fields":{"new_state":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["state.keyword"]}},"ignore_failure":false},"new_gender":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQA/HsKICAib3AiOiB7CiAgICAibmFtZSI6ICJMT1dFUiIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfQogIF0KfQ==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0,"SOURCES":[0],"DIGESTS":["gender.keyword"]}},"ignore_failure":false}}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
