Add field type to query shape and extensive unit tests (#140)

* Add query type to shape and extensive unit tests Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> * Log query shape if trace log enabled Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> * Refactor code, handle multifield with test, remove properties cache Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> * Use only first index and optimization to get properties once per query Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> * Further refactoring Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> --------- Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
opensearch-project · Oct 10, 2024 · dcf52c2 · dcf52c2
1 parent f89eb88
commit dcf52c2
Show file tree

Hide file tree

Showing 6 changed files with 957 additions and 103 deletions.
diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
@@ -58,6 +58,7 @@ public final class QueryInsightsListener extends SearchRequestOperationsListener
     private final ClusterService clusterService;
     private boolean groupingFieldNameEnabled;
     private boolean groupingFieldTypeEnabled;
+    private final QueryShapeGenerator queryShapeGenerator;
 
     /**
      * Constructor for QueryInsightsListener
@@ -87,6 +88,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.
@@ -270,9 +272,25 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
             attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap());
             attributes.put(Attribute.TASK_RESOURCE_USAGES, tasksResourceUsages);
 
-            if (queryInsightsService.isGroupingEnabled()) {
-                String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(request.source(), groupingFieldNameEnabled);
-                attributes.put(Attribute.QUERY_HASHCODE, hashcode);
+            if (queryInsightsService.isGroupingEnabled() || log.isTraceEnabled()) {
+                // Generate the query shape only if grouping is enabled or trace logging is enabled
+                final String queryShape = queryShapeGenerator.buildShape(
+                    request.source(),
+                    groupingFieldNameEnabled,
+                    groupingFieldTypeEnabled,
+                    searchRequestContext.getSuccessfulSearchShardIndices()
+                );
+
+                // Print the query shape if tracer is enabled
+                if (log.isTraceEnabled()) {
+                    log.trace("Query Shape:\n{}", queryShape);
+                }
+
+                // Add hashcode attribute when grouping is enabled
+                if (queryInsightsService.isGroupingEnabled()) {
+                    String hashcode = queryShapeGenerator.getShapeHashCodeAsString(queryShape);
+                    attributes.put(Attribute.QUERY_HASHCODE, hashcode);
+                }
             }
 
             Map<String, Object> labels = new HashMap<>();