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 c6a3e6197db..7a5bc830d05 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 @@ -8,7 +8,6 @@ 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; @@ -31,6 +30,9 @@ import org.opensearch.search.SearchModule; import org.opensearch.search.builder.PointInTimeBuilder; import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.FieldSortBuilder; +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; @@ -51,7 +53,7 @@ public class OpenSearchQueryRequest implements OpenSearchRequest { private final IndexName indexName; /** Search request source builder. */ - private SearchSourceBuilder sourceBuilder; + private final SearchSourceBuilder sourceBuilder; /** OpenSearchExprValueFactory. */ @EqualsAndHashCode.Exclude @ToString.Exclude @@ -203,12 +205,23 @@ public OpenSearchResponse searchWithPIT(Function if (searchAfter != null) { this.sourceBuilder.searchAfter(searchAfter); } - // Set sort field for search_after - if (this.sourceBuilder.sorts() == null) { + // 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 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); + this.sourceBuilder.sort(SortBuilders.shardDocSort()); + } else { + // If sort fields specified, sort by `fields` + `_doc` + `_shard_doc`. + if (this.sourceBuilder.sorts().stream() + .noneMatch( + b -> b instanceof FieldSortBuilder f && f.fieldName().equals(DOC_FIELD_NAME))) { + this.sourceBuilder.sort(DOC_FIELD_NAME, ASC); + } + if (this.sourceBuilder.sorts().stream().noneMatch(ShardDocSortBuilder.class::isInstance)) { + this.sourceBuilder.sort(SortBuilders.shardDocSort()); + } } SearchRequest searchRequest = new SearchRequest().indices(indexName.getIndexNames()).source(this.sourceBuilder); 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 52c208da156..5ce83c66ff8 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,7 +184,6 @@ 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()); @@ -215,7 +214,6 @@ 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()); @@ -237,7 +235,6 @@ 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 e4c7220a39f..0f8b298732e 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 @@ -13,7 +13,6 @@ 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; @@ -408,7 +407,7 @@ void test_push_down_project() { .size(MAX_RESULT_WINDOW) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(DOC_FIELD_NAME, ASC) - .sort(METADATA_FIELD_ID, ASC) + .sort(SortBuilders.shardDocSort()) .pointInTimeBuilder(new PointInTimeBuilder("samplePITId")) .fetchSource(new String[] {"intA"}, new String[0]), requestBuilder); @@ -421,7 +420,7 @@ void test_push_down_project() { .size(MAX_RESULT_WINDOW) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(DOC_FIELD_NAME, ASC) - .sort(METADATA_FIELD_ID, ASC) + .sort(SortBuilders.shardDocSort()) .pointInTimeBuilder(new PointInTimeBuilder("samplePITId")) .fetchSource("intA", null), exprValueFactory, @@ -536,7 +535,7 @@ void test_push_down_project_with_alias_type() { .size(MAX_RESULT_WINDOW) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(DOC_FIELD_NAME, ASC) - .sort(METADATA_FIELD_ID, ASC) + .sort(SortBuilders.shardDocSort()) .pointInTimeBuilder(new PointInTimeBuilder("samplePITId")) .fetchSource(new String[] {"intA"}, new String[0]), requestBuilder); @@ -549,7 +548,7 @@ void test_push_down_project_with_alias_type() { .size(MAX_RESULT_WINDOW) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(DOC_FIELD_NAME, ASC) - .sort(METADATA_FIELD_ID, ASC) + .sort(SortBuilders.shardDocSort()) .pointInTimeBuilder(new PointInTimeBuilder("samplePITId")) .fetchSource("intA", null), exprValueFactory,