-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: David Zane <davizane@amazon.com>
- Loading branch information
Showing
2 changed files
with
143 additions
and
109 deletions.
There are no files selected for viewing
138 changes: 138 additions & 0 deletions
138
src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,138 @@ | ||
| /* | ||
| * 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.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| import org.opensearch.index.query.QueryBuilder; | ||
| import org.opensearch.search.aggregations.AggregationBuilder; | ||
| import org.opensearch.search.aggregations.AggregatorFactories; | ||
| import org.opensearch.search.aggregations.PipelineAggregationBuilder; | ||
| import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder; | ||
| import org.opensearch.search.builder.SearchSourceBuilder; | ||
| import org.opensearch.search.sort.FieldSortBuilder; | ||
| import org.opensearch.search.sort.SortBuilder; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.Collection; | ||
| import java.util.Collections; | ||
| import java.util.List; | ||
|
|
||
| class QueryShapeService { | ||
| static final String TWO_SPACE_INDENT = " "; | ||
|
|
||
| static String buildShape(SearchSourceBuilder source, Boolean showFields) { | ||
| StringBuilder shape = new StringBuilder(); | ||
| shape.append(buildQueryShape(source.query())); | ||
| shape.append(buildAggregationShape(source.aggregations(), showFields)); | ||
| shape.append(buildSortShape(source.sorts(), showFields)); | ||
| return shape.toString(); | ||
| } | ||
|
|
||
| static String buildQueryShape(QueryBuilder queryBuilder) { | ||
| if (queryBuilder == null) { | ||
| return ""; | ||
| } | ||
| QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); | ||
| queryBuilder.visit(shapeVisitor); | ||
| return shapeVisitor.prettyPrintTree(""); | ||
| } | ||
|
|
||
| static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) { | ||
| if (aggregationsBuilder == null) { | ||
| return ""; | ||
| } | ||
| StringBuilder aggregationShape = recursiveAggregationShapeBuilder( | ||
| aggregationsBuilder.getAggregatorFactories(), | ||
| aggregationsBuilder.getPipelineAggregatorFactories(), | ||
| new StringBuilder(), | ||
| 0, | ||
| showFields | ||
| ); | ||
| return aggregationShape.toString(); | ||
| } | ||
|
|
||
| static StringBuilder recursiveAggregationShapeBuilder( | ||
| Collection<AggregationBuilder> aggregationBuilders, | ||
| Collection<PipelineAggregationBuilder> pipelineAggregations, | ||
| StringBuilder outputBuilder, | ||
| int indentCount, | ||
| Boolean showFields | ||
| ) { | ||
| String baseIndent = TWO_SPACE_INDENT.repeat(indentCount); | ||
|
|
||
| //// Normal Aggregations //// | ||
| if (aggregationBuilders.isEmpty() == false) { | ||
| outputBuilder.append(baseIndent).append("aggregation:").append("\n"); | ||
| } | ||
| List<String> aggShapeStrings = new ArrayList<>(); | ||
| for (AggregationBuilder agg : aggregationBuilders) { | ||
| StringBuilder stringBuilder = new StringBuilder(); | ||
| stringBuilder.append(baseIndent).append(TWO_SPACE_INDENT).append(agg.getType()); | ||
| if (showFields) { | ||
| if (agg instanceof ValuesSourceAggregationBuilder) { | ||
| stringBuilder.append(" [").append(((ValuesSourceAggregationBuilder<?>) agg).field()).append("]"); | ||
| } else { | ||
| stringBuilder.append(" []"); | ||
| } | ||
| } | ||
| stringBuilder.append("\n"); | ||
|
|
||
| if (agg.getSubAggregations().isEmpty() == false) { | ||
| // Recursive call on sub-aggregations | ||
| recursiveAggregationShapeBuilder(agg.getSubAggregations(), agg.getPipelineAggregations(), stringBuilder, indentCount + 2, showFields); | ||
| } | ||
| aggShapeStrings.add(stringBuilder.toString()); | ||
| } | ||
|
|
||
| // Sort aggregations | ||
| Collections.sort(aggShapeStrings); | ||
| for (String shapeString : aggShapeStrings) { | ||
| outputBuilder.append(shapeString); | ||
| } | ||
|
|
||
| //// Pipeline Aggregation (cannot have sub-aggregations) //// | ||
| if (pipelineAggregations.isEmpty() == false) { | ||
| outputBuilder.append(baseIndent).append("pipeline aggregation:").append("\n"); | ||
| } | ||
| for (PipelineAggregationBuilder pipelineAgg : pipelineAggregations) { | ||
| outputBuilder.append(baseIndent).append(TWO_SPACE_INDENT).append(pipelineAgg.getName()).append("\n"); | ||
| } | ||
|
|
||
| return outputBuilder; | ||
| } | ||
|
|
||
| static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields) { | ||
| if (sortBuilderList == null || sortBuilderList.isEmpty()) { | ||
| return ""; | ||
| } | ||
| StringBuilder sortShape = new StringBuilder(); | ||
| sortShape.append("sort:\n"); | ||
|
|
||
| List<String> shapeStrings = new ArrayList<>(); | ||
| for (SortBuilder<?> sortBuilder : sortBuilderList) { | ||
| StringBuilder stringBuilder = new StringBuilder(); | ||
| stringBuilder.append(TWO_SPACE_INDENT).append(sortBuilder.order()); | ||
| if (showFields) { | ||
| if (sortBuilder instanceof FieldSortBuilder) { | ||
| stringBuilder.append(" [").append(((FieldSortBuilder) sortBuilder).getFieldName()).append("]"); | ||
| } else { | ||
| stringBuilder.append(" []"); | ||
| } | ||
| } | ||
| shapeStrings.add(stringBuilder.toString()); | ||
| } | ||
|
|
||
| Collections.sort(shapeStrings); | ||
| for (String line : shapeStrings) { | ||
| sortShape.append(line).append("\n"); | ||
| } | ||
| return sortShape.toString(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters