Leave OpenSearchQueryRequest for manual merge

Author: penghuo

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java

@@ -8,6 +8,7 @@
 import static org.opensearch.core.xcontent.DeprecationHandler.IGNORE_DEPRECATIONS;
 import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME;
 import static org.opensearch.search.sort.SortOrder.ASC;
+import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -38,8 +39,9 @@
 import org.opensearch.search.aggregations.bucket.composite.InternalComposite;
 import org.opensearch.search.builder.PointInTimeBuilder;
 import org.opensearch.search.builder.SearchSourceBuilder;
-import org.opensearch.search.sort.FieldSortBuilder;
-import org.opensearch.search.sort.SortBuilders;
+import org.opensearch.sql.monitor.profile.MetricName;
+import org.opensearch.sql.monitor.profile.ProfileMetric;
+import org.opensearch.sql.monitor.profile.QueryProfiling;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 import org.opensearch.sql.opensearch.storage.OpenSearchIndex;
@@ -56,7 +58,6 @@
 @ToString
 public class OpenSearchQueryRequest implements OpenSearchRequest {
   private static final Logger LOG = LogManager.getLogger();
-  private static final String SHARD_DOC_FIELD_NAME = "_shard_doc";
 
   /** {@link OpenSearchRequest.IndexName}. */
   private final IndexName indexName;
@@ -206,6 +207,8 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
     } else {
+      ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE);
+      long executionStartTime = System.nanoTime();
       // Set afterKey to request, null for first round (afterKey is null in the beginning).
       if (this.sourceBuilder.aggregations() != null) {
         this.sourceBuilder.aggregations().getAggregatorFactories().stream()
@@ -238,6 +241,7 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
         searchDone = true;
       }
       needClean = searchDone;
+      metric.add(System.nanoTime() - executionStartTime);
     }
     return openSearchResponse;
   }
@@ -249,35 +253,20 @@ public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse>
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
     } else {
+      ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE);
+      long executionStartTime = System.nanoTime();
       this.sourceBuilder.pointInTimeBuilder(new PointInTimeBuilder(this.pitId));
       this.sourceBuilder.timeout(cursorKeepAlive);
       // check for search after
       if (searchAfter != null) {
         this.sourceBuilder.searchAfter(searchAfter);
       }
-      // Add sort tiebreaker for PIT search.
-      // We cannot remove it since `_shard_doc` is not added implicitly in PIT now.
-      // Ref https://github.com/opensearch-project/OpenSearch/pull/18924#issuecomment-3342365950
-      if (this.sourceBuilder.sorts() == null || this.sourceBuilder.sorts().isEmpty()) {
-        // If no sort field specified, sort by `_doc` + `_shard_doc`to get better performance
+      // Set sort field for search_after
+      if (this.sourceBuilder.sorts() == null) {
         this.sourceBuilder.sort(DOC_FIELD_NAME, ASC);
-        this.sourceBuilder.sort(SortBuilders.fieldSort(SHARD_DOC_FIELD_NAME));
-      } else {
-        // If sort fields specified, sort by `fields` + `_doc` + `_shard_doc`.
-        if (this.sourceBuilder.sorts().stream()
-            .noneMatch(
-                b ->
-                    b instanceof FieldSortBuilder
-                        && ((FieldSortBuilder) b).fieldName().equals(DOC_FIELD_NAME))) {
-          this.sourceBuilder.sort(DOC_FIELD_NAME, ASC);
-        }
-        if (this.sourceBuilder.sorts().stream()
-            .noneMatch(
-                b ->
-                    b instanceof FieldSortBuilder
-                        && SHARD_DOC_FIELD_NAME.equals(((FieldSortBuilder) b).fieldName()))) {
-          this.sourceBuilder.sort(SortBuilders.fieldSort(SHARD_DOC_FIELD_NAME));
-        }
+        // Workaround to preserve sort location more exactly,
+        // see https://github.com/opensearch-project/sql/pull/3061
+        this.sourceBuilder.sort(METADATA_FIELD_ID, ASC);
       }
       SearchRequest searchRequest =
           new SearchRequest().indices(indexName.getIndexNames()).source(this.sourceBuilder);
@@ -297,6 +286,7 @@ public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse>
           LOG.debug(sourceBuilder);
         }
       }
+      metric.add(System.nanoTime() - executionStartTime);
     }
     return openSearchResponse;
   }