Examples with IdsQueryBuilder org.elasticsearch.index.query.IdsQueryBuilder used on opensource projects

Example 6 with IdsQueryBuilder

use of org.elasticsearch.index.query.IdsQueryBuilder in project metron by apache.

the class ElasticsearchDao method searchByGuids.

/**
 * Return the search hit based on the UUID and sensor type.
 * A callback can be specified to transform the hit into a type T.
 * If more than one hit happens, the first one will be returned.
 */
<T> List<T> searchByGuids(Collection<String> guids, Collection<String> sensorTypes, Function<SearchHit, Optional<T>> callback) {
    if (guids == null || guids.isEmpty()) {
        return Collections.EMPTY_LIST;
    }
    QueryBuilder query = null;
    IdsQueryBuilder idsQuery = null;
    if (sensorTypes != null) {
        String[] types = sensorTypes.stream().map(sensorType -> sensorType + "_doc").toArray(String[]::new);
        idsQuery = QueryBuilders.idsQuery(types);
    } else {
        idsQuery = QueryBuilders.idsQuery();
    }
    for (String guid : guids) {
        query = idsQuery.addIds(guid);
    }
    SearchRequestBuilder request = client.prepareSearch().setQuery(query).setSize(guids.size());
    org.elasticsearch.action.search.SearchResponse response = request.get();
    SearchHits hits = response.getHits();
    List<T> results = new ArrayList<>();
    for (SearchHit hit : hits) {
        Optional<T> result = callback.apply(hit);
        if (result.isPresent()) {
            results.add(result.get());
        }
    }
    return results;
}

Example 7 with IdsQueryBuilder

use of org.elasticsearch.index.query.IdsQueryBuilder in project pyramid by cheng-li.

the class LabelDistribution method getLabelDistribution.

public static long[] getLabelDistribution(ESIndex index, String labelField, String[] ids, LabelTranslator labelTranslator) {
    int numClasses = labelTranslator.getNumClasses();
    long[] distribution = new long[numClasses];
    IdsQueryBuilder idsFilterBuilder = new IdsQueryBuilder(index.getDocumentType());
    idsFilterBuilder.addIds(ids);
    SearchResponse response = index.getClient().prepareSearch(index.getIndexName()).setSize(0).setTrackScores(false).setFetchSource(false).setExplain(false).setFetchSource(false).setQuery(QueryBuilders.boolQuery().must(idsFilterBuilder).should(QueryBuilders.matchAllQuery())).addAggregation(terms("agg").field(labelField).size(Integer.MAX_VALUE)).execute().actionGet();
    Terms terms = response.getAggregations().get("agg");
    Collection<Terms.Bucket> buckets = terms.getBuckets();
    for (Terms.Bucket bucket : buckets) {
        String extLabel = (String) bucket.getKey();
        long count = bucket.getDocCount();
        int classIndex = labelTranslator.toIntLabel(extLabel);
        distribution[classIndex] = count;
    }
    return distribution;
}

Example 8 with IdsQueryBuilder

use of org.elasticsearch.index.query.IdsQueryBuilder in project muikku by otavanopisto.

the class ElasticSearchProvider method getUserWorkspaces.

private Set<SchoolDataIdentifier> getUserWorkspaces(SchoolDataIdentifier userIdentifier) {
    Set<SchoolDataIdentifier> result = new HashSet<>();
    IdsQueryBuilder query = idsQuery("User");
    query.addIds(String.format("%s/%s", userIdentifier.getIdentifier(), userIdentifier.getDataSource()));
    SearchResponse response = elasticClient.prepareSearch("muikku").setTypes("User").setQuery(query).addField("workspaces").setSize(1).execute().actionGet();
    SearchHit[] hits = response.getHits().getHits();
    for (SearchHit hit : hits) {
        Map<String, SearchHitField> fields = hit.getFields();
        SearchHitField workspaceField = fields.get("workspaces");
        if (workspaceField != null && workspaceField.getValues() != null) {
            for (Object value : workspaceField.getValues()) {
                if (value instanceof Number) {
                    Long workspaceEntityId = ((Number) value).longValue();
                    WorkspaceEntity workspaceEntity = workspaceEntityController.findWorkspaceEntityById(workspaceEntityId);
                    if (workspaceEntity != null) {
                        result.add(new SchoolDataIdentifier(workspaceEntity.getIdentifier(), workspaceEntity.getDataSource().getIdentifier()));
                    }
                }
            }
        }
    }
    return result;
}

Example 9 with IdsQueryBuilder

use of org.elasticsearch.index.query.IdsQueryBuilder in project muikku by otavanopisto.

the class ElasticSearchProvider method searchWorkspaces.

@Override
public SearchResult searchWorkspaces(String schoolDataSource, List<String> subjects, List<String> identifiers, List<SchoolDataIdentifier> educationTypes, List<SchoolDataIdentifier> curriculumIdentifiers, String freeText, List<WorkspaceAccess> accesses, SchoolDataIdentifier accessUser, boolean includeUnpublished, int start, int maxResults, List<Sort> sorts) {
    if (identifiers != null && identifiers.isEmpty()) {
        return new SearchResult(0, 0, new ArrayList<Map<String, Object>>(), 0);
    }
    BoolQueryBuilder query = boolQuery();
    freeText = sanitizeSearchString(freeText);
    try {
        if (!includeUnpublished) {
            query.must(termQuery("published", Boolean.TRUE));
        }
        if (accesses != null) {
            BoolQueryBuilder accessQuery = boolQuery();
            for (WorkspaceAccess access : accesses) {
                switch(access) {
                    case LOGGED_IN:
                    case ANYONE:
                        accessQuery.should(termQuery("access", access));
                        break;
                    case MEMBERS_ONLY:
                        BoolQueryBuilder memberQuery = boolQuery();
                        IdsQueryBuilder idsQuery = idsQuery("Workspace");
                        for (SchoolDataIdentifier userWorkspace : getUserWorkspaces(accessUser)) {
                            idsQuery.addIds(String.format("%s/%s", userWorkspace.getIdentifier(), userWorkspace.getDataSource()));
                        }
                        memberQuery.must(idsQuery);
                        memberQuery.must(termQuery("access", access));
                        accessQuery.should(memberQuery);
                        break;
                }
            }
            query.must(accessQuery);
        }
        if (StringUtils.isNotBlank(schoolDataSource)) {
            query.must(termQuery("schoolDataSource", schoolDataSource.toLowerCase()));
        }
        if (subjects != null && !subjects.isEmpty()) {
            query.must(termsQuery("subjectIdentifier", subjects));
        }
        if (educationTypes != null && !educationTypes.isEmpty()) {
            List<String> educationTypeIds = new ArrayList<>(educationTypes.size());
            for (SchoolDataIdentifier educationType : educationTypes) {
                educationTypeIds.add(educationType.toId());
            }
            query.must(termsQuery("educationTypeIdentifier.untouched", educationTypeIds));
        }
        if (!CollectionUtils.isEmpty(curriculumIdentifiers)) {
            List<String> curriculumIds = new ArrayList<>(curriculumIdentifiers.size());
            for (SchoolDataIdentifier curriculumIdentifier : curriculumIdentifiers) {
                curriculumIds.add(curriculumIdentifier.toId());
            }
            query.must(boolQuery().should(termsQuery("curriculumIdentifiers.untouched", curriculumIds)).should(boolQuery().mustNot(existsQuery("curriculumIdentifiers"))).minimumNumberShouldMatch(1));
        }
        if (identifiers != null) {
            query.must(termsQuery("identifier", identifiers));
        }
        if (StringUtils.isNotBlank(freeText)) {
            String[] words = freeText.split(" ");
            for (int i = 0; i < words.length; i++) {
                if (StringUtils.isNotBlank(words[i])) {
                    query.must(boolQuery().should(prefixQuery("name", words[i])).should(prefixQuery("description", words[i])).should(prefixQuery("subject", words[i])).should(prefixQuery("staffMembers.firstName", words[i])).should(prefixQuery("staffMembers.lastName", words[i])));
                }
            }
        }
        SearchRequestBuilder requestBuilder = elasticClient.prepareSearch("muikku").setTypes("Workspace").setFrom(start).setSize(maxResults);
        if (sorts != null && !sorts.isEmpty()) {
            for (Sort sort : sorts) {
                requestBuilder.addSort(sort.getField(), SortOrder.valueOf(sort.getOrder().name()));
            }
        }
        SearchResponse response = requestBuilder.setQuery(query).execute().actionGet();
        List<Map<String, Object>> searchResults = new ArrayList<Map<String, Object>>();
        SearchHits searchHits = response.getHits();
        long totalHitCount = searchHits.getTotalHits();
        SearchHit[] results = searchHits.getHits();
        for (SearchHit hit : results) {
            Map<String, Object> hitSource = hit.getSource();
            hitSource.put("indexType", hit.getType());
            searchResults.add(hitSource);
        }
        SearchResult result = new SearchResult(start, maxResults, searchResults, totalHitCount);
        return result;
    } catch (Exception e) {
        logger.log(Level.SEVERE, "ElasticSearch query failed unexpectedly", e);
        return new SearchResult(0, 0, new ArrayList<Map<String, Object>>(), 0);
    }
}

Example 10 with IdsQueryBuilder

use of org.elasticsearch.index.query.IdsQueryBuilder in project muikku by otavanopisto.

the class ElasticSearchProvider method searchUsers.

@Override
public SearchResult searchUsers(String text, String[] textFields, Collection<EnvironmentRoleArchetype> archetypes, Collection<Long> groups, Collection<Long> workspaces, Collection<SchoolDataIdentifier> userIdentifiers, Boolean includeInactiveStudents, Boolean includeHidden, Boolean onlyDefaultUsers, int start, int maxResults, Collection<String> fields, Collection<SchoolDataIdentifier> excludeSchoolDataIdentifiers, Date startedStudiesBefore, Date studyTimeEndsBefore) {
    try {
        long now = OffsetDateTime.now().toEpochSecond();
        text = sanitizeSearchString(text);
        BoolQueryBuilder query = boolQuery();
        if (!Boolean.TRUE.equals(includeHidden)) {
            query.mustNot(termQuery("hidden", true));
        }
        if (Boolean.TRUE.equals(onlyDefaultUsers)) {
            query.must(termQuery("isDefaultIdentifier", true));
        }
        if (StringUtils.isNotBlank(text) && !ArrayUtils.isEmpty(textFields)) {
            String[] words = text.split(" ");
            for (int i = 0; i < words.length; i++) {
                if (StringUtils.isNotBlank(words[i])) {
                    BoolQueryBuilder fieldBuilder = boolQuery();
                    for (String textField : textFields) {
                        fieldBuilder.should(prefixQuery(textField, words[i]));
                    }
                    query.must(fieldBuilder);
                }
            }
        }
        if (excludeSchoolDataIdentifiers != null) {
            IdsQueryBuilder excludeIdsQuery = idsQuery("User");
            for (SchoolDataIdentifier excludeSchoolDataIdentifier : excludeSchoolDataIdentifiers) {
                excludeIdsQuery.addIds(String.format("%s/%s", excludeSchoolDataIdentifier.getIdentifier(), excludeSchoolDataIdentifier.getDataSource()));
            }
            query.mustNot(excludeIdsQuery);
        }
        if (startedStudiesBefore != null) {
            query.must(rangeQuery("studyStartDate").lt((long) startedStudiesBefore.getTime() / 1000));
        }
        if (studyTimeEndsBefore != null) {
            query.must(rangeQuery("studyTimeEnd").lt((long) studyTimeEndsBefore.getTime() / 1000));
        }
        if (archetypes != null) {
            List<String> archetypeNames = new ArrayList<>(archetypes.size());
            for (EnvironmentRoleArchetype archetype : archetypes) {
                archetypeNames.add(archetype.name().toLowerCase());
            }
            query.must(termsQuery("archetype", archetypeNames.toArray(new String[0])));
        }
        if (groups != null) {
            query.must(termsQuery("groups", ArrayUtils.toPrimitive(groups.toArray(new Long[0]))));
        }
        if (workspaces != null) {
            query.must(termsQuery("workspaces", ArrayUtils.toPrimitive(workspaces.toArray(new Long[0]))));
        }
        if (userIdentifiers != null) {
            IdsQueryBuilder includeIdsQuery = idsQuery("User");
            for (SchoolDataIdentifier userIdentifier : userIdentifiers) {
                includeIdsQuery.addIds(String.format("%s/%s", userIdentifier.getIdentifier(), userIdentifier.getDataSource()));
            }
            query.must(includeIdsQuery);
        }
        if (includeInactiveStudents == false) {
            /**
             * List only active users.
             *
             * Active user is
             * - staff member (teacher, manager, study guider, study programme leader, administrator)
             * - student that has study start date (in the past) and no study end date
             * - student that has study start date (in the past) and study end date in the future
             * - student that has no study start and end date but belongs to an active workspace
             *
             * Active workspace is
             * - published and
             * - either has no start/end date or current date falls between them
             */
            Set<Long> activeWorkspaceEntityIds = getActiveWorkspaces();
            query.must(boolQuery().should(termsQuery("archetype", EnvironmentRoleArchetype.TEACHER.name().toLowerCase(), EnvironmentRoleArchetype.MANAGER.name().toLowerCase(), EnvironmentRoleArchetype.STUDY_GUIDER.name().toLowerCase(), EnvironmentRoleArchetype.STUDY_PROGRAMME_LEADER.name().toLowerCase(), EnvironmentRoleArchetype.ADMINISTRATOR.name().toLowerCase())).should(boolQuery().must(termQuery("archetype", EnvironmentRoleArchetype.STUDENT.name().toLowerCase())).must(existsQuery("studyStartDate")).must(rangeQuery("studyStartDate").lte(now)).mustNot(existsQuery("studyEndDate"))).should(boolQuery().must(termQuery("archetype", EnvironmentRoleArchetype.STUDENT.name().toLowerCase())).must(existsQuery("studyStartDate")).must(rangeQuery("studyStartDate").lte(now)).must(existsQuery("studyEndDate")).must(rangeQuery("studyEndDate").gte(now))).should(boolQuery().must(termQuery("archetype", EnvironmentRoleArchetype.STUDENT.name().toLowerCase())).mustNot(existsQuery("studyEndDate")).mustNot(existsQuery("studyStartDate")).must(termsQuery("workspaces", ArrayUtils.toPrimitive(activeWorkspaceEntityIds.toArray(new Long[0]))))));
        }
        SearchRequestBuilder requestBuilder = elasticClient.prepareSearch("muikku").setTypes("User").setFrom(start).setSize(maxResults);
        if (CollectionUtils.isNotEmpty(fields)) {
            requestBuilder.addFields(fields.toArray(new String[0]));
        }
        SearchResponse response = requestBuilder.setQuery(query).addSort("_score", SortOrder.DESC).addSort("lastName", SortOrder.ASC).addSort("firstName", SortOrder.ASC).execute().actionGet();
        List<Map<String, Object>> searchResults = new ArrayList<Map<String, Object>>();
        SearchHits searchHits = response.getHits();
        long totalHitCount = searchHits.getTotalHits();
        SearchHit[] results = searchHits.getHits();
        for (SearchHit hit : results) {
            Map<String, Object> hitSource = hit.getSource();
            if (hitSource == null) {
                hitSource = new HashMap<>();
                for (String key : hit.getFields().keySet()) {
                    hitSource.put(key, hit.getFields().get(key).getValue().toString());
                }
            }
            hitSource.put("indexType", hit.getType());
            searchResults.add(hitSource);
        }
        SearchResult result = new SearchResult(start, maxResults, searchResults, totalHitCount);
        return result;
    } catch (Exception e) {
        logger.log(Level.SEVERE, "ElasticSearch query failed unexpectedly", e);
        return new SearchResult(0, 0, new ArrayList<Map<String, Object>>(), 0);
    }
}