In spring data elasticsearch, can aggregation queries be put in a respository implementation?

@Repository public interface SonarMetricRepository extends ElasticsearchRepository<Article, String> { @Query("{\"bool\": {\"must\": {\"match\": {\"authors.name\": \"?0\"}}, \"filter\": {\"term\": {\"tags\": \"?1\" }}}}") Page<Article> findByAuthorsNameAndFilteredTagQuery(String name, String tag, Pageable pageable); }

SerialDiffPipelineAggregationBuilder serialDiffPipelineAggregationBuilder = PipelineAggregatorBuilders .diff("Percent_Change", "avg_coverage") .lag(1); AvgAggregationBuilder averageCoverageAggregationBuilder = AggregationBuilders .avg("avg_coverage") .field("coverage"); AggregationBuilder coverageHistoryAggregationBuilder = AggregationBuilders .dateHistogram("coverage_history") .field("@timestamp") .calendarInterval(DateHistogramInterval.WEEK) .subAggregation(averageCoverageAggregationBuilder) .subAggregation(serialDiffPipelineAggregationBuilder); TermsAggregationBuilder sonarProjectKeyAggregationBuilder = AggregationBuilders .terms("project_key") .field("key.keyword") .subAggregation(coverageHistoryAggregationBuilder); BoolQueryBuilder searchQuery = new BoolQueryBuilder() .filter(matchAllQuery()) .filter(matchPhraseQuery("name.keyword", "my-sample-sonar-project")) .filter(rangeQuery("@timestamp") .format("strict_date_optional_time") .gte("2020-07-08T19:29:12.054Z") .lte("2020-07-15T19:29:12.055Z")); // Join query and aggregation together SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder() .query(searchQuery) .aggregation(sonarProjectKeyAggregationBuilder); SearchRequest searchRequest = new SearchRequest("sonarmetrics").source(searchSourceBuilder); SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);

Ok, so if I get it right, you want to add an aggregation to a repository query. This is not possible with the methods that are automatically created by Spring Data Elasticsearch, but it is not too hard to implement it.

To show you how to do this, I use a simpler example, where we have defined a Person entity:

@Document(indexName = "person")
public class Person {

    @Id
    @Nullable
    private Long id;

    @Field(type = FieldType.Text, fielddata = true)
    @Nullable
    private String lastName;

    @Field(type = FieldType.Text, fielddata = true)
    @Nullable
    private String firstName;

    // getter/setter
}

There also is a corresponding repository:

public interface PersonRepository extends ElasticsearchRepository<Person, Long>{
}

We now want to extend this repository to be able to search for persons with a first name, and also return for these persons the top 10 ten last names with the count (a terms aggs on the lastNames).

The first thing to do is to define a customization repository that describes the method you need:

interface PersonCustomRepository {
    SearchPage<Person> findByFirstNameWithLastNameCounts(String firstName, Pageable pageable);
}

We want to pass in a Pageable so that the methods returns pages of data. We return a SearchPage object check the documentation on return types that will contain paging information along with a SearchHits<Person>. This object then has the aggregations information and the result data.

We then change the PersonRepository to extend this new interface:

public interface PersonRepository extends ElasticsearchRepository<Person, Long>, PersonCustomRepository {
}

Of course we now need to provide an implementation in a class named PersonCustomRepositoryImpl (this must be named like the interface with Impl added):

public class PersonCustomRepositoryImpl implements PersonCustomRepository {

    private final ElasticsearchOperations operations;

    public PersonCustomRepositoryImpl(ElasticsearchOperations operations) { // let Spring inject an operations which we use to do the work
        this.operations = operations;
    }

    @Override
    public SearchPage<Person> findByFirstNameWithLastNameCounts(String firstName, Pageable pageable) {

        Query query = new NativeSearchQueryBuilder()                       // we build a Elasticsearch native query
            .addAggregation(terms("lastNames").field("lastName").size(10)) // add the aggregation
            .withQuery(QueryBuilders.matchQuery("firstName", firstName))   // add the query part
            .withPageable(pageable)                                        // add the requested page
            .build();

        SearchHits<Person> searchHits = operations.search(query, Person.class);  // send it of and get the result

        return SearchHitSupport.searchPageFor(searchHits, pageable);  // convert the result to a SearchPage
    }
}

That's all for the implementation of the search. Now the repository has this additional method. How to use it?

For this demo I assume we have a REST controller that takes a name and returns a pair of:

the found persons as a list of SearchHit<Person> objects
a Map<String, Long> containing the last names and their count

This could be implemented as follows, comments describe what is done:

@GetMapping("persons/firstNameWithLastNameCounts/{firstName}")
public Pair<List<SearchHit<Person>>, Map<String, Long>> firstNameWithLastNameCounts(@PathVariable("firstName") String firstName) {

    // helper function to get the lastName counts from an Elasticsearch Aggregations
    // Spring Data Elasticsearch does not have functions for that, so we need to know what is coming back
    Function<Aggregations, Map<String, Long>> getLastNameCounts = aggregations -> {
        if (aggregations != null) {
            Aggregation lastNames = aggregations.get("lastNames");
            if (lastNames != null) {
                List<? extends Terms.Bucket> buckets = ((Terms) lastNames).getBuckets();
                if (buckets != null) {
                    return buckets.stream().collect(Collectors.toMap(Terms.Bucket::getKeyAsString, Terms.Bucket::getDocCount));
                }
            }
        }
        return Collections.emptyMap();
    };

    // the parts of the returned object
    Map<String, Long> lastNameCounts = null;
    List<SearchHit<Person>> searchHits = new ArrayList<>();

    // request pages of size 1000
    Pageable pageable = PageRequest.of(0, 1000);
    boolean fetchMore = true;
    while (fetchMore) {
        // call the custom method implementation
        SearchPage<Person> searchPage = personRepository.findByFirstNameWithLastNameCounts(firstName, pageable);

        // get the aggregations on the first call, will be the same on the other pages
        if (lastNameCounts == null) {
            Aggregations aggregations = searchPage.getSearchHits().getAggregations();
            lastNameCounts = getLastNameCounts.apply(aggregations);
        }

        // collect the returned data
        if (searchPage.hasContent()) {
            searchHits.addAll(searchPage.getContent());
        }

        pageable = searchPage.nextPageable();
        fetchMore = searchPage.hasNext();
    }

    // return the collected stuff
    return Pair.of(searchHits, lastNameCounts);
}

I hope this gives an idea on how to implement custom repository functions and add functionality not provided out of the box.

Recommended topics

Hot tags