Skip to content

Commit

Permalink
Query instrumentation Minor Enhancements (opensearch-project#73)
Browse files Browse the repository at this point in the history
* Query instrumentation minor fixes

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

* Fix unit tests

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

* Rename package

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

* Spotless

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

---------

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
  • Loading branch information
deshsidd committed Aug 28, 2024
1 parent aeeaa24 commit 00efb8b
Show file tree
Hide file tree
Showing 5 changed files with 217 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
*/
public class SearchQueryAggregationCategorizer {

private static final String TYPE_TAG = "type";
static final String AGGREGATION_TYPE_TAG = "agg_type";
private final SearchQueryCounters searchQueryCounters;

/**
Expand Down Expand Up @@ -49,7 +49,7 @@ public void incrementSearchQueryAggregationCounters(
private void incrementCountersRecursively(AggregationBuilder aggregationBuilder, Map<MetricType, Number> measurements) {
// Increment counters for the current aggregation
String aggregationType = aggregationBuilder.getType();
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, aggregationType), measurements);
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(AGGREGATION_TYPE_TAG, aggregationType), measurements);

// Recursively process sub-aggregations if any
Collection<AggregationBuilder> subAggregations = aggregationBuilder.getSubAggregations();
Expand All @@ -63,7 +63,7 @@ private void incrementCountersRecursively(AggregationBuilder aggregationBuilder,
Collection<PipelineAggregationBuilder> pipelineAggregations = aggregationBuilder.getPipelineAggregations();
for (PipelineAggregationBuilder pipelineAggregation : pipelineAggregations) {
String pipelineAggregationType = pipelineAggregation.getType();
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(TYPE_TAG, pipelineAggregationType), measurements);
searchQueryCounters.incrementAggCounter(1, Tags.create().addTag(AGGREGATION_TYPE_TAG, pipelineAggregationType), measurements);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,12 @@
*/
public final class SearchQueryCounters {
private static final String LEVEL_TAG = "level";
private static final String TYPE_TAG = "type";
private static final String QUERY_TYPE_TAG = "type";
private static final String UNIT = "1";
private static final String UNIT_MILLIS = "ms";
private static final String UNIT_CPU_CYCLES = "ns";
private static final String UNIT_BYTES = "bytes";

private final MetricsRegistry metricsRegistry;
/**
* Aggregation counter
Expand Down Expand Up @@ -83,17 +87,17 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) {
this.queryTypeLatencyHistogram = metricsRegistry.createHistogram(
"search.query.type.latency.histogram",
"Histogram for the latency per query type",
UNIT
UNIT_MILLIS
);
this.queryTypeCpuHistogram = metricsRegistry.createHistogram(
"search.query.type.cpu.histogram",
"Histogram for the cpu per query type",
UNIT
UNIT_CPU_CYCLES
);
this.queryTypeMemoryHistogram = metricsRegistry.createHistogram(
"search.query.type.memory.histogram",
"Histogram for the memory per query type",
UNIT
UNIT_BYTES
);
this.queryHandlers = new HashMap<>();
}
Expand All @@ -109,7 +113,7 @@ public void incrementCounter(QueryBuilder queryBuilder, int level, Map<MetricTyp

Counter counter = nameToQueryTypeCounters.computeIfAbsent(uniqueQueryCounterName, k -> createQueryCounter(k));
counter.add(1, Tags.create().addTag(LEVEL_TAG, level));
incrementAllHistograms(Tags.create().addTag(LEVEL_TAG, level).addTag(TYPE_TAG, uniqueQueryCounterName), measurements);
incrementAllHistograms(Tags.create().addTag(LEVEL_TAG, level).addTag(QUERY_TYPE_TAG, uniqueQueryCounterName), measurements);
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,201 @@
/*
* 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.
*/

package org.opensearch.plugin.insights.core.service.categorizer;

import org.opensearch.common.hash.MurmurHash3;
import org.opensearch.plugin.insights.SearchSourceBuilderUtils;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.test.OpenSearchTestCase;

public final class QueryShapeGeneratorTests extends OpenSearchTestCase {

public void testComplexSearch() {
SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createDefaultSearchSourceBuilder();

String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
String expectedShowFieldsTrue = "bool []\n"
+ " must:\n"
+ " term [field1]\n"
+ " filter:\n"
+ " match [field2]\n"
+ " range [field4]\n"
+ " should:\n"
+ " regexp [field3]\n"
+ "aggregation:\n"
+ " significant_text []\n"
+ " terms [key]\n"
+ " aggregation:\n"
+ " terms [key.sub1]\n"
+ " terms [key.sub2]\n"
+ " pipeline aggregation:\n"
+ " max_bucket\n"
+ " terms [model]\n"
+ " terms [type]\n"
+ " aggregation:\n"
+ " terms [key.sub3]\n"
+ " pipeline aggregation:\n"
+ " derivative\n"
+ " top_hits []\n"
+ " pipeline aggregation:\n"
+ " avg_bucket\n"
+ " derivative\n"
+ " max_bucket\n"
+ "sort:\n"
+ " asc [album]\n"
+ " asc [price]\n"
+ " desc [color]\n"
+ " desc [vendor]\n";
assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);

String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
String expectedShowFieldsFalse = "bool\n"
+ " must:\n"
+ " term\n"
+ " filter:\n"
+ " match\n"
+ " range\n"
+ " should:\n"
+ " regexp\n"
+ "aggregation:\n"
+ " significant_text\n"
+ " terms\n"
+ " terms\n"
+ " aggregation:\n"
+ " terms\n"
+ " pipeline aggregation:\n"
+ " derivative\n"
+ " terms\n"
+ " aggregation:\n"
+ " terms\n"
+ " terms\n"
+ " pipeline aggregation:\n"
+ " max_bucket\n"
+ " top_hits\n"
+ " pipeline aggregation:\n"
+ " avg_bucket\n"
+ " derivative\n"
+ " max_bucket\n"
+ "sort:\n"
+ " asc\n"
+ " asc\n"
+ " desc\n"
+ " desc\n";
assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
}

public void testQueryShape() {
SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder();

String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
String expectedShowFieldsTrue = "bool []\n"
+ " must:\n"
+ " term [field1]\n"
+ " filter:\n"
+ " match [field2]\n"
+ " range [field4]\n"
+ " should:\n"
+ " regexp [field3]\n";
assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);

String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
String expectedShowFieldsFalse = "bool\n"
+ " must:\n"
+ " term\n"
+ " filter:\n"
+ " match\n"
+ " range\n"
+ " should:\n"
+ " regexp\n";
assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
}

public void testAggregationShape() {
SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createAggregationSearchSourceBuilder();

String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
String expectedShowFieldsTrue = "aggregation:\n"
+ " significant_text []\n"
+ " terms [key]\n"
+ " aggregation:\n"
+ " terms [key.sub1]\n"
+ " terms [key.sub2]\n"
+ " pipeline aggregation:\n"
+ " max_bucket\n"
+ " terms [model]\n"
+ " terms [type]\n"
+ " aggregation:\n"
+ " terms [key.sub3]\n"
+ " pipeline aggregation:\n"
+ " derivative\n"
+ " top_hits []\n"
+ " pipeline aggregation:\n"
+ " avg_bucket\n"
+ " derivative\n"
+ " max_bucket\n";
assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);

String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
String expectedShowFieldsFalse = "aggregation:\n"
+ " significant_text\n"
+ " terms\n"
+ " terms\n"
+ " aggregation:\n"
+ " terms\n"
+ " pipeline aggregation:\n"
+ " derivative\n"
+ " terms\n"
+ " aggregation:\n"
+ " terms\n"
+ " terms\n"
+ " pipeline aggregation:\n"
+ " max_bucket\n"
+ " top_hits\n"
+ " pipeline aggregation:\n"
+ " avg_bucket\n"
+ " derivative\n"
+ " max_bucket\n";
assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
}

public void testSortShape() {
SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createSortSearchSourceBuilder();

String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
String expectedShowFieldsTrue = "sort:\n" + " asc [album]\n" + " asc [price]\n" + " desc [color]\n" + " desc [vendor]\n";
assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);

String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
String expectedShowFieldsFalse = "sort:\n" + " asc\n" + " asc\n" + " desc\n" + " desc\n";
assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
}

public void testHashCode() {
// Create test source builders
SearchSourceBuilder defaultSourceBuilder = SearchSourceBuilderUtils.createDefaultSearchSourceBuilder();
SearchSourceBuilder querySourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder();

// showFields true
MurmurHash3.Hash128 defaultHashTrue = QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true);
MurmurHash3.Hash128 queryHashTrue = QueryShapeGenerator.getShapeHashCode(querySourceBuilder, true);
assertEquals(defaultHashTrue, QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true));
assertEquals(queryHashTrue, QueryShapeGenerator.getShapeHashCode(querySourceBuilder, true));
assertNotEquals(defaultHashTrue, queryHashTrue);

// showFields false
MurmurHash3.Hash128 defaultHashFalse = QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false);
MurmurHash3.Hash128 queryHashFalse = QueryShapeGenerator.getShapeHashCode(querySourceBuilder, false);
assertEquals(defaultHashFalse, QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false));
assertEquals(queryHashFalse, QueryShapeGenerator.getShapeHashCode(querySourceBuilder, false));
assertNotEquals(defaultHashFalse, queryHashFalse);

// Compare field data on vs off
assertNotEquals(defaultHashTrue, defaultHashFalse);
assertNotEquals(queryHashTrue, queryHashFalse);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.service.categorizor;
package org.opensearch.plugin.insights.core.service.categorizer;

import org.opensearch.index.query.BoolQueryBuilder;
import org.opensearch.index.query.ConstantScoreQueryBuilder;
Expand All @@ -16,7 +16,6 @@
import org.opensearch.index.query.RegexpQueryBuilder;
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.index.query.TermsQueryBuilder;
import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeVisitor;
import org.opensearch.test.OpenSearchTestCase;

public final class QueryShapeVisitorTests extends OpenSearchTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
* compatible open source license.
*/

package org.opensearch.plugin.insights.core.service.categorizor;
package org.opensearch.plugin.insights.core.service.categorizer;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
Expand All @@ -15,6 +15,7 @@
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.opensearch.plugin.insights.QueryInsightsTestUtils.generateQueryInsightRecords;
import static org.opensearch.plugin.insights.core.service.categorizer.SearchQueryAggregationCategorizer.AGGREGATION_TYPE_TAG;

import java.util.Arrays;
import java.util.HashMap;
Expand All @@ -36,7 +37,6 @@
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.index.query.WildcardQueryBuilder;
import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.opensearch.plugin.insights.core.service.categorizer.SearchQueryCategorizer;
import org.opensearch.plugin.insights.rules.model.MetricType;
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder;
Expand Down Expand Up @@ -114,7 +114,7 @@ public void testAggregationsQuery() {
verify(searchQueryCategorizer.getSearchQueryCounters().getAggCounter()).add(valueCaptor.capture(), tagsCaptor.capture());

double actualValue = valueCaptor.getValue();
String actualTag = (String) tagsCaptor.getValue().getTagsMap().get("type");
String actualTag = (String) tagsCaptor.getValue().getTagsMap().get(AGGREGATION_TYPE_TAG);

assertEquals(1.0d, actualValue, 0.0001);
assertEquals(MULTI_TERMS_AGGREGATION, actualTag);
Expand Down

0 comments on commit 00efb8b

Please sign in to comment.