Skip to content

Commit

Permalink
Add field type to query shape
Browse files Browse the repository at this point in the history 
Signed-off-by: David Zane <davizane@amazon.com>
  • Loading branch information
@dzane17
dzane17 committed Sep 26, 2024
1 parent f1827d8 commit 62d0218
Show file tree
Hide file tree
Showing 6 changed files with 266 additions and 45 deletions.
12 changes: 11 additions & 1 deletion src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
Expand Down Expand Up @@ -239,7 +240,16 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
);
}

String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(request.source(), false);
Set<String> successfulShardIndices = searchRequestContext.getSuccessfulSearchShardIndices();
if (successfulShardIndices == null) {
successfulShardIndices = Collections.emptySet();
}
String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(
request.source(),
false,
clusterService.state().getMetadata(),
successfulShardIndices
);

Map<Attribute, Object> attributes = new HashMap<>();
attributes.put(Attribute.SEARCH_TYPE, request.searchType().toString().toLowerCase(Locale.ROOT));
Expand Down
110 changes: 91 additions & 19 deletions ...ain/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,12 @@
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.lucene.util.BytesRef;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.common.hash.MurmurHash3;
import org.opensearch.core.common.io.stream.NamedWriteable;
import org.opensearch.index.query.QueryBuilder;
Expand All @@ -36,14 +41,24 @@ public class QueryShapeGenerator {
* @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 static MurmurHash3.Hash128 getShapeHashCode(
SearchSourceBuilder source,
Boolean showFields,
Metadata metadata,
Set<String> searchIndices
) {
final String shape = buildShape(source, showFields, metadata, searchIndices);
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 static String getShapeHashCodeAsString(
SearchSourceBuilder source,
Boolean showFields,
Metadata metadata,
Set<String> searchIndices
) {
MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFields, metadata, searchIndices);
String hashAsString = Long.toHexString(hashcode.h1) + Long.toHexString(hashcode.h2);
return hashAsString;
}
Expand All @@ -54,11 +69,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 static String buildShape(SearchSourceBuilder source, Boolean showFields, Metadata metadata, Set<String> searchIndices) {
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, metadata, searchIndices));
shape.append(buildAggregationShape(source.aggregations(), showFields, metadata, searchIndices));
shape.append(buildSortShape(source.sorts(), showFields, metadata, searchIndices));
return shape.toString();
}

Expand All @@ -68,11 +83,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) {
static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields, Metadata metadata, Set<String> searchIndices) {
if (queryBuilder == null) {
return EMPTY_STRING;
}
QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(metadata, searchIndices);
queryBuilder.visit(shapeVisitor);
return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFields);
}
Expand All @@ -83,7 +98,12 @@ 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) {
static String buildAggregationShape(
AggregatorFactories.Builder aggregationsBuilder,
Boolean showFields,
Metadata metadata,
Set<String> searchIndices
) {
if (aggregationsBuilder == null) {
return EMPTY_STRING;
}
Expand All @@ -92,7 +112,9 @@ static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuil
aggregationsBuilder.getPipelineAggregatorFactories(),
new StringBuilder(),
new StringBuilder(),
showFields
showFields,
metadata,
searchIndices
);
return aggregationShape.toString();
}
Expand All @@ -102,7 +124,9 @@ static StringBuilder recursiveAggregationShapeBuilder(
Collection<PipelineAggregationBuilder> pipelineAggregations,
StringBuilder outputBuilder,
StringBuilder baseIndent,
Boolean showFields
Boolean showFields,
Metadata metadata,
Set<String> searchIndices
) {
//// Normal Aggregations ////
if (aggregationBuilders.isEmpty() == false) {
Expand All @@ -113,7 +137,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, metadata, searchIndices));
}
stringBuilder.append("\n");

Expand All @@ -124,7 +148,9 @@ static StringBuilder recursiveAggregationShapeBuilder(
aggBuilder.getPipelineAggregations(),
stringBuilder,
baseIndent.append(ONE_SPACE_INDENT.repeat(4)),
showFields
showFields,
metadata,
searchIndices
);
baseIndent.delete(0, 4);
}
Expand Down Expand Up @@ -167,7 +193,7 @@ static StringBuilder recursiveAggregationShapeBuilder(
* @param showFields whether to append field data
* @return Sort shape as String
*/
static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields) {
static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields, Metadata metadata, Set<String> searchIndices) {
if (sortBuilderList == null || sortBuilderList.isEmpty()) {
return EMPTY_STRING;
}
Expand All @@ -179,7 +205,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, metadata, searchIndices));
}
shapeStrings.add(stringBuilder.toString());
}
Expand All @@ -195,11 +221,57 @@ 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) {
static String buildFieldDataString(NamedWriteable builder, Metadata metadata, Set<String> searchIndices) {
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, metadata, searchIndices));
}
return " [" + String.join(", ", fieldDataList) + "]";
}

/**
* Method to get a field's type
* @return String field type
*/
static String getFieldType(String fieldName, Metadata metadata, Set<String> searchIndices) {
if (metadata == null) {
return "none";
}
for (String searchIndex : searchIndices) {
IndexMetadata indexMetadata = metadata.index(searchIndex);
if (indexMetadata == null) {
continue;
}
MappingMetadata mappingMetadata = indexMetadata.mapping();
if (mappingMetadata == null) {
continue;
}
Map<String, Object> sourceMap = mappingMetadata.getSourceAsMap();

String fieldType = findFieldType((Map<String, Object>) sourceMap.get("properties"), fieldName);
if (fieldType != null) {
return fieldType;
}
}
return "none";
}

public static String findFieldType(Map<String, Object> properties, String fieldName) {
for (Map.Entry<String, Object> entry : properties.entrySet()) {
Map<String, Object> field = (Map<String, Object>) entry.getValue();
if (entry.getKey().equals(fieldName)) {
return (String) field.get("type");
}
// TODO: Add support for multi-fields
// if (field.containsKey("fields")) {
// String type = findFieldType((Map<String, Object>) field.get("fields"), fieldName);
// if (type != null) {
// return type;
// }
// }
}
return null;
}
}
13 changes: 10 additions & 3 deletions src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@
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.cluster.metadata.Metadata;
import org.opensearch.common.SetOnce;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilderVisitor;
Expand All @@ -28,11 +30,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 Metadata metadata;
private final Set<String> searchIndices;

@Override
public void accept(QueryBuilder queryBuilder) {
queryType.set(queryBuilder.getName());
fieldData.set(buildFieldDataString(queryBuilder));
fieldData.set(buildFieldDataString(queryBuilder, metadata, searchIndices));
}

@Override
Expand All @@ -47,7 +51,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {

@Override
public void accept(QueryBuilder qb) {
currentChild = new QueryShapeVisitor();
currentChild = new QueryShapeVisitor(metadata, searchIndices);
childVisitorList.add(currentChild);
currentChild.accept(qb);
}
Expand Down Expand Up @@ -110,5 +114,8 @@ public String prettyPrintTree(String indent, Boolean showFields) {
/**
* Default constructor
*/
public QueryShapeVisitor() {}
public QueryShapeVisitor(Metadata metadata, Set<String> searchIndices) {
this.metadata = metadata;
this.searchIndices = searchIndices;
}
}
5 changes: 0 additions & 5 deletions .../java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading

0 comments on commit 62d0218

Please sign in to comment.