feat: track total hits in ES7 if configured (#124)

* feat: track total hits in ES7 if configured

* address comments

* no need to set the const

* remove query template

dao-impl/elasticsearch-dao-7/src/main/java/com/linkedin/metadata/dao/browse/ESBrowseDAO.java

@@ -48,12 +48,21 @@
 public class ESBrowseDAO extends BaseBrowseDAO {
   private final RestHighLevelClient _client;
   private final BaseBrowseConfig _config;
+  private int _lowerBoundHits = Integer.MAX_VALUE;
 
   public ESBrowseDAO(@Nonnull RestHighLevelClient esClient, @Nonnull BaseBrowseConfig config) {
     this._client = esClient;
     this._config = config;
   }
 
+  /**
+   * Set "track_total_hits" query parameter to a custom lower bound if you do not need accurate results. It is a good
+   * trade off to speed up searches if you don’t need the accurate number of hits after a certain threshold.
+   */
+  public void setTrackTotalHits(int lowermost) {
+    _lowerBoundHits = lowermost;
+  }
+
   /**
    * Gets a list of groups/entities that match given browse request.
    *
@@ -110,6 +119,8 @@ private AggregationBuilder buildAggregations(@Nonnull String path) {
   protected SearchRequest constructGroupsSearchRequest(@Nonnull String path, @Nonnull Map<String, String> requestMap) {
     final SearchRequest searchRequest = new SearchRequest(_config.getIndexName());
     final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+    searchSourceBuilder.trackTotalHitsUpTo(_lowerBoundHits);
+
     searchSourceBuilder.query(buildQueryString(path, requestMap, true));
     searchSourceBuilder.aggregation(buildAggregations(path));
     searchRequest.source(searchSourceBuilder);
@@ -171,6 +182,9 @@ SearchRequest constructEntitiesSearchRequest(@Nonnull String path, @Nonnull Map<
     final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
     searchSourceBuilder.from(from);
     searchSourceBuilder.size(size);
+
+    searchSourceBuilder.trackTotalHitsUpTo(_lowerBoundHits);
+
     searchSourceBuilder.fetchSource(new String[]{_config.getBrowsePathFieldName(), _config.getUrnFieldName()}, null);
     searchSourceBuilder.sort(_config.getSortingField(), SortOrder.ASC);
     searchSourceBuilder.query(buildQueryString(path, requestMap, false));

dao-impl/elasticsearch-dao-7/src/main/java/com/linkedin/metadata/dao/search/ESSearchDAO.java

@@ -74,6 +74,7 @@ public class ESSearchDAO<DOCUMENT extends RecordTemplate> extends BaseSearchDAO<
   private BaseESAutoCompleteQuery _autoCompleteQueryForLowCardFields;
   private BaseESAutoCompleteQuery _autoCompleteQueryForHighCardFields;
   private int _maxTermBucketSize = DEFAULT_TERM_BUCKETS_SIZE_100;
+  private int _lowerBoundHits = Integer.MAX_VALUE;
 
   // Regex patterns for matching original field names to the highlighted field name returned by elasticsearch
   private Map<String, Pattern> _highlightedFieldNamePatterns;
@@ -94,6 +95,14 @@ public ESSearchDAO(@Nonnull RestHighLevelClient esClient, @Nonnull Class<DOCUMEN
         .collect(Collectors.toMap(Function.identity(), fieldName -> Pattern.compile(fieldName + "(\\..+)?")));
   }
 
+  /**
+   * Set "track_total_hits" query parameter to a custom lower bound if you do not need accurate results. It is a good
+   * trade off to speed up searches if you don’t need the accurate number of hits after a certain threshold.
+   */
+  public void setTrackTotalHits(int lowermost) {
+    _lowerBoundHits = lowermost;
+  }
+
   @Nonnull
   protected BaseESAutoCompleteQuery getAutocompleteQueryGenerator(@Nonnull String field) {
     if (_config.getLowCardinalityFields() != null && _config.getLowCardinalityFields().contains(field)) {
@@ -200,6 +209,7 @@ SearchRequest getFilteredSearchQuery(@Nullable Filter filters, @Nullable SortCri
     }
     final SearchRequest searchRequest = new SearchRequest(_config.getIndexName());
     final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+    searchSourceBuilder.trackTotalHitsUpTo(_lowerBoundHits);
     searchSourceBuilder.query(boolQueryBuilder);
     searchSourceBuilder.from(from).size(size);
     ESUtils.buildSortOrder(searchSourceBuilder, sortCriterion);
@@ -252,6 +262,8 @@ SearchRequest constructSearchQuery(@Nonnull String input, @Nullable Filter filte
     searchSourceBuilder.from(from);
     searchSourceBuilder.size(size);
 
+    searchSourceBuilder.trackTotalHitsUpTo(_lowerBoundHits);
+
     searchSourceBuilder.query(buildQueryString(input));
     searchSourceBuilder.postFilter(ESUtils.buildFilterQuery(filter));
     buildAggregations(searchSourceBuilder, filter);

dao-impl/elasticsearch-dao-7/src/test/java/com/linkedin/metadata/dao/search/ESSearchDAOTest.java

@@ -215,6 +215,22 @@ public void testFilteredQueryWithTermsFilter() throws IOException {
     assertEquals(searchRequest.indices(), new String[]{_testSearchConfig.getIndexName()});
   }
 
+  @Test
+  public void testFilteredQueryApproxHits() throws IOException {
+    int from = 0;
+    int size = 10;
+    _searchDAO.setTrackTotalHits(100);
+    Filter filter = newFilter(ImmutableMap.of("key1", "value1, value2 ", "key2", "urn:li:entity:(foo,bar,baz)"));
+    SortCriterion sortCriterion = new SortCriterion().setOrder(SortOrder.ASCENDING).setField("urn");
+
+    SearchRequest searchRequest = _searchDAO.getFilteredSearchQuery(filter, sortCriterion, from, size);
+    assertEquals(searchRequest.source().toString(), loadJsonFromResource("UrnFilterApproxHitsQuery.json"));
+    assertEquals(searchRequest.indices(), new String[]{_testSearchConfig.getIndexName()});
+
+    // set the value of "track_total_hits" back to the default i.e. maximum integer value
+    _searchDAO.setTrackTotalHits(Integer.MAX_VALUE);
+  }
+
   @Test
   public void testFilteredQueryWithUrnValue() throws IOException {
     int from = 0;

dao-impl/elasticsearch-dao-7/src/test/resources/DefaultSortTermsFilterQuery.json

@@ -37,5 +37,6 @@
         "order" : "asc"
       }
     }
-  ]
+  ],
+  "track_total_hits" : 2147483647
 }

dao-impl/elasticsearch-dao-7/src/test/resources/EmptyFilterQuery.json

@@ -13,5 +13,6 @@
         "order" : "asc"
       }
     }
-  ]
+  ],
+  "track_total_hits" : 2147483647
 }

dao-impl/elasticsearch-dao-7/src/test/resources/RangeFilterQuery.json

@@ -59,5 +59,6 @@
         "order" : "asc"
       }
     }
-  ]
+  ],
+  "track_total_hits" : 2147483647
 }

dao-impl/elasticsearch-dao-7/src/test/resources/SortByUrnTermsFilterQuery.json

@@ -32,5 +32,6 @@
         "order" : "asc"
       }
     }
-  ]
+  ],
+  "track_total_hits" : 2147483647
 }

dao-impl/elasticsearch-dao-7/src/test/resources/UrnFilterApproxHitsQuery.json

@@ -0,0 +1,37 @@
+{
+  "from" : 0,
+  "size" : 10,
+  "query" : {
+    "bool" : {
+      "filter" : [
+        {
+          "terms" : {
+            "key1" : [
+              "value1",
+              "value2"
+            ],
+            "boost" : 1.0
+          }
+        },
+        {
+          "terms" : {
+            "key2" : [
+              "urn:li:entity:(foo,bar,baz)"
+            ],
+            "boost" : 1.0
+          }
+        }
+      ],
+      "adjust_pure_negative" : true,
+      "boost" : 1.0
+    }
+  },
+  "sort" : [
+    {
+      "urn" : {
+        "order" : "asc"
+      }
+    }
+  ],
+  "track_total_hits" : 100
+}

dao-impl/elasticsearch-dao-7/src/test/resources/UrnFilterQuery.json

@@ -32,5 +32,6 @@
         "order" : "asc"
       }
     }
-  ]
+  ],
+  "track_total_hits" : 2147483647
 }