fix(sort): fix out of bounds exception with count 0 (#15085)

Author: david-leifker

metadata-io/src/main/java/com/linkedin/metadata/dataHubUsage/DataHubUsageServiceImpl.java

@@ -112,7 +112,8 @@ private ExternalAuditEventsSearchResponse mapExternalAuditEventsSearchResponse(
     response.count(searchHits.getHits().length);
     response.total((int) searchHits.getTotalHits().value);
     String nextScrollId = null;
-    if (searchHits.getHits().length == analyticsSearchRequest.getSize()) {
+    if (searchHits.getHits().length == analyticsSearchRequest.getSize()
+        && searchHits.getHits().length > 0) {
       Object[] sort = searchHits.getHits()[searchHits.getHits().length - 1].getSortValues();
       nextScrollId = new SearchAfterWrapper(sort, null, 0L).toScrollId();
     }

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java

@@ -364,7 +364,7 @@ public RelatedEntitiesScrollResult scrollRelatedEntities(
     SearchHit[] searchHits = response.getHits().getHits();
     // Only return next scroll ID if there are more results, indicated by full size results
     String nextScrollId = null;
-    if (searchHits.length == count) {
+    if (searchHits.length == count && searchHits.length > 0) {
       Object[] sort = searchHits[searchHits.length - 1].getSortValues();
       nextScrollId = new SearchAfterWrapper(sort, null, 0L).toScrollId();
     }

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java

@@ -452,7 +452,7 @@ public ScrollResult extractScrollResult(
     SearchHit[] searchHits = searchResponse.getHits().getHits();
     // Only return next scroll ID if there are more results, indicated by full size results
     String nextScrollId = null;
-    if (searchHits.length == size) {
+    if (searchHits.length == size && searchHits.length > 0) {
       Object[] sort = searchHits[searchHits.length - 1].getSortValues();
       long expirationTimeMs = 0L;
       if (keepAlive != null && supportsPointInTime) {

metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphServiceTest.java

@@ -1,6 +1,7 @@
 package com.linkedin.metadata.graph.elastic;
 
 import static io.datahubproject.test.search.SearchTestUtils.TEST_GRAPH_SERVICE_CONFIG;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -254,6 +255,52 @@ public void testFindRelatedEntitiesNoResultsByType() {
     assertTrue(result.getEntities().isEmpty());
   }
 
+  @Test
+  public void testScrollRelatedEntitiesWithZeroCount() {
+    // Test the edge case where count=0 is passed, which should not cause
+    // ArrayIndexOutOfBoundsException when accessing searchHits[searchHits.length - 1]
+    OperationContext mockOpContext = TestOperationContexts.systemContextNoValidate();
+    GraphFilters mockGraphFilters = GraphFilters.ALL;
+    List<SortCriterion> mockSortCriteria = Collections.emptyList();
+
+    String scrollId = "test-scroll-id";
+    String keepAlive = "1m";
+    int count = 0; // This is the edge case we're testing
+
+    // Create mock search response with empty hits
+    SearchResponse mockResponse = mock(SearchResponse.class);
+    SearchHits mockHits = mock(SearchHits.class);
+    when(mockResponse.getHits()).thenReturn(mockHits);
+
+    // Create empty search hits array (simulating no results)
+    SearchHit[] hits = new SearchHit[0];
+    when(mockHits.getHits()).thenReturn(hits);
+    when(mockHits.getTotalHits()).thenReturn(new TotalHits(0L, TotalHits.Relation.EQUAL_TO));
+
+    when(mockReadDAO.getSearchResponse(any(), any(), any(), any(), any(), anyInt()))
+        .thenReturn(mockResponse);
+    when(mockReadDAO.getGraphServiceConfig()).thenReturn(TEST_GRAPH_SERVICE_CONFIG);
+
+    // This should not throw ArrayIndexOutOfBoundsException
+    RelatedEntitiesScrollResult result =
+        test.scrollRelatedEntities(
+            mockOpContext,
+            mockGraphFilters,
+            mockSortCriteria,
+            scrollId,
+            keepAlive,
+            count,
+            null,
+            null);
+
+    // Verify the result
+    assertNotNull(result);
+    assertEquals(result.getNumResults(), 0);
+    assertEquals(result.getPageSize(), 0);
+    assertTrue(result.getEntities().isEmpty());
+    assertNull(result.getScrollId()); // No scroll ID since no results
+  }
+
   @Test
   public void testRaw() {
     // Create test edge tuples

metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java

@@ -1383,6 +1383,43 @@ public void testExtractScrollResultPaginationLogic() {
     assertFalse(result.hasScrollId());
   }
 
+  @Test
+  public void testExtractScrollResultWithZeroCount() {
+    // Test the edge case where count=0 is passed, which should not cause
+    // ArrayIndexOutOfBoundsException
+    SearchRequestHandler handler =
+        SearchRequestHandler.getBuilder(
+            operationContext,
+            TestEntitySpecBuilder.getSpec(),
+            testQueryConfig,
+            null,
+            QueryFilterRewriteChain.EMPTY,
+            TEST_SEARCH_SERVICE_CONFIG);
+
+    SearchResponse mockResponse = mock(SearchResponse.class);
+    SearchHits mockHits = mock(SearchHits.class);
+
+    // Create empty search hits array (simulating no results)
+    SearchHit[] hits = new SearchHit[0];
+
+    when(mockResponse.getHits()).thenReturn(mockHits);
+    when(mockHits.getTotalHits()).thenReturn(new TotalHits(0L, TotalHits.Relation.EQUAL_TO));
+    when(mockHits.getHits()).thenReturn(hits);
+    when(mockResponse.getAggregations()).thenReturn(null);
+    when(mockResponse.getSuggest()).thenReturn(null);
+    when(mockResponse.pointInTimeId()).thenReturn("test-pit-id");
+
+    // This should not throw ArrayIndexOutOfBoundsException
+    ScrollResult result =
+        handler.extractScrollResult(operationContext, mockResponse, null, "5m", 0, true);
+
+    // Verify the result
+    assertNotNull(result);
+    assertEquals(result.getPageSize().intValue(), 0);
+    assertEquals(result.getNumEntities().intValue(), 0);
+    assertFalse(result.hasScrollId()); // No scroll ID since no results
+  }
+
   // Helper method to create scroll results with specific sizes
   private ScrollResult verifyScrollResultSize(
       SearchRequestHandler handler,