Examples with SearchRequest org.elasticsearch.action.search.SearchRequest used on opensource projects

Example 71 with SearchRequest

use of org.elasticsearch.action.search.SearchRequest in project aws-athena-query-federation by awslabs.

the class ElasticsearchRecordHandler method readWithConstraint.

/**
 * Used to read the row data associated with the provided Split.
 *
 * @param spiller A BlockSpiller that should be used to write the row data associated with this Split.
 * The BlockSpiller automatically handles chunking the response, encrypting, and spilling to S3.
 * @param recordsRequest Details of the read request, including:
 * 1. The Split
 * 2. The Catalog, Database, and Table the read request is for.
 * 3. The filtering predicate (if any)
 * 4. The columns required for projection.
 * @param queryStatusChecker A QueryStatusChecker that you can use to stop doing work for a query that has already terminated
 * @throws RuntimeException when an error occurs while attempting to send the query, or the query timed out.
 * @note Avoid writing >10 rows per-call to BlockSpiller.writeRow(...) because this will limit the BlockSpiller's
 * ability to control Block size. The resulting increase in Block size may cause failures and reduced performance.
 */
@Override
protected void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) throws RuntimeException {
    logger.info("readWithConstraint - enter - Domain: {}, Index: {}, Mapping: {}", recordsRequest.getTableName().getSchemaName(), recordsRequest.getTableName().getTableName(), recordsRequest.getSchema());
    String domain = recordsRequest.getTableName().getSchemaName();
    String endpoint = recordsRequest.getSplit().getProperty(domain);
    String index = recordsRequest.getTableName().getTableName();
    String shard = recordsRequest.getSplit().getProperty(ElasticsearchMetadataHandler.SHARD_KEY);
    long numRows = 0;
    if (queryStatusChecker.isQueryRunning()) {
        AwsRestHighLevelClient client = clientFactory.getOrCreateClient(endpoint);
        try {
            // Create field extractors for all data types in the schema.
            GeneratedRowWriter rowWriter = createFieldExtractors(recordsRequest);
            // Create a new search-source injected with the projection, predicate, and the pagination batch size.
            SearchSourceBuilder searchSource = new SearchSourceBuilder().size(QUERY_BATCH_SIZE).timeout(new TimeValue(queryTimeout, TimeUnit.SECONDS)).fetchSource(ElasticsearchQueryUtils.getProjection(recordsRequest.getSchema())).query(ElasticsearchQueryUtils.getQuery(recordsRequest.getConstraints().getSummary()));
            // Create a new search-request for the specified index.
            SearchRequest searchRequest = new SearchRequest(index).preference(shard);
            int hitsNum;
            int currPosition = 0;
            do {
                // Process the search request injecting the search-source, and setting the from position
                // used for pagination of results.
                SearchResponse searchResponse = client.getDocuments(searchRequest.source(searchSource.from(currPosition)));
                // Throw on query timeout.
                if (searchResponse.isTimedOut()) {
                    throw new RuntimeException("Request for index (" + index + ") " + shard + " timed out.");
                }
                // Increment current position to next batch of results.
                currPosition += QUERY_BATCH_SIZE;
                // Process hits.
                Iterator<SearchHit> hitIterator = searchResponse.getHits().iterator();
                hitsNum = searchResponse.getHits().getHits().length;
                while (hitIterator.hasNext() && queryStatusChecker.isQueryRunning()) {
                    ++numRows;
                    spiller.writeRows((Block block, int rowNum) -> rowWriter.writeRow(block, rowNum, client.getDocument(hitIterator.next())) ? 1 : 0);
                }
            // if hitsNum < QUERY_BATCH_SIZE, then this is the last batch of documents.
            } while (hitsNum == QUERY_BATCH_SIZE && queryStatusChecker.isQueryRunning());
        } catch (IOException error) {
            throw new RuntimeException("Error sending search query: " + error.getMessage(), error);
        }
    }
    logger.info("readWithConstraint: numRows[{}]", numRows);
}

Example 72 with SearchRequest

use of org.elasticsearch.action.search.SearchRequest in project pipelines by gbif.

the class ExtensionTermCountRequestBuilder method getRequest.

public ExtTermCountRequest getRequest() {
    String aggsField = prefix == null || prefix.isEmpty() ? term : prefix + "." + term;
    SearchRequest request = new SearchRequest().source(new SearchSourceBuilder().size(0).query(QueryBuilders.termQuery(termName, termValue)).aggregation(AggregationBuilders.count(AGGREGATION).field(aggsField))).indices(indexName);
    return ExtTermCountRequest.create(term, request);
}

Example 73 with SearchRequest

use of org.elasticsearch.action.search.SearchRequest in project starcoin-search by starcoinorg.

the class BlockService method getUncleBlockByHash.

public UncleBlock getUncleBlockByHash(String network, String hash) {
    SearchRequest searchRequest = new SearchRequest(getIndex(network, Constant.UNCLE_BLOCK_INDEX));
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("header.block_hash", hash);
    searchSourceBuilder.query(termQueryBuilder);
    searchRequest.source(searchSourceBuilder);
    SearchResponse searchResponse;
    try {
        searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
    } catch (IOException e) {
        logger.error("get uncle block by hash error:", e);
        return null;
    }
    Result<UncleBlock> result = ServiceUtils.getSearchResult(searchResponse, UncleBlock.class);
    List<UncleBlock> blocks = result.getContents();
    if (blocks.size() == 1) {
        return blocks.get(0);
    } else {
        logger.warn("get uncle block by hash is null, network: {}, : {}", network, hash);
    }
    return null;
}

Example 74 with SearchRequest

use of org.elasticsearch.action.search.SearchRequest in project starcoin-search by starcoinorg.

the class BlockService method getUncleBlockByHeight.

public UncleBlock getUncleBlockByHeight(String network, long height) {
    SearchRequest searchRequest = new SearchRequest(getIndex(network, Constant.UNCLE_BLOCK_INDEX));
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("header.number", height);
    searchSourceBuilder.query(termQueryBuilder);
    searchRequest.source(searchSourceBuilder);
    SearchResponse searchResponse;
    try {
        searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
    } catch (IOException e) {
        logger.error("get uncle block error:", e);
        return null;
    }
    Result<UncleBlock> result = ServiceUtils.getSearchResult(searchResponse, UncleBlock.class);
    List<UncleBlock> blocks = result.getContents();
    if (blocks.size() == 1) {
        return blocks.get(0);
    } else {
        logger.warn("get uncle block by height is null, network: {}, : {}", network, height);
    }
    return null;
}

Example 75 with SearchRequest

use of org.elasticsearch.action.search.SearchRequest in project starcoin-search by starcoinorg.

the class BlockService method getRange.

public Result<Block> getRange(String network, int page, int count, int start_height) {
    SearchRequest searchRequest = new SearchRequest(getIndex(network, Constant.BLOCK_IDS_INDEX));
    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchSourceBuilder.query(QueryBuilders.matchAllQuery());
    // page size
    searchSourceBuilder.size(count);
    // begin offset
    int offset;
    if (page > 1) {
        offset = (page - 1) * count;
        if (offset >= ELASTICSEARCH_MAX_HITS && start_height > 0) {
            offset = start_height - (page - 1) * count;
            searchSourceBuilder.searchAfter(new Object[] { offset });
        } else {
            searchSourceBuilder.from(offset);
        }
    }
    searchSourceBuilder.sort("header.number", SortOrder.DESC);
    searchSourceBuilder.trackTotalHits(true);
    searchRequest.source(searchSourceBuilder);
    SearchResponse searchResponse;
    try {
        searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
    } catch (IOException e) {
        logger.error("get range block error:", e);
        return null;
    }
    return ServiceUtils.getSearchResult(searchResponse, Block.class);
}