Use only first index and optimization to get properties once per query

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>

src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java

@@ -112,32 +112,58 @@ public String buildShape(
         Boolean showFieldType,
         Set<Index> successfulSearchShardIndices
     ) {
+        Index firstIndex = successfulSearchShardIndices.iterator().next();
+        Map<String, Object> propertiesAsMap = getPropertiesMapForIndex(firstIndex);
+
         StringBuilder shape = new StringBuilder();
-        shape.append(buildQueryShape(source.query(), showFieldName, showFieldType, successfulSearchShardIndices));
-        shape.append(buildAggregationShape(source.aggregations(), showFieldName, showFieldType, successfulSearchShardIndices));
-        shape.append(buildSortShape(source.sorts(), showFieldName, showFieldType, successfulSearchShardIndices));
+        shape.append(buildQueryShape(source.query(), showFieldName, showFieldType, propertiesAsMap, firstIndex));
+        shape.append(buildAggregationShape(source.aggregations(), showFieldName, showFieldType, propertiesAsMap, firstIndex));
+        shape.append(buildSortShape(source.sorts(), showFieldName, showFieldType, propertiesAsMap, firstIndex));
         return shape.toString();
     }
 
+    private Map<String, Object> getPropertiesMapForIndex(Index index) {
+        Map<String, MappingMetadata> indexMapping;
+        try {
+            indexMapping = clusterService.state().metadata().findMappings(new String[] { index.getName() }, input -> str -> true);
+        } catch (IOException e) {
+            // If an error occurs while retrieving mappings, return an empty map
+            return Collections.emptyMap();
+        }
+
+        MappingMetadata mappingMetadata = indexMapping.get(index.getName());
+        if (mappingMetadata == null) {
+            return Collections.emptyMap();
+        }
+
+        Map<String, Object> propertiesMap = (Map<String, Object>) mappingMetadata.getSourceAsMap().get("properties");
+        if (propertiesMap == null) {
+            return Collections.emptyMap();
+        }
+        return propertiesMap;
+    }
+
     /**
      * Builds the query-section shape.
      *
      * @param queryBuilder          search request query builder
      * @param showFieldName         whether to append field names
      * @param showFieldType         whether to append field types
-     * @param successfulSearchShardIndices the set of indices that were successfully searched
+     * @param propertiesAsMap       properties
+     * @param index                 index
      * @return Query-section shape as a String
      */
     String buildQueryShape(
         QueryBuilder queryBuilder,
         Boolean showFieldName,
         Boolean showFieldType,
-        Set<Index> successfulSearchShardIndices
+        Map<String, Object> propertiesAsMap,
+        Index index
     ) {
         if (queryBuilder == null) {
             return EMPTY_STRING;
         }
-        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(this, successfulSearchShardIndices, showFieldName, showFieldType);
+        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(this, propertiesAsMap, index, showFieldName, showFieldType);
         queryBuilder.visit(shapeVisitor);
         return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFieldName, showFieldType);
     }
@@ -148,14 +174,16 @@ String buildQueryShape(
      * @param aggregationsBuilder    search request aggregation builder
      * @param showFieldName          whether to append field names
      * @param showFieldType          whether to append field types
-     * @param successfulSearchShardIndices the set of indices that were successfully searched
+     * @param propertiesAsMap       properties
+     * @param index                 index
      * @return Aggregation shape as a String
      */
     String buildAggregationShape(
         AggregatorFactories.Builder aggregationsBuilder,
         Boolean showFieldName,
         Boolean showFieldType,
-        Set<Index> successfulSearchShardIndices
+        Map<String, Object> propertiesAsMap,
+        Index index
     ) {
         if (aggregationsBuilder == null) {
             return EMPTY_STRING;
@@ -167,7 +195,8 @@ String buildAggregationShape(
             new StringBuilder(),
             showFieldName,
             showFieldType,
-            successfulSearchShardIndices
+            propertiesAsMap,
+            index
         );
         return aggregationShape.toString();
     }
@@ -179,7 +208,8 @@ StringBuilder recursiveAggregationShapeBuilder(
         StringBuilder baseIndent,
         Boolean showFieldName,
         Boolean showFieldType,
-        Set<Index> successfulSearchShardIndices
+        Map<String, Object> propertiesAsMap,
+        Index index
     ) {
         //// Normal Aggregations ////
         if (aggregationBuilders.isEmpty() == false) {
@@ -190,7 +220,7 @@ StringBuilder recursiveAggregationShapeBuilder(
             StringBuilder stringBuilder = new StringBuilder();
             stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType());
             if (showFieldName || showFieldType) {
-                stringBuilder.append(buildFieldDataString(aggBuilder, successfulSearchShardIndices, showFieldName, showFieldType));
+                stringBuilder.append(buildFieldDataString(aggBuilder, propertiesAsMap, index, showFieldName, showFieldType));
             }
             stringBuilder.append("\n");
 
@@ -203,7 +233,8 @@ StringBuilder recursiveAggregationShapeBuilder(
                     baseIndent.append(ONE_SPACE_INDENT.repeat(4)),
                     showFieldName,
                     showFieldType,
-                    successfulSearchShardIndices
+                    propertiesAsMap,
+                    index
                 );
                 baseIndent.delete(0, 4);
             }
@@ -246,14 +277,16 @@ StringBuilder recursiveAggregationShapeBuilder(
      * @param sortBuilderList        search request sort builders list
      * @param showFieldName          whether to append field names
      * @param showFieldType          whether to append field types
-     * @param successfulSearchShardIndices the set of indices that were successfully searched
+     * @param propertiesAsMap       properties
+     * @param index                 index
      * @return Sort shape as a String
      */
     String buildSortShape(
         List<SortBuilder<?>> sortBuilderList,
         Boolean showFieldName,
         Boolean showFieldType,
-        Set<Index> successfulSearchShardIndices
+        Map<String, Object> propertiesAsMap,
+        Index index
     ) {
         if (sortBuilderList == null || sortBuilderList.isEmpty()) {
             return EMPTY_STRING;
@@ -266,7 +299,7 @@ String buildSortShape(
             StringBuilder stringBuilder = new StringBuilder();
             stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order());
             if (showFieldName || showFieldType) {
-                stringBuilder.append(buildFieldDataString(sortBuilder, successfulSearchShardIndices, showFieldName, showFieldType));
+                stringBuilder.append(buildFieldDataString(sortBuilder, propertiesAsMap, index, showFieldName, showFieldType));
             }
             shapeStrings.add(stringBuilder.toString());
         }
@@ -281,15 +314,17 @@ String buildSortShape(
      * Builds a field data string from a builder.
      *
      * @param builder                  aggregation or sort builder
-     * @param successfulSearchShardIndices the set of indices that were successfully searched
+     * @param propertiesAsMap       properties
+     * @param index                 index
      * @param showFieldName            whether to include field names
      * @param showFieldType            whether to include field types
      * @return Field data string
      * Ex: " [my_field, width:5]"
      */
     String buildFieldDataString(
         NamedWriteable builder,
-        Set<Index> successfulSearchShardIndices,
+        Map<String, Object> propertiesAsMap,
+        Index index,
         Boolean showFieldName,
         Boolean showFieldType
     ) {
@@ -300,7 +335,7 @@ String buildFieldDataString(
                 fieldDataList.add(fieldName);
             }
             if (showFieldType) {
-                String fieldType = getFieldType(fieldName, successfulSearchShardIndices);
+                String fieldType = getFieldType(fieldName, propertiesAsMap, index);
                 if (fieldType != null && !fieldType.isEmpty()) {
                     fieldDataList.add(fieldType);
                 }
@@ -309,11 +344,8 @@ String buildFieldDataString(
         return " [" + String.join(", ", fieldDataList) + "]";
     }
 
-    String getFieldType(String fieldName, Set<Index> successfulSearchShardIndices) {
-        if (successfulSearchShardIndices == null) {
-            return null;
-        }
-        String fieldType = getFieldTypeFromCache(fieldName, successfulSearchShardIndices);
+    String getFieldType(String fieldName, Map<String, Object> propertiesAsMap, Index index) {
+        String fieldType = getFieldTypeFromCache(fieldName, index);
         if (fieldType != null) {
             if (fieldType.equals(NO_FIELD_TYPE_VALUE)) {
                 return null;
@@ -322,32 +354,16 @@ String getFieldType(String fieldName, Set<Index> successfulSearchShardIndices) {
             }
         }
 
-        for (Index index : successfulSearchShardIndices) {
-            Map<String, MappingMetadata> indexMapping = null;
-            try {
-                indexMapping = clusterService.state().metadata().findMappings(new String[] { index.getName() }, input -> str -> true);
-            } catch (IOException e) {
-                // Skip the current index
-                continue;
-            }
-            MappingMetadata mappingMetadata = indexMapping.get(index.getName());
-            fieldType = getFieldTypeFromMapping(index, fieldName, mappingMetadata);
-            if (fieldType != null) {
-                String finalFieldType = fieldType;
-                fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).computeIfAbsent(fieldName, k -> finalFieldType);
-                return fieldType;
-            }
+        fieldType = getFieldTypeFromMapping(index, fieldName, propertiesAsMap);
+        if (fieldType != null) {
+            String finalFieldType = fieldType;
+            fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).computeIfAbsent(fieldName, k -> finalFieldType);
+            return fieldType;
         }
         return null;
     }
 
-    String getFieldTypeFromMapping(Index index, String fieldName, MappingMetadata mappingMetadata) {
-        // Get properties directly from the mapping metadata
-        Map<String, ?> propertiesAsMap = null;
-        if (mappingMetadata != null) {
-            propertiesAsMap = (Map<String, ?>) mappingMetadata.getSourceAsMap().get("properties");
-        }
-
+    String getFieldTypeFromMapping(Index index, String fieldName, Map<String, Object> propertiesAsMap) {
         // Recursively find the field type from properties
         if (propertiesAsMap != null) {
             String fieldType = findFieldTypeInProperties(propertiesAsMap, fieldName.split("\\."), 0);
@@ -394,14 +410,12 @@ private String findFieldTypeInProperties(Map<String, ?> properties, String[] fie
         return null;
     }
 
-    String getFieldTypeFromCache(String fieldName, Set<Index> successfulSearchShardIndices) {
-        for (Index index : successfulSearchShardIndices) {
-            Map<String, String> indexMap = fieldTypeMap.get(index);
-            if (indexMap != null) {
-                String fieldType = indexMap.get(fieldName);
-                if (fieldType != null) {
-                    return fieldType;
-                }
+    String getFieldTypeFromCache(String fieldName, Index index) {
+        Map<String, String> indexMap = fieldTypeMap.get(index);
+        if (indexMap != null) {
+            String fieldType = indexMap.get(fieldName);
+            if (fieldType != null) {
+                return fieldType;
             }
         }
         return null;

src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java

@@ -15,7 +15,6 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 import org.apache.lucene.search.BooleanClause;
 import org.opensearch.common.SetOnce;
 import org.opensearch.core.index.Index;
@@ -30,14 +29,15 @@ public final class QueryShapeVisitor implements QueryBuilderVisitor {
     private final SetOnce<String> fieldData = new SetOnce<>();
     private final Map<BooleanClause.Occur, List<QueryShapeVisitor>> childVisitors = new EnumMap<>(BooleanClause.Occur.class);
     private final QueryShapeGenerator queryShapeGenerator;
-    private final Set<Index> successfulSearchShardIndices;
+    private final Map<String, Object> propertiesAsMap;
+    private final Index index;
     private final Boolean showFieldName;
     private final Boolean showFieldType;
 
     @Override
     public void accept(QueryBuilder queryBuilder) {
         queryType.set(queryBuilder.getName());
-        fieldData.set(queryShapeGenerator.buildFieldDataString(queryBuilder, successfulSearchShardIndices, showFieldName, showFieldType));
+        fieldData.set(queryShapeGenerator.buildFieldDataString(queryBuilder, propertiesAsMap, index, showFieldName, showFieldType));
     }
 
     @Override
@@ -52,7 +52,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
 
             @Override
             public void accept(QueryBuilder qb) {
-                currentChild = new QueryShapeVisitor(queryShapeGenerator, successfulSearchShardIndices, showFieldName, showFieldType);
+                currentChild = new QueryShapeVisitor(queryShapeGenerator, propertiesAsMap, index, showFieldName, showFieldType);
                 childVisitorList.add(currentChild);
                 currentChild.accept(qb);
             }
@@ -119,12 +119,14 @@ public String prettyPrintTree(String indent, Boolean showFieldName, Boolean show
      */
     public QueryShapeVisitor(
         QueryShapeGenerator queryShapeGenerator,
-        Set<Index> successfulSearchShardIndices,
+        Map<String, Object> propertiesAsMap,
+        Index index,
         Boolean showFieldName,
         Boolean showFieldType
     ) {
         this.queryShapeGenerator = queryShapeGenerator;
-        this.successfulSearchShardIndices = successfulSearchShardIndices;
+        this.propertiesAsMap = propertiesAsMap;
+        this.index = index;
         this.showFieldName = showFieldName;
         this.showFieldType = showFieldType;
     }

src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java

@@ -53,10 +53,7 @@
 import org.opensearch.test.OpenSearchTestCase;
 
 public final class QueryShapeGeneratorTests extends OpenSearchTestCase {
-    final Set<Index> successfulSearchShardIndices = Set.of(
-        new Index("index1", UUID.randomUUID().toString()),
-        new Index("index2", UUID.randomUUID().toString())
-    );
+    final Set<Index> successfulSearchShardIndices = Set.of(new Index("index1", UUID.randomUUID().toString()));
     final QueryShapeGenerator queryShapeGenerator;
 
     private ClusterService mockClusterService;
@@ -174,9 +171,9 @@ public void testEmptyMappings() throws IOException {
 
     // Field type should be inferred from both the mappings
     public void testMultipleIndexMappings() throws IOException {
-        setUpMockMappings("index1", Map.of("properties", Map.of("field1", Map.of("type", "keyword"))));
+        setUpMockMappings("index2", Map.of("properties", Map.of("field1", Map.of("type", "keyword"))));
 
-        setUpMockMappings("index2", Map.of("properties", Map.of("field2", Map.of("type", "text"), "field4", Map.of("type", "long"))));
+        setUpMockMappings("index1", Map.of("properties", Map.of("field2", Map.of("type", "text"), "field4", Map.of("type", "long"))));
 
         SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder();
 
@@ -188,7 +185,7 @@ public void testMultipleIndexMappings() throws IOException {
             + "    match [field2, text]\n"
             + "    range [field4, long]\n"
             + "  should:\n"
-            + "    regexp [field3]\n"; // No field type found
+            + "    regexp [field3]\n";
         assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);
     }
 
@@ -598,14 +595,14 @@ public void testFieldTypeCachingForNonExistentField() throws IOException {
 
         queryShapeGeneratorSpy.buildShape(sourceBuilder, true, true, successfulSearchShardIndices);
 
-        verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), eq(successfulSearchShardIndices));
+        verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class));
 
         verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromMapping(any(Index.class), eq(nonExistentField), any());
 
         queryShapeGeneratorSpy.buildShape(sourceBuilder, true, true, successfulSearchShardIndices);
 
-        verify(queryShapeGeneratorSpy, times(4)).getFieldTypeFromMapping(any(Index.class), eq(nonExistentField), any());
-        verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), eq(successfulSearchShardIndices));
+        verify(queryShapeGeneratorSpy, times(2)).getFieldTypeFromMapping(any(Index.class), eq(nonExistentField), any());
+        verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class));
     }
 
     public void testMultifieldQueryCombined() throws IOException {

src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java

@@ -10,7 +10,7 @@
 
 import static org.mockito.Mockito.mock;
 
-import java.util.Set;
+import java.util.HashMap;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.ConstantScoreQueryBuilder;
@@ -31,7 +31,13 @@ public void testQueryShapeVisitor() {
                     .mustNot(new RegexpQueryBuilder("color", "red.*"))
             )
             .must(new TermsQueryBuilder("genre", "action", "drama", "romance"));
-        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(new QueryShapeGenerator(mock(ClusterService.class)), Set.of(), false, false);
+        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(
+            new QueryShapeGenerator(mock(ClusterService.class)),
+            new HashMap<>(),
+            null,
+            false,
+            false
+        );
         builder.visit(shapeVisitor);
         assertEquals(
             "{\"type\":\"bool\",\"must\"[{\"type\":\"term\"},{\"type\":\"terms\"}],\"filter\"[{\"type\":\"constant_score\",\"filter\"[{\"type\":\"range\"}]}],\"should\"[{\"type\":\"bool\",\"must\"[{\"type\":\"match\"}],\"must_not\"[{\"type\":\"regexp\"}]}]}",