Optimizing get indices to segments map calculations for IndicesSegmen…

server/src/main/java/org/opensearch/action/admin/indices/segments/IndicesSegmentResponse.java

@@ -45,13 +45,12 @@
 import org.opensearch.index.engine.Segment;
 
 import java.io.IOException;
-import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * Transport response for retrieving indices segment information
@@ -86,21 +85,24 @@ public Map<String, IndexSegments> getIndices() {
             return indicesSegments;
         }
         Map<String, IndexSegments> indicesSegments = new HashMap<>();
-
-        Set<String> indices = new HashSet<>();
-        for (ShardSegments shard : shards) {
-            indices.add(shard.getShardRouting().getIndexName());
+        if (shards.length == 0) {
+            this.indicesSegments = indicesSegments;
+            return indicesSegments;
         }
 
-        for (String indexName : indices) {
-            List<ShardSegments> shards = new ArrayList<>();
-            for (ShardSegments shard : this.shards) {
-                if (shard.getShardRouting().getIndexName().equals(indexName)) {
-                    shards.add(shard);
-                }
+        Arrays.sort(shards, Comparator.comparing(shardSegment -> shardSegment.getShardRouting().getIndexName()));
+        int startIndexPos = 0;
+        String startIndexName = shards[startIndexPos].getShardRouting().getIndexName();
+        for (int i = 0; i < shards.length; i++) {
+            if (!shards[i].getShardRouting().getIndexName().equals(startIndexName)) {
+                indicesSegments.put(startIndexName, new IndexSegments(startIndexName, Arrays.copyOfRange(shards, startIndexPos, i)));
+                startIndexPos = i;
+                startIndexName = shards[startIndexPos].getShardRouting().getIndexName();
             }
-            indicesSegments.put(indexName, new IndexSegments(indexName, shards.toArray(new ShardSegments[0])));
         }
+        // Add the last shardSegment from shards list which would have got missed in the loop above
+        indicesSegments.put(startIndexName, new IndexSegments(startIndexName, Arrays.copyOfRange(shards, startIndexPos, shards.length)));
+
         this.indicesSegments = indicesSegments;
         return indicesSegments;
     }

server/src/test/java/org/opensearch/action/admin/indices/segments/IndicesSegmentResponseTests.java

@@ -42,7 +42,11 @@
 import org.opensearch.index.engine.Segment;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 
@@ -68,4 +72,54 @@ public void testToXContentSerialiationWithSortedFields() throws Exception {
             response.toXContent(builder, ToXContent.EMPTY_PARAMS);
         }
     }
+
+    public void testGetIndices() {
+        final int totalIndices = 5;
+        final int shardsPerIndex = 3;
+        final int segmentsPerShard = 2;
+        // Preparing a ShardSegments list, which will have (totalIndices * shardsPerIndex) shardSegments.
+        // Indices will be named -> foo1, foo2, ..., foo{totalIndices}
+        List<ShardSegments> shardSegmentsList = new ArrayList<>();
+        for (int indexName = 0; indexName < totalIndices; indexName++) {
+            for (int shardId = 0; shardId < shardsPerIndex; shardId++) {
+                ShardRouting shardRouting = TestShardRouting.newShardRouting(
+                    "foo" + indexName,
+                    shardId,
+                    "node_id",
+                    true,
+                    ShardRoutingState.STARTED
+                );
+                List<Segment> segmentList = new ArrayList<>();
+                for (int segmentNum = 0; segmentNum < segmentsPerShard; segmentNum++) {
+                    segmentList.add(new Segment("foo" + indexName + shardId + segmentNum));
+                }
+                shardSegmentsList.add(new ShardSegments(shardRouting, segmentList));
+            }
+        }
+        Collections.shuffle(shardSegmentsList);
+
+        // Prepare the IndicesSegmentResponse object and get the indicesSegments map
+        IndicesSegmentResponse response = new IndicesSegmentResponse(
+            shardSegmentsList.toArray(new ShardSegments[0]),
+            totalIndices * shardsPerIndex,
+            totalIndices * shardsPerIndex,
+            0,
+            Collections.emptyList()
+        );
+        Map<String, IndexSegments> indicesSegments = response.getIndices();
+
+        assertEquals(totalIndices, indicesSegments.size());
+        for (Map.Entry<String, IndexSegments> indexSegmentEntry : indicesSegments.entrySet()) {
+            assertEquals(shardsPerIndex, indexSegmentEntry.getValue().getShards().size());
+            for (IndexShardSegments indexShardSegment : indexSegmentEntry.getValue().getShards().values()) {
+                for (ShardSegments shardSegment : indexShardSegment.getShards()) {
+                    assertEquals(segmentsPerShard, shardSegment.getSegments().size());
+                    for (int i = 0; i < segmentsPerShard; i++) {
+                        String segmentName = indexSegmentEntry.getKey() + shardSegment.getShardRouting().getId() + i;
+                        assertEquals(segmentName, shardSegment.getSegments().get(i).getName());
+                    }
+                }
+            }
+        }
+    }
 }