From 5bcd134eba976a9eb9c93849e277485e1d15a210 Mon Sep 17 00:00:00 2001 From: Rupal Mahajan Date: Thu, 13 Jun 2024 14:33:23 -0700 Subject: [PATCH] Update remaining scroll search calls Signed-off-by: Rupal Mahajan --- .../executor/ElasticDefaultRestExecutor.java | 9 ++--- .../legacy/executor/multi/MinusExecutor.java | 33 ++++++++++--------- .../legacy/query/multi/MultiQueryAction.java | 1 + 3 files changed, 21 insertions(+), 22 deletions(-) diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/ElasticDefaultRestExecutor.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/ElasticDefaultRestExecutor.java index 95c1d470d8..a555d43105 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/executor/ElasticDefaultRestExecutor.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/executor/ElasticDefaultRestExecutor.java @@ -59,10 +59,10 @@ public void execute( Client client, Map params, QueryAction queryAction, RestChannel channel) throws Exception { ActionRequest request = requestBuilder.request(); + boolean isSearchAfter = + LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER); if (requestBuilder instanceof JoinRequestBuilder) { - boolean isSearchAfter = - LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER); if (isSearchAfter) { ((JoinRequestBuilder) requestBuilder).updateRequestWithPit(client); } @@ -73,18 +73,13 @@ public void execute( } executor.sendResponse(channel); } else if (requestBuilder instanceof MultiQueryRequestBuilder) { - - boolean isSearchAfter = - LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER); if (isSearchAfter) { ((MultiQueryRequestBuilder) requestBuilder).updateRequestWithPit(client); } - ElasticHitsExecutor executor = MultiRequestExecutorFactory.createExecutor( client, (MultiQueryRequestBuilder) requestBuilder); executor.run(); - if (isSearchAfter) { ((MultiQueryRequestBuilder) requestBuilder).getPit().deletePointInTime(client); } diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java index 52f8953083..3d6f62b5b5 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java @@ -36,7 +36,6 @@ import org.opensearch.sql.legacy.esdomain.LocalClusterState; import org.opensearch.sql.legacy.exception.SqlParseException; import org.opensearch.sql.legacy.executor.ElasticHitsExecutor; -import org.opensearch.sql.legacy.executor.join.ElasticUtils; import org.opensearch.sql.legacy.query.DefaultQueryAction; import org.opensearch.sql.legacy.query.multi.MultiQueryRequestBuilder; import org.opensearch.sql.legacy.utils.Util; @@ -192,7 +191,7 @@ private void fillMinusHitsFromResults(Set comperableHitResu private Set runWithScrollings() { - SearchResponse response = + SearchResponse scrollResp = getResponseWithHits( client, builder.getFirstSearchRequest(), @@ -202,7 +201,7 @@ private Set runWithScrollings() { builder.getPit().getPitId()); Set results = new HashSet<>(); - SearchHit[] hits = response.getHits().getHits(); + SearchHit[] hits = scrollResp.getHits().getHits(); if (hits == null || hits.length == 0) { return new HashSet<>(); } @@ -214,24 +213,26 @@ private Set runWithScrollings() { if (totalDocsFetchedFromFirstTable > this.maxDocsToFetchOnFirstTable) { break; } - response = + scrollResp = getResponseWithHits( client, builder.getFirstSearchRequest(), builder.getOriginalSelect(true), maxDocsToFetchOnEachScrollShard, - response, + scrollResp, builder.getPit().getPitId()); - hits = response.getHits().getHits(); + hits = scrollResp.getHits().getHits(); } - response = - ElasticUtils.scrollOneTimeWithHits( + scrollResp = + getResponseWithHits( this.client, this.builder.getSecondSearchRequest(), builder.getOriginalSelect(false), - this.maxDocsToFetchOnEachScrollShard); + this.maxDocsToFetchOnEachScrollShard, + null, + builder.getPit().getPitId()); - hits = response.getHits().getHits(); + hits = scrollResp.getHits().getHits(); if (hits == null || hits.length == 0) { return results; } @@ -242,15 +243,15 @@ private Set runWithScrollings() { if (totalDocsFetchedFromSecondTable > this.maxDocsToFetchOnSecondTable) { break; } - response = + scrollResp = getResponseWithHits( client, builder.getSecondSearchRequest(), builder.getOriginalSelect(false), maxDocsToFetchOnEachScrollShard, - response, + scrollResp, builder.getPit().getPitId()); - hits = response.getHits().getHits(); + hits = scrollResp.getHits().getHits(); } return results; @@ -348,11 +349,13 @@ private MinusOneFieldAndOptimizationResult runWithScrollingAndAddFilter( break; } SearchResponse responseForSecondTable = - ElasticUtils.scrollOneTimeWithHits( + getResponseWithHits( this.client, queryAction.getRequestBuilder(), secondQuerySelect, - this.maxDocsToFetchOnEachScrollShard); + this.maxDocsToFetchOnEachScrollShard, + null, + builder.getPit().getPitId()); SearchHits secondQuerySearchHits = responseForSecondTable.getHits(); SearchHit[] secondQueryHits = secondQuerySearchHits.getHits(); diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/query/multi/MultiQueryAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/query/multi/MultiQueryAction.java index a2a4d91343..a9eb6113f7 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/query/multi/MultiQueryAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/query/multi/MultiQueryAction.java @@ -40,6 +40,7 @@ public SqlElasticRequestBuilder explain() throws SqlParseException { requestBuilder.fillTableAliases( this.multiQuerySelect.getFirstSelect().getFields(), this.multiQuerySelect.getSecondSelect().getFields()); + return requestBuilder; }