Skip to content

Commit 3388dc7

Browse files
LantaoJinLantaoJin
authored
Use _doc + _shard_doc as sort tiebreaker to get better performance (#4569)
* Use _shard_doc as sort tiebreaker Signed-off-by: Lantao Jin <ltjin@amazon.com> * _doc as a part of tie-breaker have better performance Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 5630119 commit 3388dc7

3 files changed

Lines changed: 24 additions & 15 deletions

File tree

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

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import static org.opensearch.core.xcontent.DeprecationHandler.IGNORE_DEPRECATIONS;
99
import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME;
1010
import static org.opensearch.search.sort.SortOrder.ASC;
11-
import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
1211

1312
import java.io.IOException;
1413
import java.util.Collections;
@@ -31,6 +30,9 @@
3130
import org.opensearch.search.SearchModule;
3231
import org.opensearch.search.builder.PointInTimeBuilder;
3332
import org.opensearch.search.builder.SearchSourceBuilder;
33+
import org.opensearch.search.sort.FieldSortBuilder;
34+
import org.opensearch.search.sort.ShardDocSortBuilder;
35+
import org.opensearch.search.sort.SortBuilders;
3436
import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
3537
import org.opensearch.sql.opensearch.response.OpenSearchResponse;
3638
import org.opensearch.sql.opensearch.storage.OpenSearchIndex;
@@ -51,7 +53,7 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
5153
private final IndexName indexName;
5254

5355
/** Search request source builder. */
54-
private SearchSourceBuilder sourceBuilder;
56+
private final SearchSourceBuilder sourceBuilder;
5557

5658
/** OpenSearchExprValueFactory. */
5759
@EqualsAndHashCode.Exclude @ToString.Exclude
@@ -203,12 +205,23 @@ public OpenSearchResponse searchWithPIT(Function<SearchRequest, SearchResponse>
203205
if (searchAfter != null) {
204206
this.sourceBuilder.searchAfter(searchAfter);
205207
}
206-
// Set sort field for search_after
207-
if (this.sourceBuilder.sorts() == null) {
208+
// Add sort tiebreaker for PIT search.
209+
// We cannot remove it since `_shard_doc` is not added implicitly in PIT now.
210+
// Ref https://github.com/opensearch-project/OpenSearch/pull/18924#issuecomment-3342365950
211+
if (this.sourceBuilder.sorts() == null || this.sourceBuilder.sorts().isEmpty()) {
212+
// If no sort field specified, sort by `_doc` + `_shard_doc`to get better performance
208213
this.sourceBuilder.sort(DOC_FIELD_NAME, ASC);
209-
// Workaround to preserve sort location more exactly,
210-
// see https://github.com/opensearch-project/sql/pull/3061
211-
this.sourceBuilder.sort(METADATA_FIELD_ID, ASC);
214+
this.sourceBuilder.sort(SortBuilders.shardDocSort());
215+
} else {
216+
// If sort fields specified, sort by `fields` + `_doc` + `_shard_doc`.
217+
if (this.sourceBuilder.sorts().stream()
218+
.noneMatch(
219+
b -> b instanceof FieldSortBuilder f && f.fieldName().equals(DOC_FIELD_NAME))) {
220+
this.sourceBuilder.sort(DOC_FIELD_NAME, ASC);
221+
}
222+
if (this.sourceBuilder.sorts().stream().noneMatch(ShardDocSortBuilder.class::isInstance)) {
223+
this.sourceBuilder.sort(SortBuilders.shardDocSort());
224+
}
212225
}
213226
SearchRequest searchRequest =
214227
new SearchRequest().indices(indexName.getIndexNames()).source(this.sourceBuilder);

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,6 @@ void search_with_pit() {
184184
when(searchResponse.getHits()).thenReturn(searchHits);
185185
when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
186186
when(searchHit.getSortValues()).thenReturn(new String[] {"sortedValue"});
187-
when(sourceBuilder.sorts()).thenReturn(null);
188187

189188
OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
190189
assertFalse(openSearchResponse.isEmpty());
@@ -215,7 +214,6 @@ void search_with_pit_hits_null() {
215214
when(searchAction.apply(any())).thenReturn(searchResponse);
216215
when(searchResponse.getHits()).thenReturn(searchHits);
217216
when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit});
218-
when(sourceBuilder.sorts()).thenReturn(null);
219217

220218
OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
221219
assertFalse(openSearchResponse.isEmpty());
@@ -237,7 +235,6 @@ void search_with_pit_hits_empty() {
237235
when(searchAction.apply(any())).thenReturn(searchResponse);
238236
when(searchResponse.getHits()).thenReturn(searchHits);
239237
when(searchHits.getHits()).thenReturn(new SearchHit[] {});
240-
when(sourceBuilder.sorts()).thenReturn(null);
241238

242239
OpenSearchResponse openSearchResponse = request.searchWithPIT(searchAction);
243240
assertTrue(openSearchResponse.isEmpty());

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
import static org.opensearch.search.sort.SortOrder.ASC;
1414
import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
1515
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
16-
import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
1716

1817
import java.util.Collections;
1918
import java.util.List;
@@ -408,7 +407,7 @@ void test_push_down_project() {
408407
.size(MAX_RESULT_WINDOW)
409408
.timeout(DEFAULT_QUERY_TIMEOUT)
410409
.sort(DOC_FIELD_NAME, ASC)
411-
.sort(METADATA_FIELD_ID, ASC)
410+
.sort(SortBuilders.shardDocSort())
412411
.pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
413412
.fetchSource(new String[] {"intA"}, new String[0]),
414413
requestBuilder);
@@ -421,7 +420,7 @@ void test_push_down_project() {
421420
.size(MAX_RESULT_WINDOW)
422421
.timeout(DEFAULT_QUERY_TIMEOUT)
423422
.sort(DOC_FIELD_NAME, ASC)
424-
.sort(METADATA_FIELD_ID, ASC)
423+
.sort(SortBuilders.shardDocSort())
425424
.pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
426425
.fetchSource("intA", null),
427426
exprValueFactory,
@@ -536,7 +535,7 @@ void test_push_down_project_with_alias_type() {
536535
.size(MAX_RESULT_WINDOW)
537536
.timeout(DEFAULT_QUERY_TIMEOUT)
538537
.sort(DOC_FIELD_NAME, ASC)
539-
.sort(METADATA_FIELD_ID, ASC)
538+
.sort(SortBuilders.shardDocSort())
540539
.pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
541540
.fetchSource(new String[] {"intA"}, new String[0]),
542541
requestBuilder);
@@ -549,7 +548,7 @@ void test_push_down_project_with_alias_type() {
549548
.size(MAX_RESULT_WINDOW)
550549
.timeout(DEFAULT_QUERY_TIMEOUT)
551550
.sort(DOC_FIELD_NAME, ASC)
552-
.sort(METADATA_FIELD_ID, ASC)
551+
.sort(SortBuilders.shardDocSort())
553552
.pointInTimeBuilder(new PointInTimeBuilder("samplePITId"))
554553
.fetchSource("intA", null),
555554
exprValueFactory,

0 commit comments

Comments
 (0)