From e79b84f2efeec67e64bd1d5991c3cc475c289a29 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Thu, 10 Oct 2024 00:15:36 -0700 Subject: [PATCH] Use only first index and optimization to get properties once per query Signed-off-by: Siddhant Deshmukh --- .../categorizer/QueryShapeGenerator.java | 120 ++++++++++-------- .../categorizer/QueryShapeVisitor.java | 14 +- .../categorizer/QueryShapeGeneratorTests.java | 17 +-- .../categorizer/QueryShapeVisitorTests.java | 10 +- 4 files changed, 90 insertions(+), 71 deletions(-) diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java index 1d2ed94..ee896af 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java @@ -112,32 +112,58 @@ public String buildShape( Boolean showFieldType, Set successfulSearchShardIndices ) { + Index firstIndex = successfulSearchShardIndices.iterator().next(); + Map 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 getPropertiesMapForIndex(Index index) { + Map 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 propertiesMap = (Map) 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 successfulSearchShardIndices + Map 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 successfulSearchShardIndices + Map 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 successfulSearchShardIndices + Map 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> sortBuilderList, Boolean showFieldName, Boolean showFieldType, - Set successfulSearchShardIndices + Map 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,7 +314,8 @@ 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 @@ -289,7 +323,8 @@ String buildSortShape( */ String buildFieldDataString( NamedWriteable builder, - Set successfulSearchShardIndices, + Map 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 successfulSearchShardIndices) { - if (successfulSearchShardIndices == null) { - return null; - } - String fieldType = getFieldTypeFromCache(fieldName, successfulSearchShardIndices); + String getFieldType(String fieldName, Map 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 successfulSearchShardIndices) { } } - for (Index index : successfulSearchShardIndices) { - Map 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 propertiesAsMap = null; - if (mappingMetadata != null) { - propertiesAsMap = (Map) mappingMetadata.getSourceAsMap().get("properties"); - } - + String getFieldTypeFromMapping(Index index, String fieldName, Map 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 properties, String[] fie return null; } - String getFieldTypeFromCache(String fieldName, Set successfulSearchShardIndices) { - for (Index index : successfulSearchShardIndices) { - Map indexMap = fieldTypeMap.get(index); - if (indexMap != null) { - String fieldType = indexMap.get(fieldName); - if (fieldType != null) { - return fieldType; - } + String getFieldTypeFromCache(String fieldName, Index index) { + Map indexMap = fieldTypeMap.get(index); + if (indexMap != null) { + String fieldType = indexMap.get(fieldName); + if (fieldType != null) { + return fieldType; } } return null; diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java index fe553e9..a046a52 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java +++ b/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 fieldData = new SetOnce<>(); private final Map> childVisitors = new EnumMap<>(BooleanClause.Occur.class); private final QueryShapeGenerator queryShapeGenerator; - private final Set successfulSearchShardIndices; + private final Map 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 successfulSearchShardIndices, + Map 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; } diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java index a47f434..b1499f9 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java +++ b/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 successfulSearchShardIndices = Set.of( - new Index("index1", UUID.randomUUID().toString()), - new Index("index2", UUID.randomUUID().toString()) - ); + final Set 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 { diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java index 9c83736..5c9e4ba 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java +++ b/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\"}]}]}",