Java Code Examples for org.elasticsearch.search.SearchHit#sourceRef()

private SearchResponse mockSearchResponse(String searchIndex, int lookUpTime, String detectorUuid) {
    SearchResponse searchResponse = mock(SearchResponse.class);
    Map<String, DocumentField> fields = new HashMap<>();
    fields.put("_percolator_document_slot", new DocumentField("_percolator_document_slot", Arrays.asList(new Integer(searchIndex))));
    SearchHit searchHit = new SearchHit(101, "xxx", null, fields);
    BytesReference source = new BytesArray("{\"aa_user\":{\"id\":\"test-user\"},\"aa_detector\":" +
            "{\"id\":\"" + detectorUuid + "\"},\"aa_query\":{\"bool\":{\"must\":[{\"match\":" +
            "{\"name\":\"sample-web\",\"env\":\"prod\"}}]}},\"aa_enabled\":true,\"aa_lastModifiedTime\":1554828886," +
            "\"aa_createdTime\":1554828886}\n");
    searchHit.sourceRef(source);
    SearchHit[] bunchOfSearchHits = new SearchHit[1];
    bunchOfSearchHits[0] = searchHit;
    SearchHits searchHits = new SearchHits(bunchOfSearchHits, 1, 1);
    when(searchResponse.getHits()).thenReturn(searchHits);
    TimeValue timeValue = new TimeValue(lookUpTime);
    when(searchResponse.getTook()).thenReturn(timeValue);
    return searchResponse;
}
 

private void createKeyToResultsAndFillOptimizationStructure(Map<String,Map<String, List<Object>>> optimizationTermsFilterStructure, TableInJoinRequestBuilder firstTableRequest) {
    List<SearchHit> firstTableHits = fetchAllHits(firstTableRequest);

    int resultIds = 1;
    for (SearchHit hit : firstTableHits) {
        HashMap<String, List<Map.Entry<Field, Field>>> comparisons = this.hashJoinComparisonStructure.getComparisons();
        for (Map.Entry<String, List<Map.Entry<Field, Field>>> comparison : comparisons.entrySet()) {
            String comparisonID = comparison.getKey();
            List<Map.Entry<Field, Field>> t1ToT2FieldsComparison = comparison.getValue();

            String key = getComparisonKey(t1ToT2FieldsComparison, hit, true, optimizationTermsFilterStructure.get(comparisonID));

            //int docid , id
            SearchHit searchHit = new SearchHit(resultIds, hit.getId(), new Text(hit.getType()), hit.getFields(), null);
            searchHit.sourceRef(hit.getSourceRef());

            onlyReturnedFields(searchHit.getSourceAsMap(), firstTableRequest.getReturnedFields(),firstTableRequest.getOriginalSelect().isSelectAll());
            resultIds++;
            this.hashJoinComparisonStructure.insertIntoComparisonHash(comparisonID, key, searchHit);
        }
    }
}
 

private void fillMinusHitsFromOneField(String fieldName, Set<Object> fieldValues, SearchHit someHit) {
    List<SearchHit> minusHitsList = new ArrayList<>();
    int currentId = 1;
    for(Object result : fieldValues){
        Map<String,DocumentField> fields = new HashMap<>();
        ArrayList<Object> values = new ArrayList<Object>();
        values.add(result);
        fields.put(fieldName,new DocumentField(fieldName, values));
        SearchHit searchHit = new SearchHit(currentId,currentId+"", new Text(someHit.getType()), fields, null);
        searchHit.sourceRef(someHit.getSourceRef());
        searchHit.getSourceAsMap().clear();
        Map<String, Object> sourceAsMap = new HashMap<>();
        sourceAsMap.put(fieldName,result);
        searchHit.getSourceAsMap().putAll(sourceAsMap);
        currentId++;
        minusHitsList.add(searchHit);
    }
    int totalSize = currentId - 1;
    SearchHit[] unionHitsArr = minusHitsList.toArray(new SearchHit[totalSize]);
    this.minusHits = new SearchHits(unionHitsArr, new TotalHits(totalSize, TotalHits.Relation.EQUAL_TO), 1.0f);
}
 

private void fillIntersectHitsFromResults(Set<ComperableHitResult> comparableHitResults) {
    int currentId = 1;
    List<SearchHit> intersectHitsList = new ArrayList<>(comparableHitResults.size());
    Set<Map.Entry<String, String>> firstTableFieldToAlias = this.builder.getFirstTableFieldToAlias().entrySet();
    for (ComperableHitResult result : comparableHitResults) {
        SearchHit originalHit = result.getOriginalHit();
        SearchHit searchHit = new SearchHit(currentId, originalHit.getId(), new Text(originalHit.getType()), originalHit.getFields(), null);
        searchHit.sourceRef(originalHit.getSourceRef());
        searchHit.getSourceAsMap().clear();
        Map<String, Object> sourceAsMap = result.getFlattenMap();
        for (Map.Entry<String, String> entry : firstTableFieldToAlias) {
            if (sourceAsMap.containsKey(entry.getKey())) {
                Object value = sourceAsMap.get(entry.getKey());
                sourceAsMap.remove(entry.getKey());
                sourceAsMap.put(entry.getValue(), value);
            }
        }

        searchHit.getSourceAsMap().putAll(sourceAsMap);
        currentId++;
        intersectHitsList.add(searchHit);
    }
    int totalSize = currentId - 1;
    SearchHit[] unionHitsArr = intersectHitsList.toArray(new SearchHit[totalSize]);
    this.intersectHits = new SearchHits(unionHitsArr, new TotalHits(totalSize, TotalHits.Relation.EQUAL_TO), 1.0f);
}
 

private RestResponseListener<SearchResponse> search(RestChannel channel, Class<T> clazz) {
    return new RestResponseListener<SearchResponse>(channel) {
        @Override
        public RestResponse buildResponse(SearchResponse response) throws Exception {
            if (response.isTimedOut()) {
                return new BytesRestResponse(RestStatus.REQUEST_TIMEOUT, response.toString());
            }

            if (clazz == AnomalyDetector.class) {
                for (SearchHit hit : response.getHits()) {
                    XContentParser parser = XContentType.JSON
                        .xContent()
                        .createParser(
                            channel.request().getXContentRegistry(),
                            LoggingDeprecationHandler.INSTANCE,
                            hit.getSourceAsString()
                        );
                    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);

                    // write back id and version to anomaly detector object
                    ToXContentObject xContentObject = AnomalyDetector.parse(parser, hit.getId(), hit.getVersion());
                    XContentBuilder builder = xContentObject.toXContent(jsonBuilder(), EMPTY_PARAMS);
                    hit.sourceRef(BytesReference.bytes(builder));
                }
            }

            return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), EMPTY_PARAMS));
        }
    };
}
 

private SearchResponse mockSearchResponse(String searchIndex) {
    SearchResponse searchResponse = mock(SearchResponse.class);
    Map<String, DocumentField> fields = new HashMap<>();
    SearchHit searchHit = new SearchHit(101, "xxx", null, fields);
    BytesReference source = new BytesArray("{\"uuid\":\"13456565\",\"createdBy\":\"user\",\"lastUpdateTimestamp\":\"2019-05-20 12:00:00\",\"enabled\":true,\"trusted\":true,\"detectorConfig\":{\"hyperparams\":{\"alpha\":0.5,\"beta\":0.6},\"trainingMetaData\":{\"alpha\":0.5},\"params\":{\"upperWeak\":123}}}}");
    searchHit.sourceRef(source);
    SearchHit[] bunchOfSearchHits = new SearchHit[1];
    bunchOfSearchHits[0] = searchHit;
    SearchHits searchHits = new SearchHits(bunchOfSearchHits, 1, 1);
    when(searchResponse.getHits()).thenReturn(searchHits);
    return searchResponse;
}
 

private SearchResponse mockSearchResponse(String searchIndex) {
    SearchResponse searchResponse = mock(SearchResponse.class);
    Map<String, DocumentField> fields = new HashMap<>();
    SearchHit searchHit = new SearchHit(101, "xxx", null, fields);
    BytesReference source = new BytesArray("{}");
    searchHit.sourceRef(source);
    SearchHit[] bunchOfSearchHits = new SearchHit[1];
    bunchOfSearchHits[0] = searchHit;
    SearchHits searchHits = new SearchHits(bunchOfSearchHits, 1, 1);
    when(searchResponse.getHits()).thenReturn(searchHits);
    return searchResponse;
}
 

public static void givenSearchResultForDocuments(PluginClient client, String indexPattern, Map<String, BytesReference> docs) {
    List<SearchHit> hits = new ArrayList<>(docs.size());
    for (Map.Entry<String, BytesReference> entry : docs.entrySet()) {
        SearchHit hit = new SearchHit(1, entry.getKey(), null, null);
        hit.sourceRef(entry.getValue());
        hits.add(hit);
    }
    SearchHits searchHits = new SearchHits(hits.toArray(new SearchHit[hits.size()]), hits.size(), 1.0f);
    SearchResponseSections sections = new SearchResponseSections(searchHits, null, null, false, Boolean.FALSE, null,
            0);
    ShardSearchFailure[] failures = null;
    SearchResponse response = new SearchResponse(sections, "", 0, 0, 0, 0L, failures);

    when(client.search(anyString(), anyString(),anyInt(), anyBoolean())).thenReturn(response);
}
 

private void fillInternalSearchHits(List<SearchHit> unionHits, SearchHit[] hits, Map<String, String> fieldNameToAlias) {
    for(SearchHit hit : hits){
        SearchHit searchHit = new SearchHit(currentId, hit.getId(), new Text(hit.getType()), hit.getFields(), null);
        searchHit.sourceRef(hit.getSourceRef());
        searchHit.getSourceAsMap().clear();
        Map<String, Object> sourceAsMap = hit.getSourceAsMap();
        if(!fieldNameToAlias.isEmpty()){
            updateFieldNamesToAlias(sourceAsMap, fieldNameToAlias);
        }
        searchHit.getSourceAsMap().putAll(sourceAsMap);
        currentId++;
        unionHits.add(searchHit);
    }
}
 

protected SearchHit createUnmachedResult( List<Field> secondTableReturnedFields, int docId, String t1Alias, String t2Alias, SearchHit hit) {
    String unmatchedId = hit.getId() + "|0";
    Text unamatchedType = new Text(hit.getType() + "|null");

    SearchHit searchHit = new SearchHit(docId, unmatchedId, unamatchedType, hit.getFields(), null);

    searchHit.sourceRef(hit.getSourceRef());
    searchHit.getSourceAsMap().clear();
    searchHit.getSourceAsMap().putAll(hit.getSourceAsMap());
    Map<String,Object> emptySecondTableHitSource = createNullsSource(secondTableReturnedFields);

    mergeSourceAndAddAliases(emptySecondTableHitSource, searchHit,t1Alias,t2Alias);

    return searchHit;
}
 

private void fillMinusHitsFromResults(Set<ComperableHitResult> comperableHitResults) {
    int currentId = 1;
    List<SearchHit> minusHitsList = new ArrayList<>();
    for(ComperableHitResult result : comperableHitResults){
        ArrayList<Object> values = new ArrayList<Object>();
        values.add(result);
        SearchHit originalHit = result.getOriginalHit();
        SearchHit searchHit = new SearchHit(currentId,originalHit.getId(), new Text(originalHit.getType()), originalHit.getFields(), null);
        searchHit.sourceRef(originalHit.getSourceRef());
        searchHit.getSourceAsMap().clear();
        Map<String, Object> sourceAsMap = result.getFlattenMap();
        for(Map.Entry<String,String> entry : this.builder.getFirstTableFieldToAlias().entrySet()){
            if(sourceAsMap.containsKey(entry.getKey())){
                Object value = sourceAsMap.get(entry.getKey());
                sourceAsMap.remove(entry.getKey());
                sourceAsMap.put(entry.getValue(),value);
            }
        }

        searchHit.getSourceAsMap().putAll(sourceAsMap);
        currentId++;
        minusHitsList.add(searchHit);
    }
    int totalSize = currentId - 1;
    SearchHit[] unionHitsArr = minusHitsList.toArray(new SearchHit[totalSize]);
    this.minusHits = new SearchHits(unionHitsArr, new TotalHits(totalSize, TotalHits.Relation.EQUAL_TO), 1.0f);
}
 

private SearchHit getMergedHit(int currentCombinedResults, String t1Alias, String t2Alias, SearchHit hitFromFirstTable, SearchHit matchedHit) {
    onlyReturnedFields(matchedHit.getSourceAsMap(), nestedLoopsRequest.getSecondTable().getReturnedFields(),nestedLoopsRequest.getSecondTable().getOriginalSelect().isSelectAll());
    SearchHit searchHit = new SearchHit(currentCombinedResults, hitFromFirstTable.getId() + "|" + matchedHit.getId(), new Text(hitFromFirstTable.getType() + "|" + matchedHit.getType()), hitFromFirstTable.getFields(), null);
    searchHit.sourceRef(hitFromFirstTable.getSourceRef());
    searchHit.getSourceAsMap().clear();
    searchHit.getSourceAsMap().putAll(hitFromFirstTable.getSourceAsMap());

    mergeSourceAndAddAliases(matchedHit.getSourceAsMap(), searchHit, t1Alias, t2Alias);
    return searchHit;
}
 

private List<SearchHit> createCombinedResults( TableInJoinRequestBuilder secondTableRequest) {
        List<SearchHit> combinedResult = new ArrayList<>();
        int resultIds = 0;
        int totalLimit = this.requestBuilder.getTotalLimit();
        Integer hintLimit = secondTableRequest.getHintLimit();
        SearchResponse searchResponse;
        boolean finishedScrolling;
        if (hintLimit != null && hintLimit < MAX_RESULTS_ON_ONE_FETCH) {
            searchResponse = secondTableRequest.getRequestBuilder().setSize(hintLimit).get();
            finishedScrolling = true;
        } else {
            searchResponse = secondTableRequest.getRequestBuilder()
                    .setScroll(new TimeValue(60000))
                    .setSize(MAX_RESULTS_ON_ONE_FETCH).get();
            //es5.0 no need to scroll again!
//            searchResponse = client.prepareSearchScroll(searchResponse.getScrollId()).setScroll(new TimeValue(600000)).get();
            finishedScrolling = false;
        }
        updateMetaSearchResults(searchResponse);

        boolean limitReached = false;
        int fetchedSoFarFromSecondTable = 0;
        while (!limitReached) {
            SearchHit[] secondTableHits = searchResponse.getHits().getHits();
            fetchedSoFarFromSecondTable += secondTableHits.length;
            for (SearchHit secondTableHit : secondTableHits) {
                if (limitReached) break;
                //todo: need to run on comparisons. for each comparison check if exists and add.
                HashMap<String, List<Map.Entry<Field, Field>>> comparisons = this.hashJoinComparisonStructure.getComparisons();

                for (Map.Entry<String, List<Map.Entry<Field, Field>>> comparison : comparisons.entrySet()) {
                    String comparisonID = comparison.getKey();
                    List<Map.Entry<Field, Field>> t1ToT2FieldsComparison = comparison.getValue();
                    String key = getComparisonKey(t1ToT2FieldsComparison, secondTableHit, false, null);

                    SearchHitsResult searchHitsResult = this.hashJoinComparisonStructure.searchForMatchingSearchHits(comparisonID, key);

                    if (searchHitsResult != null && searchHitsResult.getSearchHits().size() > 0) {
                        searchHitsResult.setMatchedWithOtherTable(true);
                        List<SearchHit> searchHits = searchHitsResult.getSearchHits();
                        for (SearchHit matchingHit : searchHits) {
                            String combinedId = matchingHit.getId() + "|" + secondTableHit.getId();
                            //in order to prevent same matching when using OR on hashJoins.
                            if(this.alreadyMatched.contains(combinedId)){
                                continue;
                            }
                            else {
                                this.alreadyMatched.add(combinedId);
                            }

                            Map<String,Object> copiedSource = new HashMap<String,Object>();
                            copyMaps(copiedSource,secondTableHit.getSourceAsMap());
                            onlyReturnedFields(copiedSource, secondTableRequest.getReturnedFields(),secondTableRequest.getOriginalSelect().isSelectAll());



                            SearchHit searchHit = new SearchHit(matchingHit.docId(), combinedId, new Text(matchingHit.getType() + "|" + secondTableHit.getType()), matchingHit.getFields(), null);
                            searchHit.sourceRef(matchingHit.getSourceRef());
                            searchHit.getSourceAsMap().clear();
                            searchHit.getSourceAsMap().putAll(matchingHit.getSourceAsMap());
                            String t1Alias = requestBuilder.getFirstTable().getAlias();
                            String t2Alias = requestBuilder.getSecondTable().getAlias();
                            mergeSourceAndAddAliases(copiedSource, searchHit, t1Alias, t2Alias);

                            combinedResult.add(searchHit);
                            resultIds++;
                            if (resultIds >= totalLimit) {
                                limitReached = true;
                                break;
                            }
                        }
                    }
                }
            }
            if (!finishedScrolling) {
                if (secondTableHits.length > 0 && (hintLimit == null || fetchedSoFarFromSecondTable >= hintLimit)) {
                    searchResponse = client.prepareSearchScroll(searchResponse.getScrollId()).setScroll(new TimeValue(600000)).execute().actionGet();
                } else break;
            } else {
                break;
            }
        }
        return combinedResult;
    }