-
Notifications
You must be signed in to change notification settings - Fork 138
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add scroll search based OpenSearchReader (#1633)
Signed-off-by: Peng Huo <penghuo@gmail.com>
- Loading branch information
Showing
29 changed files
with
544 additions
and
473 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
18 changes: 18 additions & 0 deletions
18
flint/flint-core/src/main/scala/org/opensearch/flint/core/FlintClientBuilder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.flint.core; | ||
|
||
import org.opensearch.flint.core.storage.FlintOpenSearchClient; | ||
|
||
/** | ||
* {@link FlintClient} builder. | ||
*/ | ||
public class FlintClientBuilder { | ||
|
||
public static FlintClient build(FlintOptions options) { | ||
return new FlintOpenSearchClient(options); | ||
} | ||
} |
41 changes: 41 additions & 0 deletions
41
flint/flint-core/src/main/scala/org/opensearch/flint/core/FlintOptions.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.flint.core; | ||
|
||
import java.io.Serializable; | ||
import java.util.Map; | ||
|
||
/** | ||
* Flint Options include all the flint related configuration. | ||
*/ | ||
public class FlintOptions implements Serializable { | ||
|
||
private final Map<String, String> options; | ||
|
||
public static final String HOST = "host"; | ||
public static final String PORT = "port"; | ||
/** | ||
* Used by {@link org.opensearch.flint.core.storage.OpenSearchScrollReader} | ||
*/ | ||
public static final String SCROLL_SIZE = "scroll_size"; | ||
public static final int DEFAULT_SCROLL_SIZE = 100; | ||
|
||
public FlintOptions(Map<String, String> options) { | ||
this.options = options; | ||
} | ||
|
||
public String getHost() { | ||
return options.getOrDefault(HOST, "localhost"); | ||
} | ||
|
||
public int getPort() { | ||
return Integer.parseInt(options.getOrDefault(PORT, "9200")); | ||
} | ||
|
||
public int getScrollSize() { | ||
return Integer.parseInt(options.getOrDefault(SCROLL_SIZE, String.valueOf(DEFAULT_SCROLL_SIZE))); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
27 changes: 27 additions & 0 deletions
27
flint/flint-core/src/main/scala/org/opensearch/flint/core/storage/FlintReader.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.flint.core.storage; | ||
|
||
/** | ||
* Flint Reader Interface | ||
*/ | ||
public interface FlintReader { | ||
|
||
/** | ||
* true if next doc exist. | ||
*/ | ||
boolean hasNext(); | ||
|
||
/** | ||
* Return next doc in String. | ||
*/ | ||
String next(); | ||
|
||
/** | ||
* close. | ||
*/ | ||
void close(); | ||
} |
81 changes: 81 additions & 0 deletions
81
flint/flint-core/src/main/scala/org/opensearch/flint/core/storage/OpenSearchReader.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,81 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.flint.core.storage; | ||
|
||
import org.opensearch.action.search.SearchRequest; | ||
import org.opensearch.action.search.SearchResponse; | ||
import org.opensearch.client.RestHighLevelClient; | ||
import org.opensearch.search.SearchHit; | ||
|
||
import java.io.IOException; | ||
import java.util.Arrays; | ||
import java.util.Iterator; | ||
import java.util.List; | ||
|
||
/** | ||
* Abstract OpenSearch Reader. | ||
*/ | ||
public abstract class OpenSearchReader implements FlintReader { | ||
|
||
/** Search request source builder. */ | ||
private final SearchRequest searchRequest; | ||
|
||
protected final RestHighLevelClient client; | ||
|
||
/** | ||
* iterator of one-shot search result. | ||
*/ | ||
private Iterator<SearchHit> iterator = null; | ||
|
||
public OpenSearchReader(RestHighLevelClient client, SearchRequest searchRequest) { | ||
this.client = client; | ||
this.searchRequest = searchRequest; | ||
} | ||
|
||
@Override public boolean hasNext() { | ||
try { | ||
if (iterator == null || !iterator.hasNext()) { | ||
SearchResponse response = search(searchRequest); | ||
List<SearchHit> searchHits = Arrays.asList(response.getHits().getHits()); | ||
iterator = searchHits.iterator(); | ||
} | ||
return iterator.hasNext(); | ||
} catch (IOException e) { | ||
// todo. log error. | ||
throw new RuntimeException(e); | ||
} | ||
} | ||
|
||
@Override public String next() { | ||
return iterator.next().getSourceAsString(); | ||
} | ||
|
||
@Override public void close() { | ||
try { | ||
clean(); | ||
} catch (IOException e) { | ||
// todo. log error. | ||
} finally { | ||
if (client != null) { | ||
try { | ||
client.close(); | ||
} catch (IOException e) { | ||
throw new RuntimeException(e); | ||
} | ||
} | ||
} | ||
} | ||
|
||
/** | ||
* search. | ||
*/ | ||
abstract SearchResponse search(SearchRequest request) throws IOException; | ||
|
||
/** | ||
* clean. | ||
*/ | ||
abstract void clean() throws IOException; | ||
} |
63 changes: 63 additions & 0 deletions
63
...t/flint-core/src/main/scala/org/opensearch/flint/core/storage/OpenSearchScrollReader.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.flint.core.storage; | ||
|
||
import org.opensearch.action.search.ClearScrollRequest; | ||
import org.opensearch.action.search.SearchRequest; | ||
import org.opensearch.action.search.SearchResponse; | ||
import org.opensearch.action.search.SearchScrollRequest; | ||
import org.opensearch.client.RequestOptions; | ||
import org.opensearch.client.RestHighLevelClient; | ||
import org.opensearch.common.Strings; | ||
import org.opensearch.common.unit.TimeValue; | ||
import org.opensearch.flint.core.FlintOptions; | ||
import org.opensearch.search.builder.SearchSourceBuilder; | ||
|
||
import java.io.IOException; | ||
|
||
/** | ||
* {@link OpenSearchReader} using scroll search. https://opensearch.org/docs/latest/api-reference/scroll/ | ||
*/ | ||
public class OpenSearchScrollReader extends OpenSearchReader { | ||
|
||
/** Default scroll context timeout in minutes. */ | ||
public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(5L); | ||
|
||
private final FlintOptions options; | ||
|
||
private String scrollId = null; | ||
|
||
public OpenSearchScrollReader(RestHighLevelClient client, String indexName, SearchSourceBuilder searchSourceBuilder, FlintOptions options) { | ||
super(client, new SearchRequest().indices(indexName).source(searchSourceBuilder.size(options.getScrollSize()))); | ||
this.options = options; | ||
} | ||
|
||
/** | ||
* search. | ||
*/ | ||
SearchResponse search(SearchRequest request) throws IOException { | ||
if (Strings.isNullOrEmpty(scrollId)) { | ||
// add scroll timeout making the request as scroll search request. | ||
request.scroll(DEFAULT_SCROLL_TIMEOUT); | ||
SearchResponse response = client.search(request, RequestOptions.DEFAULT); | ||
scrollId = response.getScrollId(); | ||
return response; | ||
} else { | ||
return client.scroll(new SearchScrollRequest().scroll(DEFAULT_SCROLL_TIMEOUT).scrollId(scrollId), RequestOptions.DEFAULT); | ||
} | ||
} | ||
|
||
/** | ||
* clean the scroll context. | ||
*/ | ||
void clean() throws IOException { | ||
if (Strings.isNullOrEmpty(scrollId)) { | ||
ClearScrollRequest clearScrollRequest = new ClearScrollRequest(); | ||
clearScrollRequest.addScrollId(scrollId); | ||
client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT); | ||
} | ||
} | ||
} |
Oops, something went wrong.