feat(graphql): implement sort and facet for scroll (datahub-project#1…

…2746)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/resolver/GetMetadataAnalyticsResolver.java

@@ -78,7 +78,7 @@ private List<AnalyticsChart> getCharts(MetadataAnalyticsInput input, OperationCo
 
     SearchResult searchResult =
         _entityClient.searchAcrossEntities(
-            opContext, entities, query, filter, 0, 0, Collections.emptyList(), null);
+            opContext, entities, query, filter, 0, 0, Collections.emptyList());
 
     List<AggregationMetadata> aggregationMetadataList =
         searchResult.getMetadata().getAggregations();

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolver.java

@@ -91,8 +91,7 @@ public CompletableFuture<SearchResults> get(final DataFetchingEnvironment enviro
                                         new CriterionArray(ImmutableList.of(filterCriterion))))),
                     start,
                     count,
-                    Collections.emptyList(),
-                    null));
+                    Collections.emptyList()));
 
           } catch (Exception e) {
             throw new RuntimeException(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java

@@ -185,7 +185,6 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
                         finalFilter,
                         start,
                         count,
-                        null,
                         null));
             results
                 .getSearchResults()

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java

@@ -95,8 +95,7 @@ public CompletableFuture<SearchResults> get(final DataFetchingEnvironment enviro
                                 new ConjunctiveCriterion().setAnd(criteria))),
                     start,
                     count,
-                    Collections.emptyList(),
-                    null));
+                    Collections.emptyList()));
 
           } catch (Exception e) {
             throw new RuntimeException(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolver.java

@@ -108,8 +108,7 @@ private SearchResult getSearchResults(
             : null,
         0,
         0,
-        Collections.emptyList(),
-        null);
+        Collections.emptyList());
   }
 
   /**

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossEntitiesResolver.java

@@ -16,6 +16,7 @@
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.service.ViewService;
 import com.linkedin.view.DataHubViewInfo;
 import graphql.schema.DataFetcher;
@@ -80,6 +81,7 @@ public CompletableFuture<ScrollResults> get(DataFetchingEnvironment environment)
           } else {
             searchFlags = null;
           }
+          List<SortCriterion> sortCriteria = SearchUtils.getSortCriteria(input.getSortInput());
 
           try {
             log.debug(
@@ -108,6 +110,7 @@ public CompletableFuture<ScrollResults> get(DataFetchingEnvironment environment)
                         : baseFilter,
                     scrollId,
                     keepAlive,
+                    sortCriteria,
                     count));
           } catch (Exception e) {
             log.error(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossEntitiesResolver.java

@@ -146,8 +146,7 @@ private List<String> getStructuredPropertyFacets(final QueryContext context) {
               createStructuredPropertyFilter(),
               0,
               100,
-              Collections.emptyList(),
-              null);
+              Collections.emptyList());
       return result.getEntities().stream()
           .map(entity -> String.format("structuredProperties.%s", entity.getEntity().getId()))
           .collect(Collectors.toList());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java

@@ -378,6 +378,8 @@ public static CompletableFuture<ScrollResults> scrollAcrossEntities(
       Integer inputCount,
       String scrollId,
       String inputKeepAlive,
+      List<SortCriterion> sortCriteria,
+      List<String> facets,
       String className) {
 
     final List<EntityType> entityTypes =
@@ -431,7 +433,15 @@ public static CompletableFuture<ScrollResults> scrollAcrossEntities(
           try {
             final ScrollResult scrollResult =
                 _entityClient.scrollAcrossEntities(
-                    context, finalEntityNames, query, finalFilters, scrollId, keepAlive, count);
+                    context,
+                    finalEntityNames,
+                    query,
+                    finalFilters,
+                    scrollId,
+                    keepAlive,
+                    sortCriteria,
+                    count,
+                    facets);
             return UrnScrollResultsMapper.map(inputContext, scrollResult);
           } catch (Exception e) {
             log.warn(
@@ -518,14 +528,7 @@ public static CompletableFuture<SearchResults> searchAcrossEntities(
           try {
             final SearchResult searchResult =
                 _entityClient.searchAcrossEntities(
-                    context,
-                    finalEntityNames,
-                    query,
-                    finalFilters,
-                    start,
-                    count,
-                    sortCriteria,
-                    null);
+                    context, finalEntityNames, query, finalFilters, start, count, sortCriteria);
             return UrnSearchResultsMapper.map(inputContext, searchResult);
           } catch (Exception e) {
             log.warn(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/siblings/SiblingsSearchResolver.java

@@ -19,6 +19,7 @@
 import com.linkedin.metadata.utils.CriterionUtils;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -61,6 +62,8 @@ public CompletableFuture<ScrollResults> get(DataFetchingEnvironment environment)
         input.getCount(),
         input.getScrollId(),
         input.getKeepAlive(),
+        List.of(),
+        List.of(),
         this.getClass().getSimpleName());
   }
 }

datahub-graphql-core/src/main/resources/search.graphql

@@ -347,6 +347,11 @@ input ScrollAcrossEntitiesInput {
   Flags controlling search options
   """
   searchFlags: SearchFlags
+
+  """
+  Optional - Information on how to sort this search result
+  """
+  sortInput: SearchSortInput
 }
 
 

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/container/ContainerEntitiesResolverTest.java

@@ -60,8 +60,7 @@ public void testGetSuccess() throws Exception {
                                         new CriterionArray(ImmutableList.of(filterCriterion)))))),
                 Mockito.eq(0),
                 Mockito.eq(20),
-                Mockito.eq(Collections.emptyList()),
-                Mockito.eq(null)))
+                Mockito.eq(Collections.emptyList())))
         .thenReturn(
             new SearchResult()
                 .setFrom(0)

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolverTest.java

@@ -65,8 +65,7 @@ public void testGetSuccess() throws Exception {
                                         new CriterionArray(ImmutableList.of(filterCriterion)))))),
                 Mockito.eq(0),
                 Mockito.eq(20),
-                Mockito.eq(Collections.emptyList()),
-                Mockito.eq(null)))
+                Mockito.eq(Collections.emptyList())))
         .thenReturn(
             new SearchResult()
                 .setFrom(0)

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/GetQuickFiltersResolverTest.java

@@ -301,8 +301,7 @@ private static EntityClient initMockEntityClient(
                 Mockito.eq(filter),
                 Mockito.eq(start),
                 Mockito.eq(limit),
-                Mockito.eq(Collections.emptyList()),
-                Mockito.eq(null)))
+                Mockito.eq(Collections.emptyList())))
         .thenReturn(result);
     return client;
   }

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/versioning/VersionsSearchResolverTest.java

@@ -4,6 +4,7 @@
 import static com.linkedin.metadata.Constants.*;
 import static com.linkedin.metadata.utils.CriterionUtils.*;
 import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertThrows;
 
@@ -39,6 +40,7 @@
 import com.linkedin.view.DataHubViewInfo;
 import com.linkedin.view.DataHubViewType;
 import graphql.schema.DataFetchingEnvironment;
+import io.datahubproject.metadata.context.OperationContext;
 import java.util.List;
 import java.util.concurrent.CompletionException;
 import org.mockito.Mockito;
@@ -144,8 +146,7 @@ public void testGetSuccessBasic() throws Exception {
                 List.of(
                     new com.linkedin.metadata.query.filter.SortCriterion()
                         .setField(VERSION_SORT_ID_FIELD_NAME)
-                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))),
-            any());
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))));
   }
 
   @Test
@@ -240,8 +241,7 @@ public void testGetSuccessComplex() throws Exception {
                         .setOrder(com.linkedin.metadata.query.filter.SortOrder.ASCENDING),
                     new com.linkedin.metadata.query.filter.SortCriterion()
                         .setField(VERSION_SORT_ID_FIELD_NAME)
-                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))),
-            any());
+                        .setOrder(com.linkedin.metadata.query.filter.SortOrder.DESCENDING))));
   }
 
   @Test
@@ -251,7 +251,7 @@ public void testThrowsError() throws Exception {
 
     Mockito.when(
             mockEntityClient.searchAcrossEntities(
-                any(), any(), any(), any(), Mockito.anyInt(), Mockito.anyInt(), any(), any()))
+                any(), any(), any(), any(), Mockito.anyInt(), Mockito.anyInt(), any()))
         .thenThrow(new RemoteInvocationException());
 
     VersionsSearchResolver resolver = new VersionsSearchResolver(mockEntityClient, mockViewService);
@@ -271,23 +271,36 @@ public void testThrowsError() throws Exception {
   private EntityClient initMockEntityClient() throws Exception {
     EntityClient client = Mockito.mock(EntityClient.class);
 
+    SearchResult result =
+        new SearchResult()
+            .setEntities(new SearchEntityArray())
+            .setNumEntities(0)
+            .setFrom(0)
+            .setPageSize(0)
+            .setMetadata(new SearchResultMetadata());
+
     Mockito.when(
             client.searchAcrossEntities(
-                any(),
+                any(OperationContext.class),
                 any(),
                 Mockito.anyString(),
-                any(),
+                any(Filter.class),
                 Mockito.anyInt(),
                 Mockito.anyInt(),
+                anyList(),
+                anyList()))
+        .thenReturn(result);
+
+    Mockito.when(
+            client.searchAcrossEntities(
+                any(OperationContext.class),
                 any(),
-                Mockito.eq(null)))
-        .thenReturn(
-            new SearchResult()
-                .setEntities(new SearchEntityArray())
-                .setNumEntities(0)
-                .setFrom(0)
-                .setPageSize(0)
-                .setMetadata(new SearchResultMetadata()));
+                anyString(),
+                any(Filter.class),
+                anyInt(),
+                anyInt(),
+                anyList()))
+        .thenReturn(result);
 
     return client;
   }

docs/api/graphql/graphql-best-practices.md

@@ -33,7 +33,14 @@ This technique makes maintaining your GraphQL queries much more doable. For exam
 
 `search*` APIs such as [`searchAcrossEntities`](https://datahubproject.io/docs/GraphQL/queries/#searchacrossentities) are designed for minimal pagination (< ~50). They do not perform well for deep pagination requests. Use the equivalent `scroll*` APIs such as [`scrollAcrossEntities`](https://datahubproject.io/docs/GraphQL/queries/#scrollacrossentities) when expecting the need to paginate deeply into the result set.
 
-Note: that it is impossible to use `search*` for paginating beyond 10k results.
+:::note
+It is impossible to use `search*` for paginating beyond 10k results.
+:::
+
+:::caution
+In order to `scroll*` through the entire result set it is required to use a stable sort order. This means using `_score` as
+the first sort order cannot be used. Use the `urn` field as the sort order instead.
+:::
 
 #### Examples
 
@@ -54,7 +61,15 @@ Page 1 Request:
       orFilters: [
         { and: [{ field: "name", condition: CONTAIN, values: ["pet"] }] },
         { and: [{ field: "title", condition: CONTAIN, values: ["pet"] }] }
-      ]
+      ],
+      sortInput: {
+        sortCriteria: [
+          {
+            field: "urn",
+            sortOrder: ASCENDING
+          }
+        ]
+      }
     }
   ) {
     nextScrollId
@@ -110,7 +125,15 @@ Page 2 Request:
       orFilters: [
         { and: [{ field: "name", condition: CONTAIN, values: ["pet"] }] },
         { and: [{ field: "title", condition: CONTAIN, values: ["pet"] }] }
-      ]
+      ],
+      sortInput: {
+        sortCriteria: [
+          {
+            field: "urn",
+            sortOrder: ASCENDING
+          }
+        ]
+      }
     }
   ) {
     nextScrollId
@@ -282,7 +305,20 @@ Example for skipping highlighting and aggregates, typically used for scrolling s
 ```graphql
 {
   scrollAcrossEntities(
-    input: {types: [DATASET], count: 2, query: "pet", searchFlags: {skipAggregates: true, skipHighlighting: true}}
+    input: {
+      types: [DATASET], 
+      count: 2, 
+      query: "pet", 
+      searchFlags: {skipAggregates: true, skipHighlighting: true},
+      sortInput: {
+        sortCriteria: [
+          {
+            field: "urn",
+            sortOrder: ASCENDING
+          }
+        ]
+      },
+    }
   ) {
     searchResults {
       entity {