Skip to content

Commit

Permalink
Replace the SearchContext with QueryShardContext when building aggreg…
Browse files Browse the repository at this point in the history 
…ator factories (#46527)

This commit replaces the `SearchContext` with the `QueryShardContext` when building aggregator factories. Aggregator factories are part of the `SearchContext` so they shouldn't require a `SearchContext` to create them.
The main changes here are the signatures of `AggregationBuilder#build` that now takes a `QueryShardContext` and `AggregatorFactory#createInternal` that passes the `SearchContext` to build the `Aggregator`.

Relates #46523
  • Loading branch information
@jimczi
jimczi authored Sep 11, 2019
1 parent 4207fcf commit 5eeb3c5
Show file tree
Hide file tree
Showing 130 changed files with 1,294 additions and 911 deletions.
10 changes: 6 additions & 4 deletions ...ava/org/elasticsearch/search/aggregations/matrix/stats/MatrixStatsAggregationBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.MultiValueMode;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregatorFactories;
Expand All @@ -32,7 +33,6 @@
import org.elasticsearch.search.aggregations.support.ValuesSource.Numeric;
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.Map;
Expand Down Expand Up @@ -80,9 +80,11 @@ public MultiValueMode multiValueMode() {
}

@Override
protected MatrixStatsAggregatorFactory innerBuild(SearchContext context, Map<String, ValuesSourceConfig<Numeric>> configs,
AggregatorFactory parent, AggregatorFactories.Builder subFactoriesBuilder) throws IOException {
return new MatrixStatsAggregatorFactory(name, configs, multiValueMode, context, parent, subFactoriesBuilder, metaData);
protected MatrixStatsAggregatorFactory innerBuild(QueryShardContext queryShardContext,
Map<String, ValuesSourceConfig<Numeric>> configs,
AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder) throws IOException {
return new MatrixStatsAggregatorFactory(name, configs, multiValueMode, queryShardContext, parent, subFactoriesBuilder, metaData);
}

@Override
Expand Down
30 changes: 20 additions & 10 deletions ...java/org/elasticsearch/search/aggregations/matrix/stats/MatrixStatsAggregatorFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
*/
package org.elasticsearch.search.aggregations.matrix.stats;

import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.MultiValueMode;
import org.elasticsearch.search.aggregations.Aggregator;
import org.elasticsearch.search.aggregations.AggregatorFactories;
Expand All @@ -37,23 +38,32 @@ final class MatrixStatsAggregatorFactory extends ArrayValuesSourceAggregatorFact
private final MultiValueMode multiValueMode;

MatrixStatsAggregatorFactory(String name,
Map<String, ValuesSourceConfig<ValuesSource.Numeric>> configs, MultiValueMode multiValueMode,
SearchContext context, AggregatorFactory parent, AggregatorFactories.Builder subFactoriesBuilder,
Map<String, Object> metaData) throws IOException {
super(name, configs, context, parent, subFactoriesBuilder, metaData);
Map<String, ValuesSourceConfig<ValuesSource.Numeric>> configs,
MultiValueMode multiValueMode,
QueryShardContext queryShardContext,
AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder,
Map<String, Object> metaData) throws IOException {
super(name, configs, queryShardContext, parent, subFactoriesBuilder, metaData);
this.multiValueMode = multiValueMode;
}

@Override
protected Aggregator createUnmapped(Aggregator parent, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData)
protected Aggregator createUnmapped(SearchContext searchContext,
Aggregator parent,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData)
throws IOException {
return new MatrixStatsAggregator(name, null, context, parent, multiValueMode, pipelineAggregators, metaData);
return new MatrixStatsAggregator(name, null, searchContext, parent, multiValueMode, pipelineAggregators, metaData);
}

@Override
protected Aggregator doCreateInternal(Map<String, ValuesSource.Numeric> valuesSources, Aggregator parent,
boolean collectsFromSingleBucket, List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException {
return new MatrixStatsAggregator(name, valuesSources, context, parent, multiValueMode, pipelineAggregators, metaData);
protected Aggregator doCreateInternal(Map<String, ValuesSource.Numeric> valuesSources,
SearchContext searchContext,
Aggregator parent,
boolean collectsFromSingleBucket,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException {
return new MatrixStatsAggregator(name, valuesSources, searchContext, parent, multiValueMode, pipelineAggregators, metaData);
}
}
22 changes: 11 additions & 11 deletions ...va/org/elasticsearch/search/aggregations/support/ArrayValuesSourceAggregationBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,14 @@
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationInitializationException;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
import org.elasticsearch.search.aggregations.AggregatorFactory;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.ArrayList;
Expand Down Expand Up @@ -239,28 +239,28 @@ public Map<String, Object> missingMap() {
}

@Override
protected final ArrayValuesSourceAggregatorFactory<VS> doBuild(SearchContext context, AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder) throws IOException {
Map<String, ValuesSourceConfig<VS>> configs = resolveConfig(context);
ArrayValuesSourceAggregatorFactory<VS> factory = innerBuild(context, configs, parent, subFactoriesBuilder);
protected final ArrayValuesSourceAggregatorFactory<VS> doBuild(QueryShardContext queryShardContext, AggregatorFactory parent,
Builder subFactoriesBuilder) throws IOException {
Map<String, ValuesSourceConfig<VS>> configs = resolveConfig(queryShardContext);
ArrayValuesSourceAggregatorFactory<VS> factory = innerBuild(queryShardContext, configs, parent, subFactoriesBuilder);
return factory;
}

protected Map<String, ValuesSourceConfig<VS>> resolveConfig(SearchContext context) {
protected Map<String, ValuesSourceConfig<VS>> resolveConfig(QueryShardContext queryShardContext) {
HashMap<String, ValuesSourceConfig<VS>> configs = new HashMap<>();
for (String field : fields) {
ValuesSourceConfig<VS> config = config(context, field, null);
ValuesSourceConfig<VS> config = config(queryShardContext, field, null);
configs.put(field, config);
}
return configs;
}

protected abstract ArrayValuesSourceAggregatorFactory<VS> innerBuild(SearchContext context,
protected abstract ArrayValuesSourceAggregatorFactory<VS> innerBuild(QueryShardContext queryShardContext,
Map<String, ValuesSourceConfig<VS>> configs,
AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder) throws IOException;

public ValuesSourceConfig<VS> config(SearchContext context, String field, Script script) {
public ValuesSourceConfig<VS> config(QueryShardContext queryShardContext, String field, Script script) {

ValueType valueType = this.valueType != null ? this.valueType : targetValueType;

Expand All @@ -282,7 +282,7 @@ public ValuesSourceConfig<VS> config(SearchContext context, String field, Script
return config.format(resolveFormat(format, valueType));
}

MappedFieldType fieldType = context.smartNameFieldType(field);
MappedFieldType fieldType = queryShardContext.getMapperService().fullName(field);
if (fieldType == null) {
ValuesSourceType valuesSourceType = valueType != null ? valueType.getValuesSourceType() : this.valuesSourceType;
ValuesSourceConfig<VS> config = new ValuesSourceConfig<>(valuesSourceType);
Expand All @@ -291,7 +291,7 @@ public ValuesSourceConfig<VS> config(SearchContext context, String field, Script
return config.unmapped(true);
}

IndexFieldData<?> indexFieldData = context.getForField(fieldType);
IndexFieldData<?> indexFieldData = queryShardContext.getForField(fieldType);

ValuesSourceConfig<VS> config;
if (valuesSourceType == ValuesSourceType.ANY) {
Expand Down
34 changes: 22 additions & 12 deletions ...ava/org/elasticsearch/search/aggregations/support/ArrayValuesSourceAggregatorFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

package org.elasticsearch.search.aggregations.support;

import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.aggregations.Aggregator;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.AggregatorFactory;
Expand All @@ -36,35 +37,44 @@ public abstract class ArrayValuesSourceAggregatorFactory<VS extends ValuesSource
protected Map<String, ValuesSourceConfig<VS>> configs;

public ArrayValuesSourceAggregatorFactory(String name, Map<String, ValuesSourceConfig<VS>> configs,
SearchContext context, AggregatorFactory parent,
QueryShardContext queryShardContext, AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder,
Map<String, Object> metaData) throws IOException {
super(name, context, parent, subFactoriesBuilder, metaData);
super(name, queryShardContext, parent, subFactoriesBuilder, metaData);
this.configs = configs;
}

@Override
public Aggregator createInternal(Aggregator parent, boolean collectsFromSingleBucket, List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException {
public Aggregator createInternal(SearchContext searchContext,
Aggregator parent,
boolean collectsFromSingleBucket,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException {
HashMap<String, VS> valuesSources = new HashMap<>();

for (Map.Entry<String, ValuesSourceConfig<VS>> config : configs.entrySet()) {
VS vs = config.getValue().toValuesSource(context.getQueryShardContext());
VS vs = config.getValue().toValuesSource(queryShardContext);
if (vs != null) {
valuesSources.put(config.getKey(), vs);
}
}
if (valuesSources.isEmpty()) {
return createUnmapped(parent, pipelineAggregators, metaData);
return createUnmapped(searchContext, parent, pipelineAggregators, metaData);
}
return doCreateInternal(valuesSources, parent, collectsFromSingleBucket, pipelineAggregators, metaData);
return doCreateInternal(valuesSources, searchContext, parent,
collectsFromSingleBucket, pipelineAggregators, metaData);
}

protected abstract Aggregator createUnmapped(Aggregator parent, List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException;
protected abstract Aggregator createUnmapped(SearchContext searchContext,
Aggregator parent,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException;

protected abstract Aggregator doCreateInternal(Map<String, VS> valuesSources, Aggregator parent, boolean collectsFromSingleBucket,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException;
protected abstract Aggregator doCreateInternal(Map<String, VS> valuesSources,
SearchContext searchContext,
Aggregator parent,
boolean collectsFromSingleBucket,
List<PipelineAggregator> pipelineAggregators,
Map<String, Object> metaData) throws IOException;

}
2 changes: 1 addition & 1 deletion .../lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -551,9 +551,9 @@ private static Response prepareRamIndex(Request request,
ParsedDocument parsedDocument = indexService.mapperService().documentMapper().parse(sourceToParse);
indexWriter.addDocuments(parsedDocument.docs());
try (IndexReader indexReader = DirectoryReader.open(indexWriter)) {
final long absoluteStartMillis = System.currentTimeMillis();
final IndexSearcher searcher = new IndexSearcher(indexReader);
searcher.setQueryCache(null);
final long absoluteStartMillis = System.currentTimeMillis();
QueryShardContext context =
indexService.newQueryShardContext(0, searcher, () -> absoluteStartMillis, null);
return handler.apply(context, indexReader.leaves().get(0));
Expand Down
22 changes: 11 additions & 11 deletions ...nt-join/src/main/java/org/elasticsearch/join/aggregations/ChildrenAggregationBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.fielddata.plain.SortedSetDVOrdinalsIndexFieldData;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.join.mapper.ParentIdFieldMapper;
import org.elasticsearch.join.mapper.ParentJoinFieldMapper;
import org.elasticsearch.search.aggregations.AggregationBuilder;
Expand All @@ -39,7 +40,6 @@
import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.Map;
Expand Down Expand Up @@ -95,29 +95,29 @@ protected void innerWriteTo(StreamOutput out) throws IOException {
}

@Override
protected ValuesSourceAggregatorFactory<WithOrdinals> innerBuild(SearchContext context,
ValuesSourceConfig<WithOrdinals> config,
AggregatorFactory parent,
Builder subFactoriesBuilder) throws IOException {
return new ChildrenAggregatorFactory(name, config, childFilter, parentFilter, context, parent,
protected ValuesSourceAggregatorFactory<WithOrdinals> innerBuild(QueryShardContext queryShardContext,
ValuesSourceConfig<WithOrdinals> config,
AggregatorFactory parent,
Builder subFactoriesBuilder) throws IOException {
return new ChildrenAggregatorFactory(name, config, childFilter, parentFilter, queryShardContext, parent,
subFactoriesBuilder, metaData);
}

@Override
protected ValuesSourceConfig<WithOrdinals> resolveConfig(SearchContext context) {
protected ValuesSourceConfig<WithOrdinals> resolveConfig(QueryShardContext queryShardContext) {
ValuesSourceConfig<WithOrdinals> config = new ValuesSourceConfig<>(ValuesSourceType.BYTES);
joinFieldResolveConfig(context, config);
joinFieldResolveConfig(queryShardContext, config);
return config;
}

private void joinFieldResolveConfig(SearchContext context, ValuesSourceConfig<WithOrdinals> config) {
ParentJoinFieldMapper parentJoinFieldMapper = ParentJoinFieldMapper.getMapper(context.mapperService());
private void joinFieldResolveConfig(QueryShardContext queryShardContext, ValuesSourceConfig<WithOrdinals> config) {
ParentJoinFieldMapper parentJoinFieldMapper = ParentJoinFieldMapper.getMapper(queryShardContext.getMapperService());
ParentIdFieldMapper parentIdFieldMapper = parentJoinFieldMapper.getParentIdFieldMapper(childType, false);
if (parentIdFieldMapper != null) {
parentFilter = parentIdFieldMapper.getParentFilter();
childFilter = parentIdFieldMapper.getChildFilter(childType);
MappedFieldType fieldType = parentIdFieldMapper.fieldType();
final SortedSetDVOrdinalsIndexFieldData fieldData = context.getForField(fieldType);
final SortedSetDVOrdinalsIndexFieldData fieldData = queryShardContext.getForField(fieldType);
config.fieldContext(new FieldContext(fieldType.name(), fieldData, fieldType));
} else {
config.unmapped(true);
Expand Down
Loading

0 comments on commit 5eeb3c5

Please sign in to comment.