Examples with CompositeAggregation org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregation used on opensource projects

Example 1 with CompositeAggregation

use of org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregation in project openk9 by smclab.

the class SuggestionsHTTPHandler method _getSuggestionsResponse.

private SuggestionsResponse _getSuggestionsResponse(List<Datasource> datasourceList, PluginDriverDTOList pluginDriverDTOList, SearchRequest searchRequest, SearchResponse searchResponse, List<SuggestionCategoryField> fields, Map<String, String[]> entityMap) {
    Aggregations aggregations = searchResponse.getAggregations();
    if (aggregations == null) {
        return SuggestionsResponse.of(List.of(), null);
    }
    Map<String, Aggregation> aggregationMap = aggregations.asMap();
    if (!aggregationMap.containsKey("composite")) {
        return SuggestionsResponse.of(List.of(), null);
    }
    CompositeAggregation compositeAggregation = (CompositeAggregation) aggregationMap.get("composite");
    Map<String, Long> fieldNameCategoryIdMap = fields.stream().collect(Collectors.toMap(SuggestionCategoryField::getFieldName, SuggestionCategoryField::getCategoryId, (a1, a2) -> a2));
    Long datasourceIdCategoryId = fieldNameCategoryIdMap.getOrDefault("datasourceId", 1L);
    Long entityIdCategoryId = fieldNameCategoryIdMap.getOrDefault("entities.id", 2L);
    Long entitiesContextCategoryId = fieldNameCategoryIdMap.getOrDefault("entities.context", entityIdCategoryId);
    Long documentTypesCategoryId = fieldNameCategoryIdMap.getOrDefault("documentTypes", 4L);
    List<? extends CompositeAggregation.Bucket> buckets = compositeAggregation.getBuckets();
    LinkedList<Suggestions> suggestions = new LinkedList<>();
    String suggestKeyword = searchRequest.getSuggestKeyword();
    BiConsumer<String, Suggestions> addSuggestions;
    if (suggestKeyword != null) {
        addSuggestions = (key, sugg) -> {
            if (!suggestions.contains(sugg)) {
                if (key.contains(suggestKeyword)) {
                    suggestions.addFirst(sugg);
                }
            }
        };
    } else {
        addSuggestions = (key, sugg) -> {
            if (!suggestions.contains(sugg)) {
                suggestions.add(sugg);
            }
        };
    }
    for (CompositeAggregation.Bucket bucket : buckets) {
        Map<String, Object> keys = new HashMap<>(bucket.getKey());
        for (Map.Entry<String, Object> entry : keys.entrySet()) {
            String key = entry.getKey();
            String value = (String) entry.getValue();
            if (value == null) {
                continue;
            }
            switch(key) {
                case "datasourceId":
                    long datasourceIdL = Long.parseLong(value);
                    _datasource(datasourceList, pluginDriverDTOList, datasourceIdCategoryId, addSuggestions, datasourceIdL);
                    break;
                case "entities.context":
                    break;
                case "entities.id":
                    String[] typeName = entityMap.get(value);
                    if (typeName != null) {
                        String type = typeName[0];
                        String name = typeName[1];
                        String entitiesContext = (String) keys.get("entities.context");
                        if (entitiesContext != null) {
                            addSuggestions.accept(name, Suggestions.entity(value, entitiesContextCategoryId, type, name, entitiesContext));
                        } else {
                            addSuggestions.accept(name, Suggestions.entity(value, entityIdCategoryId, type, name));
                        }
                    }
                    break;
                case "documentTypes":
                    addSuggestions.accept(value, Suggestions.docType(value, documentTypesCategoryId));
                    break;
                default:
                    Long textCategoryId = fieldNameCategoryIdMap.getOrDefault(key, 5L);
                    addSuggestions.accept(value, Suggestions.text(value, textCategoryId, key));
            }
        }
    }
    Map<String, Object> map = compositeAggregation.afterKey();
    String afterKey = null;
    int[] range = searchRequest.getRange();
    if (map != null) {
        afterKey = _jsonFactory.toJson(map);
        afterKey = Base64.getEncoder().encodeToString(afterKey.getBytes(StandardCharsets.UTF_8));
    }
    return SuggestionsResponse.of(suggestions, afterKey);
}

Example 2 with CompositeAggregation

use of org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregation in project openk9 by smclab.

the class SuggestionsHTTPHandler method searchHitToResponseMono.

@Override
protected Mono<Object> searchHitToResponseMono(Tenant tenant, List<Datasource> datasourceList, PluginDriverDTOList pluginDriverDTOList, HttpServerRequest httpServerRequest, SearchRequest searchRequest, SearchResponse searchResponse) {
    Long suggestionCategoryId = searchRequest.getSuggestionCategoryId();
    Mono<List<SuggestionCategoryField>> suggestionCategoryFields;
    if (suggestionCategoryId == null) {
        suggestionCategoryFields = _datasourceClient.findSuggestionCategoryFieldsByTenantId(tenant.getTenantId());
    } else {
        suggestionCategoryFields = _datasourceClient.findSuggestionCategoryFieldsByTenantIdAndCategoryId(tenant.getTenantId(), suggestionCategoryId);
    }
    if (_enableEntityAggregation) {
        return _search.search(factory -> {
            org.elasticsearch.action.search.SearchRequest searchRequestEntity = factory.createSearchRequestEntity(tenant.getTenantId());
            Aggregations aggregations = searchResponse.getAggregations();
            CompositeAggregation compositeAggregation = aggregations.get("composite");
            List<? extends CompositeAggregation.Bucket> buckets = compositeAggregation.getBuckets();
            BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
            buckets.stream().map(bucket -> (String) bucket.getKey().get("entities.id")).filter(Objects::nonNull).distinct().forEach(entityId -> boolQueryBuilder.should(QueryBuilders.matchQuery("id", entityId)));
            if (_log.isDebugEnabled()) {
                _log.debug("entities query: " + boolQueryBuilder);
            }
            SearchSourceBuilder ssb = new SearchSourceBuilder();
            ssb.query(boolQueryBuilder);
            ssb.size(1000);
            ssb.fetchSource(new String[] { "name", "id", "type" }, null);
            return searchRequestEntity.source(ssb);
        }).flatMap(entityResponse -> suggestionCategoryFields.map(fields -> {
            Map<String, String[]> entityMap = new HashMap<>();
            for (SearchHit hit : entityResponse.getHits()) {
                Map<String, Object> sourceAsMap = hit.getSourceAsMap();
                String name = (String) sourceAsMap.get("name");
                String type = (String) sourceAsMap.get("type");
                String entityId = (String) sourceAsMap.get("id");
                entityMap.put(entityId, new String[] { type, name });
            }
            return _getSuggestionsResponse(datasourceList, pluginDriverDTOList, searchRequest, searchResponse, fields, entityMap);
        }));
    } else {
        return suggestionCategoryFields.map(fields -> _getSuggestionsResponse(datasourceList, pluginDriverDTOList, searchRequest, searchResponse, fields, Map.of()));
    }
}

Example 3 with CompositeAggregation

use of org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregation in project openk9 by smclab.

the class SuggestionsHTTPHandler method customizeSearchSourceBuilderMono.

@Override
protected Mono<org.elasticsearch.action.search.SearchRequest> customizeSearchSourceBuilderMono(Tenant tenant, List<Datasource> datasources, SearchRequest searchRequest, List<PluginDriverDTO> documentTypeList, SearchSourceBuilder searchSourceBuilder, org.elasticsearch.action.search.SearchRequest elasticSearchQuery) {
    return Mono.defer(() -> {
        Long suggestionCategoryId = searchRequest.getSuggestionCategoryId();
        if (suggestionCategoryId == null) {
            return _datasourceClient.findSuggestionCategoryFieldsByTenantId(tenant.getTenantId());
        } else {
            return _datasourceClient.findSuggestionCategoryFieldsByTenantIdAndCategoryId(tenant.getTenantId(), suggestionCategoryId);
        }
    }).map(fields -> {
        if (!(fields == null || fields.isEmpty())) {
            Function<String, CompositeValuesSourceBuilder<?>> fieldToTerms = nameField -> new TermsValuesSourceBuilder(nameField).field(nameField).missingBucket(true);
            CompositeAggregationBuilder compositeAggregation = fields.stream().map(SuggestionCategoryField::getFieldName).map(fieldToTerms).collect(Collectors.collectingAndThen(Collectors.toList(), list -> AggregationBuilders.composite("composite", list)));
            String afterKey = searchRequest.getAfterKey();
            if (afterKey != null) {
                byte[] afterKeyDecoded = Base64.getDecoder().decode(afterKey);
                Map<String, Object> map = _jsonFactory.fromJsonMap(new String(afterKeyDecoded), Object.class);
                compositeAggregation.aggregateAfter(map);
            }
            int[] range = searchRequest.getRange();
            if (range != null && range.length == 2) {
                int size = range[1];
                compositeAggregation.size(size);
            }
            searchSourceBuilder.aggregation(compositeAggregation);
        }
        searchSourceBuilder.from(0);
        searchSourceBuilder.size(0);
        searchSourceBuilder.highlighter(null);
        return elasticSearchQuery.source(searchSourceBuilder);
    });
}