diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 4ab5b40e4ac..c6a3e6197db 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/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 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); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index 5ce83c66ff8..52c208da156 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/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()); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index fd418c73a60..e4c7220a39f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/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,