Add query shape field type

Signed-off-by: David Zane <davizane@amazon.com>

src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

@@ -15,6 +15,7 @@
 import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNSizeSetting;
 import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNWindowSizeSetting;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -28,6 +29,7 @@
 import org.opensearch.action.search.SearchRequestContext;
 import org.opensearch.action.search.SearchRequestOperationsListener;
 import org.opensearch.action.search.SearchTask;
+import org.opensearch.cluster.metadata.MappingMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
@@ -54,6 +56,7 @@ public final class QueryInsightsListener extends SearchRequestOperationsListener
 
     private final QueryInsightsService queryInsightsService;
     private final ClusterService clusterService;
+    private final QueryShapeGenerator queryShapeGenerator;
 
     /**
      * Constructor for QueryInsightsListener
@@ -81,6 +84,7 @@ public QueryInsightsListener(
         super(initiallyEnabled);
         this.clusterService = clusterService;
         this.queryInsightsService = queryInsightsService;
+        this.queryShapeGenerator = new QueryShapeGenerator(clusterService);
 
         // Setting endpoints set up for top n queries, including enabling top n queries, window size, and top n size
         // Expected metricTypes are Latency, CPU, and Memory.
@@ -241,7 +245,15 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
                 );
             }
 
-            String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(request.source(), false);
+            String hashcode = queryShapeGenerator.getShapeHashCodeAsString(
+                request.source(),
+                false,
+                searchRequestContext.getSuccessfulSearchShardIndices()
+            );
+
+            String shape = queryShapeGenerator.buildShape(request.source(), true, searchRequestContext.getSuccessfulSearchShardIndices());
+            System.out.println("===== Query Shape =====");
+            System.out.println(shape);
 
             Map<Attribute, Object> attributes = new HashMap<>();
             attributes.put(Attribute.SEARCH_TYPE, request.searchType().toString().toLowerCase(Locale.ROOT));

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

@@ -0,0 +1,76 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.plugin.insights.core.service.categorizer;
+
+import java.util.Map;
+
+/**
+ * Visitor pattern for obtaining index mappings
+ *
+ * @opensearch.internal
+ */
+final class MappingVisitor {
+
+    private MappingVisitor() {}
+
+    @SuppressWarnings("unchecked")
+    static String visitMapping(Map<String, ?> propertiesAsMap, String[] fieldName, int idx) {
+        if (propertiesAsMap.containsKey(fieldName[idx])) {
+            Map<String, ?> fieldMapping = (Map<String, ?>) propertiesAsMap.get(fieldName[idx]);
+            if (idx == fieldName.length - 1) {
+                return (String) fieldMapping.get("type");
+            } else {
+                // Multi field case
+                if (fieldMapping.containsKey("fields")) {
+                    return visitMapping((Map<String, ?>) fieldMapping.get("fields"), fieldName, idx + 1);
+                } else {
+                    return null;
+                }
+            }
+        }
+
+        // fieldName not found at current level
+        // call visitMapping() for any fields with subfields (contains "properties" key)
+        for (Object v : propertiesAsMap.values()) {
+            if (v instanceof Map) {
+                Map<String, ?> fieldMapping = (Map<String, ?>) v;
+                if (fieldMapping.containsKey("properties")) {
+                    return visitMapping((Map<String, ?>) fieldMapping.get("properties"), fieldName, idx);
+                }
+            }
+        }
+        // no mapping found
+        return null;
+    }
+
+}

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

@@ -8,13 +8,20 @@
 
 package org.opensearch.plugin.insights.core.service.categorizer;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.cluster.metadata.MappingMetadata;
+import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.hash.MurmurHash3;
 import org.opensearch.core.common.io.stream.NamedWriteable;
+import org.opensearch.core.index.Index;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.WithFieldName;
 import org.opensearch.search.aggregations.AggregationBuilder;
@@ -29,21 +36,28 @@
 public class QueryShapeGenerator {
     static final String EMPTY_STRING = "";
     static final String ONE_SPACE_INDENT = " ";
+    private final ClusterService clusterService;
+    private final ConcurrentHashMap<Index, ConcurrentHashMap<String, String>> fieldTypeMap;
+
+    public QueryShapeGenerator(ClusterService clusterService) {
+        this.clusterService = clusterService;
+        this.fieldTypeMap = new ConcurrentHashMap<>();
+    }
 
     /**
      * Method to get query shape hash code given a source
      * @param source search request source
      * @param showFields whether to include field data in query shape
      * @return Hash code of query shape as a MurmurHash3.Hash128 object (128-bit)
      */
-    public static MurmurHash3.Hash128 getShapeHashCode(SearchSourceBuilder source, Boolean showFields) {
-        final String shape = buildShape(source, showFields);
+    public MurmurHash3.Hash128 getShapeHashCode(SearchSourceBuilder source, Boolean showFields, Set<Index> successfulSearchShardIndices) {
+        final String shape = buildShape(source, showFields, successfulSearchShardIndices);
         final BytesRef shapeBytes = new BytesRef(shape);
         return MurmurHash3.hash128(shapeBytes.bytes, 0, shapeBytes.length, 0, new MurmurHash3.Hash128());
     }
 
-    public static String getShapeHashCodeAsString(SearchSourceBuilder source, Boolean showFields) {
-        MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFields);
+    public String getShapeHashCodeAsString(SearchSourceBuilder source, Boolean showFields, Set<Index> successfulSearchShardIndices) {
+        MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFields, successfulSearchShardIndices);
         String hashAsString = Long.toHexString(hashcode.h1) + Long.toHexString(hashcode.h2);
         return hashAsString;
     }
@@ -54,11 +68,11 @@ public static String getShapeHashCodeAsString(SearchSourceBuilder source, Boolea
      * @param showFields whether to append field data
      * @return Search query shape as String
      */
-    public static String buildShape(SearchSourceBuilder source, Boolean showFields) {
+    public String buildShape(SearchSourceBuilder source, Boolean showFields, Set<Index> successfulSearchShardIndices) {
         StringBuilder shape = new StringBuilder();
-        shape.append(buildQueryShape(source.query(), showFields));
-        shape.append(buildAggregationShape(source.aggregations(), showFields));
-        shape.append(buildSortShape(source.sorts(), showFields));
+        shape.append(buildQueryShape(source.query(), showFields, successfulSearchShardIndices));
+        shape.append(buildAggregationShape(source.aggregations(), showFields, successfulSearchShardIndices));
+        shape.append(buildSortShape(source.sorts(), showFields, successfulSearchShardIndices));
         return shape.toString();
     }
 
@@ -68,11 +82,11 @@ public static String buildShape(SearchSourceBuilder source, Boolean showFields)
      * @param showFields whether to append field data
      * @return Query-section shape as String
      */
-    static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) {
+    String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields, Set<Index> successfulSearchShardIndices) {
         if (queryBuilder == null) {
             return EMPTY_STRING;
         }
-        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
+        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(this, successfulSearchShardIndices);
         queryBuilder.visit(shapeVisitor);
         return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFields);
     }
@@ -83,7 +97,11 @@ static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) {
      * @param showFields whether to append field data
      * @return Aggregation shape as String
      */
-    static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) {
+    String buildAggregationShape(
+        AggregatorFactories.Builder aggregationsBuilder,
+        Boolean showFields,
+        Set<Index> successfulSearchShardIndices
+    ) {
         if (aggregationsBuilder == null) {
             return EMPTY_STRING;
         }
@@ -92,17 +110,19 @@ static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuil
             aggregationsBuilder.getPipelineAggregatorFactories(),
             new StringBuilder(),
             new StringBuilder(),
-            showFields
+            showFields,
+            successfulSearchShardIndices
         );
         return aggregationShape.toString();
     }
 
-    static StringBuilder recursiveAggregationShapeBuilder(
+    StringBuilder recursiveAggregationShapeBuilder(
         Collection<AggregationBuilder> aggregationBuilders,
         Collection<PipelineAggregationBuilder> pipelineAggregations,
         StringBuilder outputBuilder,
         StringBuilder baseIndent,
-        Boolean showFields
+        Boolean showFields,
+        Set<Index> successfulSearchShardIndices
     ) {
         //// Normal Aggregations ////
         if (aggregationBuilders.isEmpty() == false) {
@@ -113,7 +133,7 @@ static StringBuilder recursiveAggregationShapeBuilder(
             StringBuilder stringBuilder = new StringBuilder();
             stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType());
             if (showFields) {
-                stringBuilder.append(buildFieldDataString(aggBuilder));
+                stringBuilder.append(buildFieldDataString(aggBuilder, successfulSearchShardIndices));
             }
             stringBuilder.append("\n");
 
@@ -124,7 +144,8 @@ static StringBuilder recursiveAggregationShapeBuilder(
                     aggBuilder.getPipelineAggregations(),
                     stringBuilder,
                     baseIndent.append(ONE_SPACE_INDENT.repeat(4)),
-                    showFields
+                    showFields,
+                    successfulSearchShardIndices
                 );
                 baseIndent.delete(0, 4);
             }
@@ -167,7 +188,7 @@ static StringBuilder recursiveAggregationShapeBuilder(
      * @param showFields whether to append field data
      * @return Sort shape as String
      */
-    static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields) {
+    String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields, Set<Index> successfulSearchShardIndices) {
         if (sortBuilderList == null || sortBuilderList.isEmpty()) {
             return EMPTY_STRING;
         }
@@ -179,7 +200,7 @@ static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showF
             StringBuilder stringBuilder = new StringBuilder();
             stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order());
             if (showFields) {
-                stringBuilder.append(buildFieldDataString(sortBuilder));
+                stringBuilder.append(buildFieldDataString(sortBuilder, successfulSearchShardIndices));
             }
             shapeStrings.add(stringBuilder.toString());
         }
@@ -195,11 +216,47 @@ static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showF
      * @return String: comma separated list with leading space in square brackets
      * Ex: " [my_field, width:5]"
      */
-    static String buildFieldDataString(NamedWriteable builder) {
+    String buildFieldDataString(NamedWriteable builder, Set<Index> successfulSearchShardIndices) {
         List<String> fieldDataList = new ArrayList<>();
         if (builder instanceof WithFieldName) {
-            fieldDataList.add(((WithFieldName) builder).fieldName());
+            String fieldName = ((WithFieldName) builder).fieldName();
+            fieldDataList.add(fieldName);
+            fieldDataList.add(getFieldType(fieldName, successfulSearchShardIndices));
         }
         return " [" + String.join(", ", fieldDataList) + "]";
     }
+
+    String getFieldType(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;
+                }
+            }
+        }
+        Map<String, MappingMetadata> allIndicesMap;
+        try {
+            allIndicesMap = clusterService.state()
+                .metadata()
+                .findMappings(successfulSearchShardIndices.stream().map(Index::getName).toArray(String[]::new), input -> (str -> true));
+        } catch (IOException e) {
+            return null;
+        }
+        for (Index index : successfulSearchShardIndices) {
+            MappingMetadata mappingMetadata = allIndicesMap.get(index.getName());
+            if (mappingMetadata != null) {
+                @SuppressWarnings("unchecked")
+                Map<String, ?> propertiesAsMap = (Map<String, ?>) mappingMetadata.getSourceAsMap().get("properties");
+                String fieldType = MappingVisitor.visitMapping(propertiesAsMap, fieldName.split("\\."), 0);
+                if (fieldType != null) {
+                    // add item to cache
+                    fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).put(fieldName, fieldType);
+                    return fieldType;
+                }
+            }
+        }
+        return null;
+    }
 }

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

@@ -9,15 +9,16 @@
 package org.opensearch.plugin.insights.core.service.categorizer;
 
 import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.ONE_SPACE_INDENT;
-import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.buildFieldDataString;
 
 import java.util.ArrayList;
 import java.util.EnumMap;
 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;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilderVisitor;
 
@@ -28,11 +29,13 @@ public final class QueryShapeVisitor implements QueryBuilderVisitor {
     private final SetOnce<String> queryType = new SetOnce<>();
     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;
 
     @Override
     public void accept(QueryBuilder queryBuilder) {
         queryType.set(queryBuilder.getName());
-        fieldData.set(buildFieldDataString(queryBuilder));
+        fieldData.set(queryShapeGenerator.buildFieldDataString(queryBuilder, successfulSearchShardIndices));
     }
 
     @Override
@@ -47,7 +50,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
 
             @Override
             public void accept(QueryBuilder qb) {
-                currentChild = new QueryShapeVisitor();
+                currentChild = new QueryShapeVisitor(queryShapeGenerator, successfulSearchShardIndices);
                 childVisitorList.add(currentChild);
                 currentChild.accept(qb);
             }
@@ -110,5 +113,8 @@ public String prettyPrintTree(String indent, Boolean showFields) {
     /**
      * Default constructor
      */
-    public QueryShapeVisitor() {}
+    public QueryShapeVisitor(QueryShapeGenerator queryShapeGenerator, Set<Index> successfulSearchShardIndices) {
+        this.queryShapeGenerator = queryShapeGenerator;
+        this.successfulSearchShardIndices = successfulSearchShardIndices;
+    }
 }

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

@@ -87,11 +87,6 @@ public void categorize(SearchQueryRecord record) {
         incrementQueryTypeCounters(source.query(), measurements);
         incrementQueryAggregationCounters(source.aggregations(), measurements);
         incrementQuerySortCounters(source.sorts(), measurements);
-
-        if (logger.isTraceEnabled()) {
-            String searchShape = QueryShapeGenerator.buildShape(source, true);
-            logger.trace(searchShape);
-        }
     }
 
     private void incrementQuerySortCounters(List<SortBuilder<?>> sorts, Map<MetricType, Measurement> measurements) {