[Backport 2.19-dev] Support pushdown dedup with expression (#4957)

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -2044,7 +2044,7 @@ public void testComplexDedup() throws IOException {
                 + " dedup 2 gender, state"));
   }
 
-  @Ignore("https://github.com/opensearch-project/sql/issues/4789")
+  @Test
   public void testDedupExpr() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     String expected = loadExpectedPlan("explain_dedup_expr1.yaml");
@@ -2054,30 +2054,36 @@ public void testDedupExpr() throws IOException {
             "source=opensearch-sql_test_index_account | eval new_gender = lower(gender) | dedup 1"
                 + " new_gender"));
     expected = loadExpectedPlan("explain_dedup_expr2.yaml");
+    String alternative = loadExpectedPlan("explain_dedup_expr2_alternative.yaml");
     assertYamlEqualsIgnoreId(
         expected,
+        alternative,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval new_gender = lower(gender), new_state = lower(state) | dedup 1 new_gender,"
                 + " new_state"));
     expected = loadExpectedPlan("explain_dedup_expr3.yaml");
+    alternative = loadExpectedPlan("explain_dedup_expr3_alternative.yaml");
     assertYamlEqualsIgnoreId(
         expected,
+        alternative,
         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,
         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,"
                 + " -state | dedup 2 new_gender, new_state"));
   }
 
-  @Ignore("https://github.com/opensearch-project/sql/issues/4789")
+  @Test
   public void testDedupRename() throws IOException {
-    // rename changes nothing, reuse the same yaml files of testDedupExpr()
+    // rename changes nothing, reuse the same yaml files
     enabledOnlyWhenPushdownIsEnabled();
     String expected = loadExpectedPlan("explain_dedup_expr1.yaml");
     assertYamlEqualsIgnoreId(
@@ -2086,41 +2092,96 @@ public void testDedupRename() throws IOException {
             "source=opensearch-sql_test_index_account | eval tmp_gender = lower(gender) | rename"
                 + " tmp_gender as new_gender | dedup 1 new_gender"));
     expected = loadExpectedPlan("explain_dedup_expr2.yaml");
+    String alternative = loadExpectedPlan("explain_dedup_expr2_alternative.yaml");
     assertYamlEqualsIgnoreId(
         expected,
+        alternative,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval tmp_gender = lower(gender), tmp_state = lower(state) | rename tmp_gender"
                 + " as new_gender | rename tmp_state as new_state | dedup 1 new_gender,"
                 + " new_state"));
     expected = loadExpectedPlan("explain_dedup_expr3.yaml");
+    alternative = loadExpectedPlan("explain_dedup_expr3_alternative.yaml");
     assertYamlEqualsIgnoreId(
         expected,
+        alternative,
         explainQueryYaml(
             "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,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval tmp_gender = lower(gender) | eval tmp_state = lower(state) | rename"
                 + " tmp_gender as new_gender | rename tmp_state as new_state | sort gender,"
                 + " -state | dedup 2 new_gender, new_state"));
   }
 
-  @Ignore("SortExprIndexScanRule not work?")
-  public void testDedupRename2() throws IOException {
+  @Test
+  public void testRenameDedupThenSortExpr() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
-    String expected = loadExpectedPlan("explain_dedup_expr4.yaml");
+    String expected = loadExpectedPlan("explain_dedup_expr_complex1.yaml");
+    String alternative = loadExpectedPlan("explain_dedup_expr_complex1_alternative.yaml");
     assertYamlEqualsIgnoreId(
         expected,
+        alternative,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
                 + " eval tmp_gender = lower(gender) | eval tmp_state = lower(state) | rename"
                 + " tmp_gender as new_gender | rename tmp_state as new_state | sort new_gender,"
                 + " -new_state | dedup 2 new_gender, new_state"));
+    expected = loadExpectedPlan("explain_dedup_expr_complex2.yaml");
+    alternative = loadExpectedPlan("explain_dedup_expr_complex2_alternative.yaml");
+    assertYamlEqualsIgnoreId(
+        expected,
+        alternative,
+        explainQueryYaml(
+            "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
+                + " 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 | sort new_gender, -new_state"));
+  }
+
+  @Test
+  public void testDedupWithExpr() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    String expected = loadExpectedPlan("explain_dedup_with_expr1.yaml");
+    assertYamlEqualsIgnoreId(
+        expected,
+        explainQueryYaml(
+            "source=opensearch-sql_test_index_account | eval new_gender = lower(gender) | dedup 1"
+                + " age"));
+    expected = loadExpectedPlan("explain_dedup_with_expr2.yaml");
+    String alternative = loadExpectedPlan("explain_dedup_with_expr2_alternative.yaml");
+    assertYamlEqualsIgnoreId(
+        expected,
+        alternative,
+        explainQueryYaml(
+            "source=opensearch-sql_test_index_account | fields account_number, gender, age, state |"
+                + " eval new_gender = lower(gender), new_state = lower(state) | dedup 1 age,"
+                + " new_state"));
+    expected = loadExpectedPlan("explain_dedup_with_expr3.yaml");
+    alternative = loadExpectedPlan("explain_dedup_with_expr3_alternative.yaml");
+    assertYamlEqualsIgnoreId(
+        expected,
+        alternative,
+        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,
+        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,"
+                + " -state | dedup 2 gender, state"));
   }
 
   @Test

integ-test/src/test/java/org/opensearch/sql/util/MatcherUtils.java

@@ -421,6 +421,25 @@ public static void assertYamlEqualsIgnoreId(String expectedYaml, String actualYa
     assertYamlEquals(cleanUpYaml(expectedYaml), cleanUpYaml(actualYaml));
   }
 
+  /**
+   * Compare actual YAML with two expected YAML strings, using the second as a fallback. This is
+   * useful when the DSL implementation can produce multiple valid plan variants. If the first
+   * comparison fails, attempts the second comparison instead.
+   *
+   * @param expectedYaml1 the primary expected YAML string
+   * @param expectedYaml2 the fallback expected YAML string
+   * @param actualYaml the actual YAML string to compare
+   * @throws AssertionError if both comparisons fail (reports only the second failure)
+   */
+  public static void assertYamlEqualsIgnoreId(
+      String expectedYaml1, String expectedYaml2, String actualYaml) {
+    try {
+      assertYamlEquals(cleanUpYaml(expectedYaml1), cleanUpYaml(actualYaml));
+    } catch (AssertionError e) {
+      assertYamlEquals(cleanUpYaml(expectedYaml2), cleanUpYaml(actualYaml));
+    }
+  }
+
   public static void assertYamlEquals(String expected, String actual) {
     String normalizedExpected = normalizeLineBreaks(expected).trim();
     String normalizedActual = normalizeLineBreaks(actual).trim();

integ-test/src/test/resources/expectedOutput/calcite/explain_complex_sort_expr_pushdown_for_smj_w_max_option.yaml

@@ -18,4 +18,4 @@ calcite:
         EnumerableMergeJoin(condition=[=($0, $7)], joinType=[left])
           EnumerableCalc(expr#0=[{inputs}], expr#1=['(?<lastname>^[A-Z])'], expr#2=['lastname'], expr#3=[REX_EXTRACT($t0, $t1, $t2)], $f0=[$t3])
             CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[lastname], LIMIT->10000, SORT_EXPR->[REX_EXTRACT($0, '(?<lastname>^[A-Z])', 'lastname') ASCENDING NULLS_LAST]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["lastname"],"excludes":[]},"sort":[{"_script":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQC63sKICAib3AiOiB7CiAgICAibmFtZSI6ICJSRVhfRVhUUkFDVCIsCiAgICAia2luZCI6ICJPVEhFUl9GVU5DVElPTiIsCiAgICAic3ludGF4IjogIkZVTkNUSU9OIgogIH0sCiAgIm9wZXJhbmRzIjogWwogICAgewogICAgICAiZHluYW1pY1BhcmFtIjogMCwKICAgICAgInR5cGUiOiB7CiAgICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfSwKICAgIHsKICAgICAgImR5bmFtaWNQYXJhbSI6IDEsCiAgICAgICJ0eXBlIjogewogICAgICAgICJ0eXBlIjogIlZBUkNIQVIiLAogICAgICAgICJudWxsYWJsZSI6IHRydWUsCiAgICAgICAgInByZWNpc2lvbiI6IC0xCiAgICAgIH0KICAgIH0sCiAgICB7CiAgICAgICJkeW5hbWljUGFyYW0iOiAyLAogICAgICAidHlwZSI6IHsKICAgICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAgICJwcmVjaXNpb24iOiAtMQogICAgICB9CiAgICB9CiAgXSwKICAiY2xhc3MiOiAib3JnLm9wZW5zZWFyY2guc3FsLmV4cHJlc3Npb24uZnVuY3Rpb24uVXNlckRlZmluZWRGdW5jdGlvbkJ1aWxkZXIkMSIsCiAgInR5cGUiOiB7CiAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICJudWxsYWJsZSI6IHRydWUsCiAgICAicHJlY2lzaW9uIjogMjAwMAogIH0sCiAgImRldGVybWluaXN0aWMiOiB0cnVlLAogICJkeW5hbWljIjogZmFsc2UKfQ==\"}","lang":"opensearch_compounded_script","params":{"MISSING_MAX":true,"utcTimestamp": 0,"SOURCES":[0,2,2],"DIGESTS":["lastname","(?<lastname>^[A-Z])","lastname"]}},"type":"string","order":"asc"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
-          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0},agg#0=LITERAL_AGG(1)), LIMIT->50000, SORT->[6]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"lastname":{"terms":{"field":"lastname","missing_bucket":false,"order":"asc"}}}]},"aggregations":{"$f1":{"top_hits":{"from":0,"size":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["lastname","account_number","firstname","address","birthdate","gender","city","balance","employer","state","age","email","male"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=50000, pageSize=null, startFrom=0)])
+          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0},agg#0=LITERAL_AGG(1)), LIMIT->50000, SORT->[6]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"lastname":{"terms":{"field":"lastname","missing_bucket":false,"order":"asc"}}}]},"aggregations":{"$f1":{"top_hits":{"from":0,"size":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["lastname","account_number","firstname","address","birthdate","gender","city","balance","employer","state","age","email","male","_id","_index","_score","_maxscore","_sort","_routing"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=50000, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_complex1.yaml

@@ -7,4 +7,4 @@ calcite:
             LogicalFilter(condition=[IS NOT NULL($4)])
               CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
   physical: |
-    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0},agg#0=LITERAL_AGG(1)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","_source":{"includes":["account_number","firstname","address","balance","gender","city","employer","state","age","email","lastname"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":false,"order":"asc"}}}]},"aggregations":{"$f1":{"top_hits":{"from":0,"size":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["gender","account_number","firstname","address","balance","city","employer","state","age","email","lastname"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0},agg#0=LITERAL_AGG(1)), 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"}}}]},"aggregations":{"$f1":{"top_hits":{"from":0,"size":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["gender","account_number","firstname","address","balance","city","employer","state","age","email","lastname","_id","_index","_score","_maxscore","_sort","_routing"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_complex2.yaml

@@ -8,4 +8,4 @@ calcite:
               LogicalProject(account_number=[$0], gender=[$4], age=[$8], state=[$7])
                 CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
   physical: |
-    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, gender, age, state], AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0, 1},agg#0=LITERAL_AGG(1)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","_source":{"includes":["account_number","gender","age","state"],"excludes":[]},"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":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["gender","state","account_number","age"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0, 1},agg#0=LITERAL_AGG(1)), 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":1,"version":false,"seq_no_primary_term":false,"explain":false,"_source":{"includes":["gender","state","account_number","firstname","address","balance","city","employer","age","email","lastname","_id","_index","_score","_maxscore","_sort","_routing"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_dedup_complex3.yaml

@@ -7,4 +7,4 @@ calcite:
             LogicalFilter(condition=[AND(IS NOT NULL($4), IS NOT NULL($7))])
               CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
   physical: |
-    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], AGGREGATION->rel#:LogicalAggregate.NONE.[](input=LogicalProject#,group={0, 1},agg#0=LITERAL_AGG(2)), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","_source":{"includes":["account_number","firstname","address","balance","gender","city","employer","state","age","email","lastname"],"excludes":[]},"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","firstname","address","balance","city","employer","age","email","lastname"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
+    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","firstname","address","balance","city","employer","age","email","lastname","_id","_index","_score","_maxscore","_sort","_routing"],"excludes":[]},"script_fields":{}}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])