diff --git a/src/main/java/querqy/elasticsearch/QuerqyPlugin.java b/src/main/java/querqy/elasticsearch/QuerqyPlugin.java index d5bf8a4..7a818f7 100644 --- a/src/main/java/querqy/elasticsearch/QuerqyPlugin.java +++ b/src/main/java/querqy/elasticsearch/QuerqyPlugin.java @@ -1,9 +1,5 @@ package querqy.elasticsearch; -import static java.util.Arrays.asList; -import static java.util.Collections.unmodifiableList; -import static querqy.elasticsearch.rewriterstore.Constants.SETTINGS_QUERQY_INDEX_NUM_REPLICAS; - import org.elasticsearch.action.ActionRequest; import org.elasticsearch.action.ActionResponse; import org.elasticsearch.client.Client; @@ -16,7 +12,6 @@ import org.elasticsearch.common.settings.Setting; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.settings.SettingsFilter; -import org.elasticsearch.xcontent.NamedXContentRegistry; import org.elasticsearch.env.Environment; import org.elasticsearch.env.NodeEnvironment; import org.elasticsearch.index.IndexModule; @@ -30,25 +25,33 @@ import org.elasticsearch.script.ScriptService; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.watcher.ResourceWatcherService; +import org.elasticsearch.xcontent.NamedXContentRegistry; +import querqy.elasticsearch.aggregation.InternalDecorationAggregation; +import querqy.elasticsearch.aggregation.QuerqyDecorationAggregationBuilder; import querqy.elasticsearch.infologging.Log4jSink; import querqy.elasticsearch.query.QuerqyQueryBuilder; import querqy.elasticsearch.rewriterstore.DeleteRewriterAction; import querqy.elasticsearch.rewriterstore.NodesClearRewriterCacheAction; import querqy.elasticsearch.rewriterstore.NodesReloadRewriterAction; +import querqy.elasticsearch.rewriterstore.PutRewriterAction; import querqy.elasticsearch.rewriterstore.RestDeleteRewriterAction; import querqy.elasticsearch.rewriterstore.RestPutRewriterAction; -import querqy.elasticsearch.rewriterstore.PutRewriterAction; import querqy.elasticsearch.rewriterstore.TransportDeleteRewriterAction; import querqy.elasticsearch.rewriterstore.TransportNodesClearRewriterCacheAction; import querqy.elasticsearch.rewriterstore.TransportNodesReloadRewriterAction; import querqy.elasticsearch.rewriterstore.TransportPutRewriterAction; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.function.Supplier; +import static java.util.Arrays.asList; +import static java.util.Collections.unmodifiableList; +import static querqy.elasticsearch.rewriterstore.Constants.SETTINGS_QUERQY_INDEX_NUM_REPLICAS; + public class QuerqyPlugin extends Plugin implements SearchPlugin, ActionPlugin { @@ -111,8 +114,8 @@ public Collection createComponents(final Client client, final ClusterSer final NamedXContentRegistry xContentRegistry, final Environment environment, final NodeEnvironment nodeEnvironment, final NamedWriteableRegistry namedWriteableRegistry, - IndexNameExpressionResolver indexNameExpressionResolver, - Supplier repositoriesServiceSupplier) { + final IndexNameExpressionResolver indexNameExpressionResolver, + final Supplier repositoriesServiceSupplier) { return Arrays.asList(rewriterShardContexts, querqyProcessor); } @@ -122,4 +125,18 @@ public List> getSettings() { Setting.Property.NodeScope)); } + + @Override + public List getAggregations() { + final List r = new ArrayList<>(); + r.add( + new AggregationSpec( + QuerqyDecorationAggregationBuilder.NAME, + QuerqyDecorationAggregationBuilder::new, + QuerqyDecorationAggregationBuilder.PARSER + ).addResultReader(InternalDecorationAggregation::new) + ); + return r; + } + } diff --git a/src/main/java/querqy/elasticsearch/QuerqyProcessor.java b/src/main/java/querqy/elasticsearch/QuerqyProcessor.java index a81d94d..32573bd 100644 --- a/src/main/java/querqy/elasticsearch/QuerqyProcessor.java +++ b/src/main/java/querqy/elasticsearch/QuerqyProcessor.java @@ -4,6 +4,7 @@ import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.Query; import org.elasticsearch.index.query.SearchExecutionContext; +import querqy.elasticsearch.aggregation.DecoratedQuery; import querqy.elasticsearch.infologging.LogPayloadType; import querqy.elasticsearch.infologging.SingleSinkInfoLogging; import querqy.elasticsearch.query.InfoLoggingSpec; @@ -15,6 +16,7 @@ import querqy.lucene.LuceneSearchEngineRequestAdapter; import querqy.lucene.QueryParsingController; import querqy.rewrite.RewriteChain; +import querqy.rewrite.commonrules.model.DecorateInstruction; import java.util.Collections; import java.util.List; @@ -75,10 +77,6 @@ public Query parseQuery(final QuerqyQueryBuilder queryBuilder, final SearchExecu final QueryParsingController controller = new QueryParsingController(requestAdapter); final LuceneQueries queries = controller.process(); - -// // TODO: make decos part of the general Querqy object model -// final Set decorations = (Set) requestAdapter.getContext().get(DecorateInstruction.CONTEXT_KEY); - if ((queries.querqyBoostQueries == null || queries.querqyBoostQueries.isEmpty()) && (queries.filterQueries == null || queries.filterQueries.isEmpty()) && queries.mainQuery instanceof BooleanQuery) { @@ -104,15 +102,19 @@ public Query parseQuery(final QuerqyQueryBuilder queryBuilder, final SearchExecu appendFilterQueries(queries, builder); - final BooleanQuery query = builder.build(); + final Set decorations = (Set) requestAdapter.getContext().get(DecorateInstruction.DECORATION_CONTEXT_KEY); + final Query query = + decorations != null && !decorations.isEmpty() ? + new DecoratedQuery<>(builder.build(), decorations) : + builder.build(); if (infoLogging != null) { infoLogging.endOfRequest(requestAdapter); } + return query; } - void appendFilterQueries(final LuceneQueries queries, final BooleanQuery.Builder builder) { if (queries.filterQueries != null) { diff --git a/src/main/java/querqy/elasticsearch/aggregation/DecoratedQuery.java b/src/main/java/querqy/elasticsearch/aggregation/DecoratedQuery.java new file mode 100644 index 0000000..014c227 --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/DecoratedQuery.java @@ -0,0 +1,83 @@ +package querqy.elasticsearch.aggregation; + +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Weight; + +import java.io.IOException; +import java.util.Objects; +import java.util.Set; + +public class DecoratedQuery extends Query { + + final private T query; + final private Set decorations; + + public DecoratedQuery(final T query, final Set decorations) { + this.query = Objects.requireNonNull(query); + this.decorations = Objects.requireNonNull(decorations); + } + + public T getQuery() { + return query; + } + + public Set getDecorations() { + return decorations; + } + + @Override + public Weight createWeight(final IndexSearcher searcher, final ScoreMode scoreMode, final float boost) throws IOException { + return query.createWeight(searcher, scoreMode, boost); + } + + @Override + public Query rewrite(final IndexReader reader) throws IOException { + return query.rewrite(reader); + } + + @Override + public String toString(final String field) { + return query.toString(field); + } + + @Override + public boolean equals(final Object object) { + if (!sameClassAs(object)) return false; + final DecoratedQuery other = castObject(object); + return isEqualQueriesAndDecorations(other); + } + + private boolean isEqualQueriesAndDecorations(final DecoratedQuery other) { + final Query otherQuery = other.getQuery(); + final Set otherDecorations = other.getDecorations(); + return getQuery().equals(otherQuery) && getDecorations().equals(otherDecorations); + } + + private DecoratedQuery castObject(final Object object) { + return getClass().cast(object); + } + + private int computeHashCode() { + int hashCode = Objects.hash(query, decorations); + if (hashCode == 0) { + hashCode = 1; + } + return hashCode; + } + + // cached hash code is ok since boolean queries are immutable + private int hashCode; + + @Override + public int hashCode() { + // no need for synchronization, in the worst case we would just compute the hash several times. + if (hashCode == 0) { + hashCode = computeHashCode(); + } + return hashCode; + } + +} \ No newline at end of file diff --git a/src/main/java/querqy/elasticsearch/aggregation/DecorationAggregation.java b/src/main/java/querqy/elasticsearch/aggregation/DecorationAggregation.java new file mode 100644 index 0000000..e9c2f85 --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/DecorationAggregation.java @@ -0,0 +1,15 @@ +package querqy.elasticsearch.aggregation; + +import org.elasticsearch.search.aggregations.Aggregation; + +/** + * A {@code DecorationAggregation} aggregation. Defines a single bucket the holds all the querqy info in the search context. + */ +public interface DecorationAggregation extends Aggregation { + + /** + * The result of the aggregation. The type of the object depends on the aggregation that was run. + */ + Object aggregation(); + +} diff --git a/src/main/java/querqy/elasticsearch/aggregation/InternalDecorationAggregation.java b/src/main/java/querqy/elasticsearch/aggregation/InternalDecorationAggregation.java new file mode 100644 index 0000000..9b9535d --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/InternalDecorationAggregation.java @@ -0,0 +1,110 @@ +package querqy.elasticsearch.aggregation; + +import org.elasticsearch.Version; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.search.aggregations.InternalAggregation; +import org.elasticsearch.search.aggregations.metrics.ScriptedMetric; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import static java.util.Collections.singletonList; + +public class InternalDecorationAggregation extends InternalAggregation implements ScriptedMetric { + + private final List aggregations; + + InternalDecorationAggregation(final String name, final List aggregations, final Map metadata) { + super(name, metadata); + this.aggregations = aggregations; + } + + public InternalDecorationAggregation(final StreamInput in) throws IOException { + super(in); + aggregations = in.readList(StreamInput::readGenericValue); + } + + @Override + protected void doWriteTo(final StreamOutput out) throws IOException { + out.writeCollection(aggregations, StreamOutput::writeGenericValue); + } + + @Override + public String getWriteableName() { + return QuerqyDecorationAggregationBuilder.NAME; + } + + @Override + public Object aggregation() { + if (aggregations.size() != 1) { + throw new IllegalStateException("aggregation was not reduced"); + } + return aggregations.get(0); + } + + List aggregationsList() { + return aggregations; + } + + @Override + public InternalAggregation reduce(final List aggregations, final ReduceContext reduceContext) { + final List aggregationObjects = new ArrayList<>(); + for (final InternalAggregation aggregation : aggregations) { + final InternalDecorationAggregation mapReduceAggregation = (InternalDecorationAggregation) aggregation; + aggregationObjects.addAll(mapReduceAggregation.aggregations); + } + final InternalDecorationAggregation firstAggregation = ((InternalDecorationAggregation) aggregations.get(0)); + final List aggregation; + if (reduceContext.isFinalReduce()) { + aggregation = Collections.singletonList(aggregationObjects); + } else { + // if we are not an final reduce we have to maintain all the aggs from all the incoming one + // until we hit the final reduce phase. + aggregation = aggregationObjects; + } + return new InternalDecorationAggregation(firstAggregation.getName(), aggregation, getMetadata()); + } + + @Override + protected boolean mustReduceOnSingleInternalAgg() { + return true; + } + + @Override + public Object getProperty(final List path) { + if (path.isEmpty()) { + return this; + } else if (path.size() == 1 && "value".equals(path.get(0))) { + return aggregation(); + } else { + throw new IllegalArgumentException("path not supported for [" + getName() + "]: " + path); + } + } + + @Override + public XContentBuilder doXContentBody(final XContentBuilder builder, final Params params) throws IOException { + return builder.field(CommonFields.VALUE.getPreferredName(), aggregation()); + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) return true; + if (obj == null || getClass() != obj.getClass()) return false; + if (!super.equals(obj)) return false; + + final InternalDecorationAggregation other = (InternalDecorationAggregation) obj; + return Objects.equals(aggregations, other.aggregations); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), aggregations); + } + +} diff --git a/src/main/java/querqy/elasticsearch/aggregation/ParsedDecorationAggregation.java b/src/main/java/querqy/elasticsearch/aggregation/ParsedDecorationAggregation.java new file mode 100644 index 0000000..3b08947 --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/ParsedDecorationAggregation.java @@ -0,0 +1,80 @@ +package querqy.elasticsearch.aggregation; + +import org.elasticsearch.common.bytes.BytesArray; +import org.elasticsearch.search.aggregations.ParsedAggregation; +import org.elasticsearch.xcontent.ObjectParser; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; + +public class ParsedDecorationAggregation extends ParsedAggregation implements DecorationAggregation { + + private List aggregation; + + @Override + public String getType() { + return QuerqyDecorationAggregationBuilder.NAME; + } + + @Override + public Object aggregation() { + assert aggregation.size() == 1; + return aggregation.get(0); + } + + @Override + public XContentBuilder doXContentBody(final XContentBuilder builder, final Params params) throws IOException { + return builder.field(CommonFields.VALUE.getPreferredName(), aggregation()); + } + + private static final ObjectParser PARSER = new ObjectParser<>( + ParsedDecorationAggregation.class.getSimpleName(), + true, + ParsedDecorationAggregation::new + ); + + static { + declareAggregationFields(PARSER); + PARSER.declareField( + (agg, value) -> agg.aggregation = Collections.singletonList(value), + ParsedDecorationAggregation::parseValue, + CommonFields.VALUE, + ObjectParser.ValueType.VALUE_OBJECT_ARRAY + ); + } + + private static Object parseValue(final XContentParser parser) throws IOException { + final XContentParser.Token token = parser.currentToken(); + Object value = null; + if (token != XContentParser.Token.VALUE_NULL) { + if (token.isValue()) { + if (token == XContentParser.Token.VALUE_STRING) { + // binary values will be parsed back and returned as base64 strings when reading from json and yaml + value = parser.text(); + } else if (token == XContentParser.Token.VALUE_NUMBER) { + value = parser.numberValue(); + } else if (token == XContentParser.Token.VALUE_BOOLEAN) { + value = parser.booleanValue(); + } else if (token == XContentParser.Token.VALUE_EMBEDDED_OBJECT) { + // binary values will be parsed back and returned as BytesArray when reading from cbor and smile + value = new BytesArray(parser.binaryValue()); + } + } else if (token == XContentParser.Token.START_OBJECT) { + value = parser.map(); + } else if (token == XContentParser.Token.START_ARRAY) { + value = parser.list(); + } + } + return value; + } + + public static ParsedDecorationAggregation fromXContent(final XContentParser parser, final String name) { + final ParsedDecorationAggregation aggregation = PARSER.apply(parser, null); + aggregation.setName(name); + return aggregation; + } + +} diff --git a/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationBuilder.java b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationBuilder.java new file mode 100644 index 0000000..028a34e --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationBuilder.java @@ -0,0 +1,79 @@ +package querqy.elasticsearch.aggregation; + +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.search.aggregations.AbstractAggregationBuilder; +import org.elasticsearch.search.aggregations.AggregationBuilder; +import org.elasticsearch.search.aggregations.AggregatorFactories.Builder; +import org.elasticsearch.search.aggregations.AggregatorFactory; +import org.elasticsearch.search.aggregations.support.AggregationContext; +import org.elasticsearch.xcontent.ObjectParser; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Map; +import java.util.Objects; + +public class QuerqyDecorationAggregationBuilder extends AbstractAggregationBuilder { + + public static final String NAME = "decorations"; + + public static final ObjectParser PARSER = + new ObjectParser<>(NAME, QuerqyDecorationAggregationBuilder::new); + + public QuerqyDecorationAggregationBuilder() { + super(NAME); + } + + public QuerqyDecorationAggregationBuilder(final StreamInput in) throws IOException { + super(in); + } + + protected QuerqyDecorationAggregationBuilder(final QuerqyDecorationAggregationBuilder clone, final Builder factoriesBuilder, final Map metadata) { + super(clone, factoriesBuilder, metadata); + } + + @Override + protected AggregationBuilder shallowCopy(final Builder factoriesBuilder, final Map metadata) { + return new QuerqyDecorationAggregationBuilder(this, factoriesBuilder, metadata); + } + + @Override + public BucketCardinality bucketCardinality() { + return BucketCardinality.NONE; + } + + @Override + protected AggregatorFactory doBuild(final AggregationContext context, final AggregatorFactory parent, final Builder subFactoriesBuilder) + throws IOException { + return new QuerqyDecorationAggregatorFactory(name, context, parent, subFactoriesBuilder, metadata); + } + + @Override + protected XContentBuilder internalXContent(final XContentBuilder builder, final Params params) { + return builder; + } + + @Override + protected void doWriteTo(final StreamOutput out) { + // no state to write out + } + + @Override + public String getType() { + return NAME; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode()); + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) return true; + if (obj == null || getClass() != obj.getClass()) return false; + return super.equals(obj); + } + +} diff --git a/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregator.java b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregator.java new file mode 100644 index 0000000..68f4c4f --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregator.java @@ -0,0 +1,54 @@ +package querqy.elasticsearch.aggregation; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.ScoreMode; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.search.aggregations.InternalAggregation; +import org.elasticsearch.search.aggregations.LeafBucketCollector; +import org.elasticsearch.search.aggregations.metrics.MetricsAggregator; +import org.elasticsearch.search.aggregations.support.AggregationContext; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Map; +import java.util.Set; + +public final class QuerqyDecorationAggregator extends MetricsAggregator { + + final Set decorations; + + public QuerqyDecorationAggregator(final String name, final AggregationContext context, final Map metadata, final Set decorations) + throws IOException { + super(name, context, null, metadata); + this.decorations = decorations; + } + + @Override + public ScoreMode scoreMode() { + return ScoreMode.COMPLETE; + } + + @Override + protected LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final LeafBucketCollector sub) { + // No sub-aggregations + return LeafBucketCollector.NO_OP_COLLECTOR; + } + + @Override + public InternalAggregation buildAggregation(final long l) { + StreamOutput.checkWriteable(decorations); + return new InternalDecorationAggregation( + name, + new ArrayList<>(decorations), + metadata() + ); + } + + @Override + public InternalAggregation buildEmptyAggregation() { + throw new UnsupportedOperationException( + "querqy_decoration aggregations cannot serve as sub-aggregations, hence should never be called on #buildEmptyAggregations" + ); + } + +} diff --git a/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregatorFactory.java b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregatorFactory.java new file mode 100644 index 0000000..13e78d5 --- /dev/null +++ b/src/main/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregatorFactory.java @@ -0,0 +1,73 @@ +package querqy.elasticsearch.aggregation; + +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.elasticsearch.search.aggregations.AggregationExecutionException; +import org.elasticsearch.search.aggregations.Aggregator; +import org.elasticsearch.search.aggregations.AggregatorFactories; +import org.elasticsearch.search.aggregations.AggregatorFactory; +import org.elasticsearch.search.aggregations.CardinalityUpperBound; +import org.elasticsearch.search.aggregations.support.AggregationContext; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +public class QuerqyDecorationAggregatorFactory extends AggregatorFactory { + + public QuerqyDecorationAggregatorFactory( + final String name, + final AggregationContext context, + final AggregatorFactory parent, + final AggregatorFactories.Builder subFactories, + final Map metadata + ) throws IOException { + super(name, context, parent, subFactories, metadata); + } + + @Override + public Aggregator createInternal(final Aggregator parent, final CardinalityUpperBound cardinality, final Map metadata) + throws IOException { + if (parent != null) { + throw new IllegalArgumentException( + "Aggregation [" + + parent.name() + + "] cannot have a querqy_decoration " + + "sub-aggregation [" + + name + + "]. querqy_decoration aggregations can only be defined as top level aggregations" + ); + } + if (cardinality != CardinalityUpperBound.ONE) { + throw new AggregationExecutionException("Aggregation [" + name() + "] must have cardinality 1 but was [" + cardinality + "]"); + } + final Query query = context.subSearchContext() == null ? null : context.subSearchContext().query(); + final Set> decoratedQueries = getDecoratedQueries(query); + return new QuerqyDecorationAggregator(name, context, metadata, collectAllDecorations(decoratedQueries)); + } + + + private Set collectAllDecorations(final Set> decoratedQueries) { + return decoratedQueries.stream().flatMap(decoratedQuery -> decoratedQuery.getDecorations().stream()).collect(Collectors.toSet()); + } + + private Set> getDecoratedQueries(final Query query) { + if (query == null) { + return Collections.emptySet(); + } + final Set> decoratedQueries = new HashSet<>(); + query.visit(new QueryVisitor() { + @Override + public void visitLeaf(final Query query) { + if (query instanceof DecoratedQuery) { + decoratedQueries.add((DecoratedQuery) query); + } + } + }); + return decoratedQueries; + } + +} diff --git a/src/test/java/querqy/elasticsearch/RewriterIntegrationTest.java b/src/test/java/querqy/elasticsearch/RewriterIntegrationTest.java index 5394c34..14ee20e 100644 --- a/src/test/java/querqy/elasticsearch/RewriterIntegrationTest.java +++ b/src/test/java/querqy/elasticsearch/RewriterIntegrationTest.java @@ -1,15 +1,5 @@ package querqy.elasticsearch; -import java.io.BufferedReader; -import java.io.InputStreamReader; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; -import java.util.concurrent.ExecutionException; -import java.util.stream.Collectors; - import org.elasticsearch.action.search.SearchRequestBuilder; import org.elasticsearch.action.search.SearchResponse; import org.elasticsearch.action.support.WriteRequest; @@ -27,6 +17,16 @@ import querqy.elasticsearch.rewriterstore.PutRewriterRequest; import querqy.elasticsearch.rewriterstore.PutRewriterResponse; +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; + public class RewriterIntegrationTest extends ESSingleNodeTestCase { private final String INDEX_NAME = "test_index"; diff --git a/src/test/java/querqy/elasticsearch/aggregation/InternalDecorationAggregationTests.java b/src/test/java/querqy/elasticsearch/aggregation/InternalDecorationAggregationTests.java new file mode 100644 index 0000000..c84f636 --- /dev/null +++ b/src/test/java/querqy/elasticsearch/aggregation/InternalDecorationAggregationTests.java @@ -0,0 +1,242 @@ +package querqy.elasticsearch.aggregation; + +import org.elasticsearch.Version; +import org.elasticsearch.common.geo.GeoPoint; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.plugins.SearchPlugin; +import org.elasticsearch.search.aggregations.Aggregation; +import org.elasticsearch.search.aggregations.Aggregation.CommonFields; +import org.elasticsearch.search.aggregations.InternalAggregation; +import org.elasticsearch.search.aggregations.InternalAggregation.ReduceContext; +import org.elasticsearch.search.aggregations.ParsedAggregation; +import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree; +import org.elasticsearch.test.InternalAggregationTestCase; +import org.elasticsearch.test.VersionUtils; +import org.elasticsearch.xcontent.ContextParser; +import org.elasticsearch.xcontent.NamedXContentRegistry; +import org.elasticsearch.xcontent.ParseField; +import querqy.elasticsearch.QuerqyPlugin; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.function.Predicate; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +import static java.util.Collections.singletonList; +import static org.hamcrest.Matchers.equalTo; + +public class InternalDecorationAggregationTests extends InternalAggregationTestCase { + + private Supplier[] valueTypes; + @SuppressWarnings({ "rawtypes", "unchecked" }) + private final Supplier[] leafValueSuppliers = new Supplier[] { + () -> randomInt(), + () -> randomLong(), + () -> randomDouble(), + () -> randomFloat(), + () -> randomBoolean(), + () -> randomAlphaOfLength(5), + () -> new GeoPoint(randomDouble(), randomDouble()), + () -> null }; + @SuppressWarnings({ "rawtypes", "unchecked" }) + private final Supplier[] nestedValueSuppliers = new Supplier[] { () -> new HashMap(), () -> new ArrayList<>() }; + + private static final List namedXContents = getDefaultNamedXContents(); + static { + Map> map = new HashMap<>(); + map.put(QuerqyDecorationAggregationBuilder.NAME, (p, c) -> ParsedDecorationAggregation.fromXContent(p, (String) c)); + + List namedXContentsToAdd = map.entrySet() + .stream() + .map(entry -> new NamedXContentRegistry.Entry(Aggregation.class, new ParseField(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()); + namedXContents.addAll(namedXContentsToAdd); + } + + @Override + @SuppressWarnings({ "rawtypes", "unchecked" }) + public void setUp() throws Exception { + super.setUp(); + // we want the same value types (also for nested lists, maps) for all random aggregations + int levels = randomIntBetween(1, 3); + valueTypes = new Supplier[levels]; + for (int i = 0; i < levels; i++) { + if (i < levels - 1) { + valueTypes[i] = randomFrom(nestedValueSuppliers); + } else { + // the last one needs to be a leaf value, not map or list + valueTypes[i] = randomFrom(leafValueSuppliers); + } + } + } + + @Override + protected SearchPlugin registerPlugin() { + return new QuerqyPlugin(Settings.EMPTY); + } + + @Override + protected InternalDecorationAggregation createTestInstance(String name, Map metadata) { + Map params = new HashMap<>(); + if (randomBoolean()) { + params.put(randomAlphaOfLength(5), randomAlphaOfLength(5)); + } + return new InternalDecorationAggregation(name, randomAggregations(), metadata); + } + + private List randomAggregations() { + return randomList(randomBoolean() ? 1 : 5, this::randomAggregation); + } + + @SuppressWarnings({ "rawtypes", "unchecked" }) + private Object randomAggregation() { + int levels = randomIntBetween(1, 3); + Supplier[] valueTypes = new Supplier[levels]; + for (int l = 0; l < levels; l++) { + if (l < levels - 1) { + valueTypes[l] = randomFrom(nestedValueSuppliers); + } else { + // the last one needs to be a leaf value, not map or + // list + valueTypes[l] = randomFrom(leafValueSuppliers); + } + } + return randomValue(valueTypes, 0); + } + + @SuppressWarnings("unchecked") + private static Object randomValue(Supplier[] valueTypes, int level) { + Object value = valueTypes[level].get(); + if (value instanceof Map) { + int elements = randomIntBetween(1, 5); + Map map = (Map) value; + for (int i = 0; i < elements; i++) { + map.put(randomAlphaOfLength(5), randomValue(valueTypes, level + 1)); + } + } else if (value instanceof List) { + int elements = randomIntBetween(1, 5); + List list = (List) value; + for (int i = 0; i < elements; i++) { + list.add(randomValue(valueTypes, level + 1)); + } + } + return value; + } + + @Override + protected void assertReduced(InternalDecorationAggregation reduced, List inputs) { + InternalDecorationAggregation firstAgg = inputs.get(0); + assertEquals(firstAgg.getName(), reduced.getName()); + assertEquals(firstAgg.getMetadata(), reduced.getMetadata()); + int size = (int) inputs.stream().mapToLong(i -> i.aggregationsList().size()).sum(); + assertEquals(size, ((List) reduced.aggregation()).size()); + } + + @Override + public InternalDecorationAggregation createTestInstanceForXContent() { + InternalDecorationAggregation aggregation = createTestInstance(); + return (InternalDecorationAggregation) aggregation.reduce( + singletonList(aggregation), + ReduceContext.forFinalReduction(null, mockScriptService(), null, PipelineTree.EMPTY, () -> false) + ); + } + + @Override + protected void assertFromXContent(InternalDecorationAggregation aggregation, ParsedAggregation parsedAggregation) throws IOException { + assertTrue(parsedAggregation instanceof ParsedDecorationAggregation); + ParsedDecorationAggregation parsed = (ParsedDecorationAggregation) parsedAggregation; + + assertValues(aggregation.aggregation(), parsed.aggregation()); + } + + private static void assertValues(Object expected, Object actual) { + if (expected instanceof Long) { + // longs that fit into the integer range are parsed back as integer + if (actual instanceof Integer) { + assertEquals(((Long) expected).intValue(), actual); + } else { + assertEquals(expected, actual); + } + } else if (expected instanceof Float) { + // based on the xContent type, floats are sometimes parsed back as doubles + if (actual instanceof Double) { + assertEquals(expected, ((Double) actual).floatValue()); + } else { + assertEquals(expected, actual); + } + } else if (expected instanceof GeoPoint) { + assertTrue(actual instanceof Map); + GeoPoint point = (GeoPoint) expected; + @SuppressWarnings("unchecked") + Map pointMap = (Map) actual; + assertEquals(point.getLat(), pointMap.get("lat")); + assertEquals(point.getLon(), pointMap.get("lon")); + } else if (expected instanceof Map) { + @SuppressWarnings("unchecked") + Map expectedMap = (Map) expected; + @SuppressWarnings("unchecked") + Map actualMap = (Map) actual; + assertEquals(expectedMap.size(), actualMap.size()); + for (String key : expectedMap.keySet()) { + assertValues(expectedMap.get(key), actualMap.get(key)); + } + } else if (expected instanceof List) { + @SuppressWarnings("unchecked") + List expectedList = (List) expected; + @SuppressWarnings("unchecked") + List actualList = (List) actual; + assertEquals(expectedList.size(), actualList.size()); + Iterator actualIterator = actualList.iterator(); + for (Object element : expectedList) { + assertValues(element, actualIterator.next()); + } + } else { + assertEquals(expected, actual); + } + } + + @Override + protected Predicate excludePathsFromXContentInsertion() { + return path -> path.contains(CommonFields.VALUE.getPreferredName()); + } + + public void testOldSerialization() throws IOException { + // A single element list looks like a fully reduced agg + InternalDecorationAggregation original = new InternalDecorationAggregation( + "test", + org.elasticsearch.core.List.of("foo"), + null + ); + original.mergePipelineTreeForBWCSerialization(PipelineTree.EMPTY); + InternalDecorationAggregation roundTripped = (InternalDecorationAggregation) copyNamedWriteable( + original, + getNamedWriteableRegistry(), + InternalAggregation.class, + VersionUtils.randomVersionBetween(random(), Version.V_7_0_0, VersionUtils.getPreviousVersion(Version.V_7_8_0)) + ); + assertThat(roundTripped, equalTo(original)); + + // A multi-element list looks like a non-reduced agg + InternalDecorationAggregation unreduced = new InternalDecorationAggregation( + "test", + org.elasticsearch.core.List.of("foo", "bar"), + null + ); + unreduced.mergePipelineTreeForBWCSerialization(PipelineTree.EMPTY); + Exception e = expectThrows( + IllegalArgumentException.class, + () -> copyNamedWriteable( + unreduced, + getNamedWriteableRegistry(), + InternalAggregation.class, + VersionUtils.randomVersionBetween(random(), Version.V_7_0_0, VersionUtils.getPreviousVersion(Version.V_7_8_0)) + ) + ); + assertThat(e.getMessage(), equalTo("querqy doesn't support cross cluster search until 7.8.0")); + } +} diff --git a/src/test/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationIntegrationTest.java b/src/test/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationIntegrationTest.java new file mode 100644 index 0000000..a14a01c --- /dev/null +++ b/src/test/java/querqy/elasticsearch/aggregation/QuerqyDecorationAggregationIntegrationTest.java @@ -0,0 +1,147 @@ +package querqy.elasticsearch.aggregation; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.search.SearchRequestBuilder; +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.action.support.WriteRequest; +import org.elasticsearch.common.bytes.BytesArray; +import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.ConstantScoreQueryBuilder; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.test.ESSingleNodeTestCase; +import org.elasticsearch.xcontent.XContentType; +import org.junit.After; +import querqy.elasticsearch.QuerqyPlugin; +import querqy.elasticsearch.QuerqyProcessor; +import querqy.elasticsearch.query.MatchingQuery; +import querqy.elasticsearch.query.QuerqyQueryBuilder; +import querqy.elasticsearch.query.Rewriter; +import querqy.elasticsearch.rewriterstore.PutRewriterAction; +import querqy.elasticsearch.rewriterstore.PutRewriterRequest; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutionException; + +public class QuerqyDecorationAggregationIntegrationTest extends ESSingleNodeTestCase { + + public static String DECORATIONS_TEMPLATE = "{\"decorations\":{\"value\":[%s]}}"; + private final static Logger LOGGER = LogManager.getLogger(QuerqyDecorationAggregationIntegrationTest.class); + + private final String INDEX_NAME = "test_index"; + + @Override + protected Collection> getPlugins() { + return Collections.singleton(QuerqyPlugin.class); + } + + public void testSearchWithConfig() throws Exception { + + index(); + + final Map content = new HashMap<>(); + content.put("class", querqy.elasticsearch.rewriter.SimpleCommonRulesRewriterFactory.class.getName()); + + final Map config = new HashMap<>(); + config.put("rules", "k =>\nSYNONYM: c\na =>\nDECORATE: REDIRECT /faq/a\ny =>\nDECORATE: REDIRECT /faq/y"); + config.put("ignoreCase", true); + config.put("querqyParser", querqy.rewrite.commonrules.WhiteSpaceQuerqyParserFactory.class.getName()); + content.put("config", config); + + final PutRewriterRequest request = new PutRewriterRequest("common_rules", content); + + client().execute(PutRewriterAction.INSTANCE, request).get(); + + QuerqyProcessor querqyProcessor = getInstanceFromNode(QuerqyProcessor.class); + + SearchRequestBuilder searchRequestBuilder = client().prepareSearch(INDEX_NAME); + + QuerqyQueryBuilder querqyQuery = new QuerqyQueryBuilder(querqyProcessor); + querqyQuery.setRewriters(Collections.singletonList(new Rewriter("common_rules"))); + querqyQuery.setQueryFieldsAndBoostings(Arrays.asList("field1", "field2")); + querqyQuery.setMinimumShouldMatch("1"); + + QuerqyDecorationAggregationBuilder aggregationBuilder = new QuerqyDecorationAggregationBuilder(); + String[] expectedDecoration = new String[] {"REDIRECT /faq/a","REDIRECT /faq/y"}; + + // with decorations + querqyQuery.setMatchingQuery(new MatchingQuery("a k")); + searchRequestBuilder.setQuery(querqyQuery); + searchRequestBuilder.addAggregation(aggregationBuilder); + testSearchRequest(searchRequestBuilder, 2L, Collections.singleton(expectedDecoration[0])); + + // without hits, without decorations + querqyQuery.setMatchingQuery(new MatchingQuery("x z")); + searchRequestBuilder.setQuery(querqyQuery); + testSearchRequest(searchRequestBuilder, 0L, Collections.emptySet()); + + // without hits, with decorations + querqyQuery.setMatchingQuery(new MatchingQuery("x y")); + searchRequestBuilder.setQuery(querqyQuery); + testSearchRequest(searchRequestBuilder, 0L, Collections.singleton(expectedDecoration[1])); + + // with hits, without decorations + querqyQuery.setMatchingQuery(new MatchingQuery("k x")); + searchRequestBuilder.setQuery(querqyQuery); + testSearchRequest(searchRequestBuilder, 2L, Collections.emptySet()); + + // inner boolean query + querqyQuery.setMatchingQuery(new MatchingQuery("a k")); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should(querqyQuery); + searchRequestBuilder.setQuery(boolQueryBuilder); + testSearchRequest(searchRequestBuilder, 2L, Collections.singleton(expectedDecoration[0])); + + // inner constant score && inner bool query + querqyQuery.setMatchingQuery(new MatchingQuery("a k")); + + QuerqyQueryBuilder querqyQuery2 = new QuerqyQueryBuilder(querqyProcessor); + querqyQuery2.setRewriters(Collections.singletonList(new Rewriter("common_rules"))); + querqyQuery2.setQueryFieldsAndBoostings(Arrays.asList("field1", "field2")); + querqyQuery2.setMinimumShouldMatch("1"); + querqyQuery2.setMatchingQuery(new MatchingQuery("x y")); + + ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(querqyQuery2); + boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should(querqyQuery); + boolQueryBuilder.should(constantScoreQueryBuilder); + searchRequestBuilder.setQuery(boolQueryBuilder); + testSearchRequest(searchRequestBuilder, 2L, new HashSet<>(Arrays.asList(expectedDecoration))); + + } + + private void testSearchRequest(SearchRequestBuilder searchRequestBuilder, long expectedHits, Set expectedDecorations) throws ExecutionException, InterruptedException { + SearchResponse response = client().search(searchRequestBuilder.request()).get(); + LOGGER.info("Response:\n{}", response); + assertEquals(expectedHits, response.getHits().getTotalHits().value); + InternalDecorationAggregation aggregation = response.getAggregations().get(QuerqyDecorationAggregationBuilder.NAME); + assertEquals( + expectedDecorations, + new HashSet<>((List) aggregation.aggregation()) + ); + } + + @After + public void deleteRewriterIndex() { + client().admin().indices().prepareDelete(".querqy").get(); + } + + public void index() { + client().admin().indices().prepareCreate(INDEX_NAME).get(); + client().prepareIndex(INDEX_NAME, null) + .setSource("field1", "a b", "field2", "a c") + .get(); + client().prepareIndex(INDEX_NAME, null) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .setSource("field1", "b c") + .get(); + } +}