Pushdown the top rare commands to nested aggregation (opensearch-project#4707)

* Optimize the top rare commands to nested aggregation

Signed-off-by: Lantao Jin <ltjin@amazon.com>

* Fix cost computation

Signed-off-by: Lantao Jin <ltjin@amazon.com>

---------

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -15,9 +15,9 @@
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_DEDUP;
-import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_MAIN;
-import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_SUBSEARCH;
-import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_TOP_RARE;
+import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_MAIN;
+import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_RARE_TOP;
+import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_SUBSEARCH;
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRelation;
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRexCall;
 import static org.opensearch.sql.calcite.utils.PlanUtils.transformPlanToAttachChild;
@@ -1643,7 +1643,7 @@ public RelNode visitAppendCol(AppendCol node, CalcitePlanContext context) {
             List.of(),
             WindowFrame.toCurrentRow());
     context.relBuilder.projectPlus(
-        context.relBuilder.alias(mainRowNumber, ROW_NUMBER_COLUMN_NAME_MAIN));
+        context.relBuilder.alias(mainRowNumber, ROW_NUMBER_COLUMN_FOR_MAIN));
 
     // 3. build subsearch tree (attach relation to subsearch)
     UnresolvedPlan relation = getRelation(node);
@@ -1661,7 +1661,7 @@ public RelNode visitAppendCol(AppendCol node, CalcitePlanContext context) {
             List.of(),
             WindowFrame.toCurrentRow());
     context.relBuilder.projectPlus(
-        context.relBuilder.alias(subsearchRowNumber, ROW_NUMBER_COLUMN_NAME_SUBSEARCH));
+        context.relBuilder.alias(subsearchRowNumber, ROW_NUMBER_COLUMN_FOR_SUBSEARCH));
 
     List<String> subsearchFields = context.relBuilder.peek().getRowType().getFieldNames();
     List<String> mainFields = context.relBuilder.peek(1).getRowType().getFieldNames();
@@ -1675,17 +1675,17 @@ public RelNode visitAppendCol(AppendCol node, CalcitePlanContext context) {
     // 7. join with condition `_row_number_main_ = _row_number_subsearch_`
     RexNode joinCondition =
         context.relBuilder.equals(
-            context.relBuilder.field(2, 0, ROW_NUMBER_COLUMN_NAME_MAIN),
-            context.relBuilder.field(2, 1, ROW_NUMBER_COLUMN_NAME_SUBSEARCH));
+            context.relBuilder.field(2, 0, ROW_NUMBER_COLUMN_FOR_MAIN),
+            context.relBuilder.field(2, 1, ROW_NUMBER_COLUMN_FOR_SUBSEARCH));
     context.relBuilder.join(
         JoinAndLookupUtils.translateJoinType(Join.JoinType.FULL), joinCondition);
 
     if (!node.isOverride()) {
       // 8. if override = false, drop both _row_number_ columns
       context.relBuilder.projectExcept(
           List.of(
-              context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_MAIN),
-              context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_SUBSEARCH)));
+              context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_MAIN),
+              context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_SUBSEARCH)));
       return context.relBuilder.peek();
     } else {
       // 9. if override = true, override the duplicated columns in main by subsearch values
@@ -1697,11 +1697,11 @@ public RelNode visitAppendCol(AppendCol node, CalcitePlanContext context) {
           mainFields.stream().filter(subsearchFields::contains).collect(Collectors.toSet());
       RexNode caseCondition =
           context.relBuilder.equals(
-              context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_MAIN),
-              context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_SUBSEARCH));
+              context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_MAIN),
+              context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_SUBSEARCH));
       for (int mainFieldIndex = 0; mainFieldIndex < mainFields.size(); mainFieldIndex++) {
         String mainFieldName = mainFields.get(mainFieldIndex);
-        if (mainFieldName.equals(ROW_NUMBER_COLUMN_NAME_MAIN)) {
+        if (mainFieldName.equals(ROW_NUMBER_COLUMN_FOR_MAIN)) {
           continue;
         }
         finalFieldNames.add(mainFieldName);
@@ -1726,7 +1726,7 @@ public RelNode visitAppendCol(AppendCol node, CalcitePlanContext context) {
           subsearchFieldIndex < subsearchFields.size();
           subsearchFieldIndex++) {
         String subsearchFieldName = subsearchFields.get(subsearchFieldIndex);
-        if (subsearchFieldName.equals(ROW_NUMBER_COLUMN_NAME_SUBSEARCH)) {
+        if (subsearchFieldName.equals(ROW_NUMBER_COLUMN_FOR_SUBSEARCH)) {
           continue;
         }
         if (!duplicatedFields.contains(subsearchFieldName)) {
@@ -1887,14 +1887,15 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
       addIgnoreNullBucketHintToAggregate(context);
     }
 
-    // 2. add a window column
+    // 2. add count() column with sort direction
     List<RexNode> partitionKeys = rexVisitor.analyze(node.getGroupExprList(), context);
     RexNode countField;
     if (node.getCommandType() == RareTopN.CommandType.TOP) {
       countField = context.relBuilder.desc(context.relBuilder.field(countFieldName));
     } else {
       countField = context.relBuilder.field(countFieldName);
     }
+
     RexNode rowNumberWindowOver =
         PlanUtils.makeOver(
             context,
@@ -1905,22 +1906,22 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
             List.of(countField),
             WindowFrame.toCurrentRow());
     context.relBuilder.projectPlus(
-        context.relBuilder.alias(rowNumberWindowOver, ROW_NUMBER_COLUMN_NAME_TOP_RARE));
+        context.relBuilder.alias(rowNumberWindowOver, ROW_NUMBER_COLUMN_FOR_RARE_TOP));
 
     // 3. filter row_number() <= k in each partition
     int k = node.getNoOfResults();
     context.relBuilder.filter(
         context.relBuilder.lessThanOrEqual(
-            context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_TOP_RARE),
+            context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_RARE_TOP),
             context.relBuilder.literal(k)));
 
     // 4. project final output. the default output is group by list + field list
     Boolean showCount = (Boolean) argumentMap.get(RareTopN.Option.showCount.name()).getValue();
     if (showCount) {
-      context.relBuilder.projectExcept(context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_TOP_RARE));
+      context.relBuilder.projectExcept(context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_RARE_TOP));
     } else {
       context.relBuilder.projectExcept(
-          context.relBuilder.field(ROW_NUMBER_COLUMN_NAME_TOP_RARE),
+          context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_RARE_TOP),
           context.relBuilder.field(countFieldName));
     }
     return context.relBuilder.peek();

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

@@ -62,9 +62,9 @@ public interface PlanUtils {
   /** this is only for dedup command, do not reuse it in other command */
   String ROW_NUMBER_COLUMN_FOR_DEDUP = "_row_number_dedup_";
 
-  String ROW_NUMBER_COLUMN_NAME_TOP_RARE = "_row_number_top_rare_";
-  String ROW_NUMBER_COLUMN_NAME_MAIN = "_row_number_main_";
-  String ROW_NUMBER_COLUMN_NAME_SUBSEARCH = "_row_number_subsearch_";
+  String ROW_NUMBER_COLUMN_FOR_RARE_TOP = "_row_number_rare_top_";
+  String ROW_NUMBER_COLUMN_FOR_MAIN = "_row_number_main_";
+  String ROW_NUMBER_COLUMN_FOR_SUBSEARCH = "_row_number_subsearch_";
 
   static SpanUnit intervalUnitToSpanUnit(IntervalUnit unit) {
     return switch (unit) {
@@ -447,10 +447,18 @@ static RexNode derefMapCall(RexNode rexNode) {
     return rexNode;
   }
 
-  /** Check if contains RexOver */
+  /** Check if contains RexOver introduced by dedup */
   static boolean containsRowNumberDedup(LogicalProject project) {
     return project.getProjects().stream()
-        .anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER);
+            .anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER)
+        && project.getRowType().getFieldNames().contains(ROW_NUMBER_COLUMN_FOR_DEDUP);
+  }
+
+  /** Check if contains RexOver introduced by dedup top/rare */
+  static boolean containsRowNumberRareTop(LogicalProject project) {
+    return project.getProjects().stream()
+            .anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER)
+        && project.getRowType().getFieldNames().contains(ROW_NUMBER_COLUMN_FOR_RARE_TOP);
   }
 
   /** Get all RexWindow list from LogicalProject */

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

@@ -1049,29 +1049,29 @@ public void testExplainCountsByAgg() throws IOException {
   }
 
   @Test
-  public void testExplainSortOnMetrics() throws IOException {
+  public void testExplainSortOnMeasure() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
-    String expected = loadExpectedPlan("explain_agg_sort_on_metrics1.yaml");
+    String expected = loadExpectedPlan("explain_agg_sort_on_measure1.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() by"
                 + " state | sort `count()`"));
-    expected = loadExpectedPlan("explain_agg_sort_on_metrics2.yaml");
+    expected = loadExpectedPlan("explain_agg_sort_on_measure2.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false sum(balance)"
                 + " as sum by state | sort - sum"));
     // TODO limit should pushdown to non-composite agg
-    expected = loadExpectedPlan("explain_agg_sort_on_metrics3.yaml");
+    expected = loadExpectedPlan("explain_agg_sort_on_measure3.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
             String.format(
                 "source=%s | stats count() as cnt by span(birthdate, 1d) | sort - cnt",
                 TEST_INDEX_BANK)));
-    expected = loadExpectedPlan("explain_agg_sort_on_metrics4.yaml");
+    expected = loadExpectedPlan("explain_agg_sort_on_measure4.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
@@ -1082,9 +1082,9 @@ public void testExplainSortOnMetrics() throws IOException {
   }
 
   @Test
-  public void testExplainSortOnMetricsMultiTerms() throws IOException {
+  public void testExplainSortOnMeasureMultiTerms() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
-    String expected = loadExpectedPlan("explain_agg_sort_on_metrics_multi_terms.yaml");
+    String expected = loadExpectedPlan("explain_agg_sort_on_measure_multi_terms.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
@@ -1093,11 +1093,11 @@ public void testExplainSortOnMetricsMultiTerms() throws IOException {
   }
 
   @Test
-  public void testExplainCompositeMultiBucketsAutoDateThenSortOnMetricsNotPushdown()
+  public void testExplainCompositeMultiBucketsAutoDateThenSortOnMeasureNotPushdown()
       throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("agg_composite_multi_terms_autodate_sort_agg_metric_not_push.yaml"),
+        loadExpectedPlan("agg_composite_multi_terms_autodate_sort_agg_measure_not_push.yaml"),
         explainQueryYaml(
             String.format(
                 "source=%s | bin timestamp bins=3 | stats bucket_nullable=false avg(value), count()"
@@ -1106,10 +1106,10 @@ public void testExplainCompositeMultiBucketsAutoDateThenSortOnMetricsNotPushdown
   }
 
   @Test
-  public void testExplainCompositeRangeThenSortOnMetricsNotPushdown() throws IOException {
+  public void testExplainCompositeRangeThenSortOnMeasureNotPushdown() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("agg_composite_range_sort_agg_metric_not_push.yaml"),
+        loadExpectedPlan("agg_composite_range_sort_agg_measure_not_push.yaml"),
         explainQueryYaml(
             String.format(
                 "source=%s | eval value_range = case(value < 7000, 'small'"
@@ -1119,10 +1119,10 @@ public void testExplainCompositeRangeThenSortOnMetricsNotPushdown() throws IOExc
   }
 
   @Test
-  public void testExplainCompositeAutoDateThenSortOnMetricsNotPushdown() throws IOException {
+  public void testExplainCompositeAutoDateThenSortOnMeasureNotPushdown() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("agg_composite_autodate_sort_agg_metric_not_push.yaml"),
+        loadExpectedPlan("agg_composite_autodate_sort_agg_measure_not_push.yaml"),
         explainQueryYaml(
             String.format(
                 "source=%s | bin timestamp bins=3 | stats bucket_nullable=false avg(value), count()"
@@ -1131,10 +1131,10 @@ public void testExplainCompositeAutoDateThenSortOnMetricsNotPushdown() throws IO
   }
 
   @Test
-  public void testExplainCompositeRangeAutoDateThenSortOnMetricsNotPushdown() throws IOException {
+  public void testExplainCompositeRangeAutoDateThenSortOnMeasureNotPushdown() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     assertYamlEqualsIgnoreId(
-        loadExpectedPlan("agg_composite_autodate_range_metric_sort_agg_metric_not_push.yaml"),
+        loadExpectedPlan("agg_composite_autodate_range_metric_sort_agg_measure_not_push.yaml"),
         explainQueryYaml(
             String.format(
                 "source=%s | bin timestamp bins=3 | eval value_range = case(value < 7000, 'small'"
@@ -1144,16 +1144,16 @@ public void testExplainCompositeRangeAutoDateThenSortOnMetricsNotPushdown() thro
   }
 
   @Test
-  public void testExplainMultipleAggregatorsWithSortOnOneMetricNotPushDown() throws IOException {
+  public void testExplainMultipleAggregatorsWithSortOnOneMeasureNotPushDown() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
     String expected =
-        loadExpectedPlan("explain_multiple_agg_with_sort_on_one_metric_not_push1.yaml");
+        loadExpectedPlan("explain_multiple_agg_with_sort_on_one_measure_not_push1.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(
             "source=opensearch-sql_test_index_account | stats bucket_nullable=false count() as c,"
                 + " sum(balance) as s by state | sort c"));
-    expected = loadExpectedPlan("explain_multiple_agg_with_sort_on_one_metric_not_push2.yaml");
+    expected = loadExpectedPlan("explain_multiple_agg_with_sort_on_one_measure_not_push2.yaml");
     assertYamlEqualsIgnoreId(
         expected,
         explainQueryYaml(