Merge #792

792: Support facet distribution for federated search r=curquiza a=Barabasbalazs

# Pull Request

## Related issue
Fixes #787 

## What does this PR do?
- Adds support for `federation.facetsByIndex` in the `POST /multi-search` route.
- Allows for merging of returned facets with `federation.mergeFacets` in the `POST /multi-search` route

## PR checklist
Please check if your PR fulfills the following requirements:
- [x] Does this PR fix an existing issue, or have you listed the changes applied in the PR description (and why they are needed)?
- [x] Have you read the contributing guidelines?
- [x] Have you made sure that the title is accurate and descriptive of the changes?

Thank you so much for contributing to Meilisearch!


Co-authored-by: Balazs Barabas <[email protected]>

Author: meili-bors[bot]

src/main/java/com/meilisearch/sdk/MergeFacets.java

@@ -0,0 +1,14 @@
+package com.meilisearch.sdk;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.Setter;
+
+@Getter
+@Setter
+@AllArgsConstructor
+@NoArgsConstructor
+public class MergeFacets {
+    private Integer maxValuesPerFacet;
+}

src/main/java/com/meilisearch/sdk/MultiSearchFederation.java

@@ -1,11 +1,16 @@
 package com.meilisearch.sdk;
 
+import java.util.Map;
+import lombok.Getter;
 import org.json.JSONObject;
 
+@Getter
 public class MultiSearchFederation {
 
     private Integer limit;
     private Integer offset;
+    private MergeFacets mergeFacets;
+    private Map<String, String[]> facetsByIndex;
 
     public MultiSearchFederation setLimit(Integer limit) {
         this.limit = limit;
@@ -17,12 +22,14 @@ public MultiSearchFederation setOffset(Integer offset) {
         return this;
     }
 
-    public Integer getLimit() {
-        return this.limit;
+    public MultiSearchFederation setMergeFacets(MergeFacets mergeFacets) {
+        this.mergeFacets = mergeFacets;
+        return this;
     }
 
-    public Integer getOffset() {
-        return this.offset;
+    public MultiSearchFederation setFacetsByIndex(Map<String, String[]> facetsByIndex) {
+        this.facetsByIndex = facetsByIndex;
+        return this;
     }
 
     /**

src/main/java/com/meilisearch/sdk/MultiSearchRequest.java

@@ -1,7 +1,6 @@
 package com.meilisearch.sdk;
 
 import java.util.ArrayList;
-import lombok.*;
 
 public class MultiSearchRequest {
     private ArrayList<IndexSearchRequest> queries;

src/main/java/com/meilisearch/sdk/model/FacetsByIndexInfo.java

@@ -0,0 +1,10 @@
+package com.meilisearch.sdk.model;
+
+import java.util.HashMap;
+import lombok.Getter;
+
+@Getter
+public class FacetsByIndexInfo {
+    private HashMap<String, HashMap<String, Integer>> distribution;
+    private HashMap<String, FacetRating> stats;
+}

src/main/java/com/meilisearch/sdk/model/MultiSearchResult.java

@@ -15,13 +15,14 @@
 public class MultiSearchResult implements Searchable {
     String indexUid;
     ArrayList<HashMap<String, Object>> hits;
-    Object facetDistribution;
+    HashMap<String, HashMap<String, Integer>> facetDistribution;
     HashMap<String, FacetRating> facetStats;
     int processingTimeMs;
     String query;
     int offset;
     int limit;
     int estimatedTotalHits;
+    HashMap<String, FacetsByIndexInfo> facetsByIndex;
 
     public MultiSearchResult() {}
 }

src/test/java/com/meilisearch/integration/SearchTest.java

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
@@ -86,8 +87,8 @@ public void testSearchOffset() throws Exception {
         SearchRequest searchRequest = SearchRequest.builder().q("a").offset(20).build();
         SearchResult searchResult = (SearchResult) index.search(searchRequest);
 
-        assertThat(searchResult.getHits(), hasSize(10));
-        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(30)));
+        assertThat(searchResult.getHits(), hasSize(11));
+        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(31)));
     }
 
     /** Test search limit */
@@ -105,7 +106,7 @@ public void testSearchLimit() throws Exception {
         SearchResult searchResult = (SearchResult) index.search(searchRequest);
 
         assertThat(searchResult.getHits(), hasSize(2));
-        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(30)));
+        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(31)));
     }
 
     /** Test search attributesToRetrieve */
@@ -310,7 +311,7 @@ public void testSearchWithMatchingStrategy() throws Exception {
         SearchResult searchResult = (SearchResult) index.search(searchRequest);
 
         assertThat(searchResult.getHits(), hasSize(20));
-        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(21)));
+        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(22)));
     }
 
     /** Test search with frequency matching strategy */
@@ -685,12 +686,10 @@ public void testRawSearchSortWithPlaceHolder() throws Exception {
         Results resGson = jsonGson.decode(index.rawSearch(searchRequest), Results.class);
 
         assertThat(resGson.hits, is(arrayWithSize(20)));
-        assertThat(resGson.hits[0].getId(), is(equalTo("155")));
-        assertThat(resGson.hits[0].getTitle(), is(equalTo("The Dark Knight")));
-        assertThat(resGson.hits[1].getId(), is(equalTo("671")));
-        assertThat(
-                resGson.hits[1].getTitle(),
-                is(equalTo("Harry Potter and the Philosopher's Stone")));
+        assertThat(resGson.hits[0].getId(), is(equalTo("2")));
+        assertThat(resGson.hits[0].getTitle(), is(equalTo("Hobbit")));
+        assertThat(resGson.hits[1].getId(), is(equalTo("155")));
+        assertThat(resGson.hits[1].getTitle(), is(equalTo("The Dark Knight")));
     }
 
     /** Test search matches */
@@ -728,7 +727,7 @@ public void testSearchPage() throws Exception {
         assertThat(searchResult.getHits(), hasSize(20));
         assertThat(searchResult.getPage(), is(equalTo(1)));
         assertThat(searchResult.getHitsPerPage(), is(equalTo(20)));
-        assertThat(searchResult.getTotalHits(), is(equalTo(30)));
+        assertThat(searchResult.getTotalHits(), is(equalTo(31)));
         assertThat(searchResult.getTotalPages(), is(equalTo(2)));
     }
 
@@ -749,8 +748,8 @@ public void testSearchPagination() throws Exception {
         assertThat(searchResult.getHits(), hasSize(2));
         assertThat(searchResult.getPage(), is(equalTo(2)));
         assertThat(searchResult.getHitsPerPage(), is(equalTo(2)));
-        assertThat(searchResult.getTotalHits(), is(equalTo(30)));
-        assertThat(searchResult.getTotalPages(), is(equalTo(15)));
+        assertThat(searchResult.getTotalHits(), is(equalTo(31)));
+        assertThat(searchResult.getTotalPages(), is(equalTo(16)));
     }
 
     /** Test place holder search */
@@ -931,6 +930,162 @@ public void testMultiSearchWithDistinct() throws Exception {
         }
     }
 
+    @Test
+    public void testMultiSearchWithFacetsByIndex() {
+        HashSet<String> indexUids = new HashSet<String>();
+        indexUids.add("movies");
+        indexUids.add("nestedMovies");
+
+        for (String indexUid : indexUids) {
+
+            Index index = client.index(indexUid);
+
+            TestData<Movie> nestedTestData = this.getTestData(NESTED_MOVIES, Movie.class);
+            TaskInfo task1 = index.addDocuments(nestedTestData.getRaw());
+
+            index.waitForTask(task1.getTaskUid());
+
+            Settings settings = new Settings();
+            settings.setFilterableAttributes(new String[] {"id", "title"});
+            settings.setSortableAttributes(new String[] {"id"});
+
+            index.waitForTask(index.updateSettings(settings).getTaskUid());
+
+            TestData<Movie> moviesTestData = this.getTestData(MOVIES_INDEX, Movie.class);
+            TaskInfo task2 = index.addDocuments(moviesTestData.getRaw());
+
+            index.waitForTask(task2.getTaskUid());
+        }
+
+        MultiSearchRequest search = new MultiSearchRequest();
+
+        for (String indexUid : indexUids) {
+            search.addQuery(new IndexSearchRequest(indexUid).setQuery("Hobbit"));
+        }
+
+        MultiSearchFederation federation = new MultiSearchFederation();
+        federation.setLimit(20);
+        federation.setOffset(0);
+        Map<String, String[]> facetsByIndex = new HashMap<String, String[]>();
+        facetsByIndex.put("nestedMovies", new String[] {"title"});
+        facetsByIndex.put("movies", new String[] {"title", "id"});
+        federation.setFacetsByIndex(facetsByIndex);
+
+        MultiSearchResult results = client.multiSearch(search, federation);
+
+        assertThat(results.getHits().size(), is(4));
+
+        HashMap<String, FacetRating> facetStats = results.getFacetStats();
+        HashMap<String, HashMap<String, Integer>> facetDistribution =
+                results.getFacetDistribution();
+
+        HashMap<String, FacetsByIndexInfo> facetsByIndexInfo = results.getFacetsByIndex();
+
+        assertThat(facetDistribution, is(nullValue()));
+        assertThat(facetStats, is(nullValue()));
+        assertThat(facetsByIndexInfo, is(not(nullValue())));
+
+        for (String indexUid : indexUids) {
+            FacetsByIndexInfo indexInfo = facetsByIndexInfo.get(indexUid);
+            assertThat(indexInfo.getDistribution(), is(not(nullValue())));
+            assertThat(indexInfo.getStats(), is(not(nullValue())));
+        }
+
+        HashMap<String, HashMap<String, Integer>> moviesIndexDistribution =
+                facetsByIndexInfo.get("movies").getDistribution();
+
+        assertThat(moviesIndexDistribution.get("id"), is(not(nullValue())));
+        assertThat(moviesIndexDistribution.get("id").get("2"), is(equalTo(1)));
+        assertThat(moviesIndexDistribution.get("id").get("5"), is(equalTo(1)));
+        assertThat(moviesIndexDistribution.get("title"), is(not(nullValue())));
+        assertThat(moviesIndexDistribution.get("title").get("Hobbit"), is(equalTo(1)));
+        assertThat(moviesIndexDistribution.get("title").get("The Hobbit"), is(equalTo(1)));
+
+        HashMap<String, FacetRating> moviesFacetRating = facetsByIndexInfo.get("movies").getStats();
+        FacetRating idMoviesFacetRating = moviesFacetRating.get("id");
+
+        assertThat(idMoviesFacetRating, is(not(nullValue())));
+        assertThat(idMoviesFacetRating.getMin(), is(equalTo(2.0)));
+        assertThat(idMoviesFacetRating.getMax(), is(equalTo(5.0)));
+
+        HashMap<String, HashMap<String, Integer>> nestedMoviesIndexDistribution =
+                facetsByIndexInfo.get("nestedMovies").getDistribution();
+
+        assertThat(nestedMoviesIndexDistribution.get("title"), is(not(nullValue())));
+        assertThat(nestedMoviesIndexDistribution.get("title").get("Hobbit"), is(equalTo(1)));
+        assertThat(nestedMoviesIndexDistribution.get("title").get("The Hobbit"), is(equalTo(1)));
+
+        HashMap<String, FacetRating> nestedMoviesFacetRating =
+                facetsByIndexInfo.get("nestedMovies").getStats();
+        assertThat(nestedMoviesFacetRating.size(), is(equalTo((0))));
+    }
+
+    @Test
+    public void testMultiSearchWithMergeFacets() {
+        HashSet<String> indexUids = new HashSet<String>();
+        indexUids.add("movies");
+        indexUids.add("nestedMovies");
+
+        for (String indexUid : indexUids) {
+
+            Index index = client.index(indexUid);
+
+            TestData<Movie> nestedTestData = this.getTestData(NESTED_MOVIES, Movie.class);
+            TaskInfo task1 = index.addDocuments(nestedTestData.getRaw());
+
+            index.waitForTask(task1.getTaskUid());
+
+            Settings settings = new Settings();
+            settings.setFilterableAttributes(new String[] {"id", "title"});
+            settings.setSortableAttributes(new String[] {"id"});
+
+            index.waitForTask(index.updateSettings(settings).getTaskUid());
+
+            TestData<Movie> moviesTestData = this.getTestData(MOVIES_INDEX, Movie.class);
+            TaskInfo task2 = index.addDocuments(moviesTestData.getRaw());
+
+            index.waitForTask(task2.getTaskUid());
+        }
+
+        MultiSearchRequest search = new MultiSearchRequest();
+
+        for (String indexUid : indexUids) {
+            search.addQuery(new IndexSearchRequest(indexUid).setQuery("Hobbit"));
+        }
+
+        MultiSearchFederation federation = new MultiSearchFederation();
+        federation.setLimit(20);
+        federation.setOffset(0);
+        federation.setMergeFacets(new MergeFacets(10));
+        Map<String, String[]> facetsByIndex = new HashMap<String, String[]>();
+        facetsByIndex.put("nestedMovies", new String[] {"title"});
+        facetsByIndex.put("movies", new String[] {"title", "id"});
+        federation.setFacetsByIndex(facetsByIndex);
+
+        MultiSearchResult results = client.multiSearch(search, federation);
+
+        assertThat(results.getHits().size(), is(4));
+
+        HashMap<String, FacetRating> facetStats = results.getFacetStats();
+        HashMap<String, HashMap<String, Integer>> facetDistribution =
+                results.getFacetDistribution();
+
+        assertThat(facetDistribution, is(not(nullValue())));
+        assertThat(facetStats, is(not(nullValue())));
+        assertThat(results.getFacetsByIndex(), is(nullValue()));
+
+        FacetRating idFacet = facetStats.get("id");
+
+        assertThat(idFacet.getMin(), is(equalTo(2.0)));
+        assertThat(idFacet.getMax(), is(equalTo(5.0)));
+
+        assertThat(facetDistribution.get("id").get("2"), is(equalTo(1)));
+        assertThat(facetDistribution.get("id").get("5"), is(equalTo(1)));
+
+        assertThat(facetDistribution.get("title").get("Hobbit"), is(equalTo(2)));
+        assertThat(facetDistribution.get("title").get("The Hobbit"), is(equalTo(2)));
+    }
+
     @Test
     public void testSimilarDocuments() throws Exception {
         HashMap<String, Boolean> features = new HashMap();

src/test/java/com/meilisearch/integration/TenantTokenTest.java

@@ -108,7 +108,7 @@ public void testGenerateTenantTokenWithFilter() throws Exception {
 
         assertThat(searchResult.getHits().size(), is(equalTo(20)));
         assertThat(searchResult.getLimit(), is(equalTo(20)));
-        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(30)));
+        assertThat(searchResult.getEstimatedTotalHits(), is(equalTo(31)));
     }
 
     /** Test Create Tenant Token with expiration date */

src/test/resources/movies.json

@@ -396,5 +396,17 @@
 		  "Thriller",
 		  "Drama"
 	   ]
+	},
+	{
+		"id": 2,
+		"title": "Hobbit",
+		"poster": "https://www.imdb.com/title/tt0903624/mediaviewer/rm3577719808/?ref_=tt_ov_i",
+		"overview": "A reluctant Hobbit, Bilbo Baggins, sets out to the Lonely Mountain with a spirited group of dwarves to reclaim their mountain home and the gold within it from the dragon Smaug.",
+		"release_date": "2021-12-14",
+		"language": "en",
+		"genres": [
+			"Adventure",
+			"Fantasy"
+		]
 	}
 ]