reverting to _doc + _id (#4435)

Author: ahkcs

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

@@ -6,6 +6,9 @@
 package org.opensearch.sql.opensearch.request;
 
 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;
@@ -28,8 +31,6 @@
 import org.opensearch.search.SearchModule;
 import org.opensearch.search.builder.PointInTimeBuilder;
 import org.opensearch.search.builder.SearchSourceBuilder;
-import org.opensearch.search.sort.ShardDocSortBuilder;
-import org.opensearch.search.sort.SortBuilders;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 import org.opensearch.sql.opensearch.storage.OpenSearchIndex;
@@ -202,14 +203,13 @@ public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse>
       if (searchAfter != null) {
         this.sourceBuilder.searchAfter(searchAfter);
       }
-      // Add sort tiebreaker `_shard_doc` for PIT search.
-      // Actually, we can remove it since `_shard_doc` should be added implicitly in PIT.
-      // https://github.com/opensearch-project/OpenSearch/pull/18924#issuecomment-3342365950
-      if (this.sourceBuilder.sorts() == null
-          || this.sourceBuilder.sorts().stream().noneMatch(ShardDocSortBuilder.class::isInstance)) {
-        this.sourceBuilder.sort(SortBuilders.shardDocSort());
+      // Set sort field for search_after
+      if (this.sourceBuilder.sorts() == null) {
+        this.sourceBuilder.sort(DOC_FIELD_NAME, ASC);
+        // 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);
       this.searchResponse = searchAction.apply(searchRequest);

opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java

@@ -184,6 +184,7 @@ void search_with_pit() {
     when(searchResponse.getHits()).thenReturn(searchHits);
     when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
     when(searchHit.getSortValues()).thenReturn(new String[] {"sortedValue"});
+    when(sourceBuilder.sorts()).thenReturn(null);
 
     OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
     assertFalse(openSearchResponse.isEmpty());
@@ -214,6 +215,7 @@ void search_with_pit_hits_null() {
     when(searchAction.apply(any())).thenReturn(searchResponse);
     when(searchResponse.getHits()).thenReturn(searchHits);
     when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
+    when(sourceBuilder.sorts()).thenReturn(null);
 
     OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
     assertFalse(openSearchResponse.isEmpty());
@@ -235,6 +237,7 @@ void search_with_pit_hits_empty() {
     when(searchAction.apply(any())).thenReturn(searchResponse);
     when(searchResponse.getHits()).thenReturn(searchHits);
     when(searchHits.getHits()).thenReturn(new SearchHit[] {});
+    when(sourceBuilder.sorts()).thenReturn(null);
 
     OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
     assertTrue(openSearchResponse.isEmpty());

opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java

@@ -9,8 +9,11 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.Mockito.*;
 import static org.opensearch.index.query.QueryBuilders.*;
+import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME;
+import static org.opensearch.search.sort.SortOrder.ASC;
 import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
 
 import java.util.Collections;
 import java.util.List;
@@ -404,7 +407,8 @@ void test_push_down_project() {
             .from(DEFAULT_OFFSET)
             .size(MAX_RESULT_WINDOW)
             .timeout(DEFAULT_QUERY_TIMEOUT)
-            .sort(SortBuilders.shardDocSort())
+            .sort(DOC_FIELD_NAME, ASC)
+            .sort(METADATA_FIELD_ID, ASC)
             .pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
             .fetchSource(new String[] {"intA"}, new String[0]),
         requestBuilder);
@@ -416,7 +420,8 @@ void test_push_down_project() {
                 .from(DEFAULT_OFFSET)
                 .size(MAX_RESULT_WINDOW)
                 .timeout(DEFAULT_QUERY_TIMEOUT)
-                .sort(SortBuilders.shardDocSort())
+                .sort(DOC_FIELD_NAME, ASC)
+                .sort(METADATA_FIELD_ID, ASC)
                 .pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
                 .fetchSource("intA", null),
             exprValueFactory,
@@ -530,7 +535,8 @@ void test_push_down_project_with_alias_type() {
             .from(DEFAULT_OFFSET)
             .size(MAX_RESULT_WINDOW)
             .timeout(DEFAULT_QUERY_TIMEOUT)
-            .sort(SortBuilders.shardDocSort())
+            .sort(DOC_FIELD_NAME, ASC)
+            .sort(METADATA_FIELD_ID, ASC)
             .pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
             .fetchSource(new String[] {"intA"}, new String[0]),
         requestBuilder);
@@ -542,7 +548,8 @@ void test_push_down_project_with_alias_type() {
                 .from(DEFAULT_OFFSET)
                 .size(MAX_RESULT_WINDOW)
                 .timeout(DEFAULT_QUERY_TIMEOUT)
-                .sort(SortBuilders.shardDocSort())
+                .sort(DOC_FIELD_NAME, ASC)
+                .sort(METADATA_FIELD_ID, ASC)
                 .pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
                 .fetchSource("intA", null),
             exprValueFactory,