use of org.opensearch.search.aggregations.InternalMultiBucketAggregation in project OpenSearch by opensearch-project.
the class BucketSelectorPipelineAggregator method reduce.
@Override
public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext reduceContext) {
InternalMultiBucketAggregation<InternalMultiBucketAggregation, InternalMultiBucketAggregation.InternalBucket> originalAgg = (InternalMultiBucketAggregation<InternalMultiBucketAggregation, InternalMultiBucketAggregation.InternalBucket>) aggregation;
List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = originalAgg.getBuckets();
BucketAggregationSelectorScript.Factory factory = reduceContext.scriptService().compile(script, BucketAggregationSelectorScript.CONTEXT);
List<InternalMultiBucketAggregation.InternalBucket> newBuckets = new ArrayList<>();
for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
Map<String, Object> vars = new HashMap<>();
if (script.getParams() != null) {
vars.putAll(script.getParams());
}
for (Map.Entry<String, String> entry : bucketsPathsMap.entrySet()) {
String varName = entry.getKey();
String bucketsPath = entry.getValue();
Double value = resolveBucketValue(originalAgg, bucket, bucketsPath, gapPolicy);
vars.put(varName, value);
}
// TODO: can we use one instance of the script for all buckets? it should be stateless?
BucketAggregationSelectorScript executableScript = factory.newInstance(vars);
if (executableScript.execute()) {
newBuckets.add(bucket);
}
}
return originalAgg.create(newBuckets);
}
use of org.opensearch.search.aggregations.InternalMultiBucketAggregation in project OpenSearch by opensearch-project.
the class CumulativeSumPipelineAggregator method reduce.
@Override
public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext reduceContext) {
InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket> histo = (InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket>) aggregation;
List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = histo.getBuckets();
HistogramFactory factory = (HistogramFactory) histo;
List<Bucket> newBuckets = new ArrayList<>(buckets.size());
double sum = 0;
for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
Double thisBucketValue = resolveBucketValue(histo, bucket, bucketsPaths()[0], GapPolicy.INSERT_ZEROS);
// Only increment the sum if it's a finite value, otherwise "increment by zero" is correct
if (thisBucketValue != null && thisBucketValue.isInfinite() == false && thisBucketValue.isNaN() == false) {
sum += thisBucketValue;
}
List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false).map((p) -> (InternalAggregation) p).collect(Collectors.toList());
aggs.add(new InternalSimpleValue(name(), sum, formatter, metadata()));
Bucket newBucket = factory.createBucket(factory.getKey(bucket), bucket.getDocCount(), InternalAggregations.from(aggs));
newBuckets.add(newBucket);
}
return factory.createAggregation(newBuckets);
}
use of org.opensearch.search.aggregations.InternalMultiBucketAggregation in project OpenSearch by opensearch-project.
the class MovAvgPipelineAggregator method reduce.
@Override
public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext reduceContext) {
InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket> histo = (InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket>) aggregation;
List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = histo.getBuckets();
HistogramFactory factory = (HistogramFactory) histo;
List<Bucket> newBuckets = new ArrayList<>();
EvictingQueue<Double> values = new EvictingQueue<>(this.window);
Number lastValidKey = 0;
int lastValidPosition = 0;
int counter = 0;
// Do we need to fit the model parameters to the data?
if (minimize) {
assert (model.canBeMinimized());
model = minimize(buckets, histo, model);
}
for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
Double thisBucketValue = resolveBucketValue(histo, bucket, bucketsPaths()[0], gapPolicy);
// Default is to reuse existing bucket. Simplifies the rest of the logic,
// since we only change newBucket if we can add to it
Bucket newBucket = bucket;
if ((thisBucketValue == null || thisBucketValue.equals(Double.NaN)) == false) {
// Some models (e.g. HoltWinters) have certain preconditions that must be met
if (model.hasValue(values.size())) {
double movavg = model.next(values);
List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false).map((p) -> (InternalAggregation) p).collect(Collectors.toList());
aggs.add(new InternalSimpleValue(name(), movavg, formatter, metadata()));
newBucket = factory.createBucket(factory.getKey(bucket), bucket.getDocCount(), InternalAggregations.from(aggs));
}
if (predict > 0) {
lastValidKey = factory.getKey(bucket);
lastValidPosition = counter;
}
values.offer(thisBucketValue);
}
counter += 1;
newBuckets.add(newBucket);
}
if (buckets.size() > 0 && predict > 0) {
double[] predictions = model.predict(values, predict);
for (int i = 0; i < predictions.length; i++) {
List<InternalAggregation> aggs;
Number newKey = factory.nextKey(lastValidKey);
if (lastValidPosition + i + 1 < newBuckets.size()) {
Bucket bucket = newBuckets.get(lastValidPosition + i + 1);
// Get the existing aggs in the bucket so we don't clobber data
aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false).map((p) -> (InternalAggregation) p).collect(Collectors.toList());
aggs.add(new InternalSimpleValue(name(), predictions[i], formatter, metadata()));
Bucket newBucket = factory.createBucket(newKey, bucket.getDocCount(), InternalAggregations.from(aggs));
// Overwrite the existing bucket with the new version
newBuckets.set(lastValidPosition + i + 1, newBucket);
} else {
// Not seen before, create fresh
aggs = new ArrayList<>();
aggs.add(new InternalSimpleValue(name(), predictions[i], formatter, metadata()));
Bucket newBucket = factory.createBucket(newKey, 0, InternalAggregations.from(aggs));
// Since this is a new bucket, simply append it
newBuckets.add(newBucket);
}
lastValidKey = newKey;
}
}
return factory.createAggregation(newBuckets);
}
use of org.opensearch.search.aggregations.InternalMultiBucketAggregation in project OpenSearch by opensearch-project.
the class TermsAggregatorTests method testWithNestedAggregations.
public void testWithNestedAggregations() throws IOException {
try (Directory directory = newDirectory()) {
try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
for (int i = 0; i < 10; i++) {
int[] nestedValues = new int[i];
for (int j = 0; j < i; j++) {
nestedValues[j] = j;
}
indexWriter.addDocuments(generateDocsWithNested(Integer.toString(i), i, nestedValues));
}
indexWriter.commit();
for (Aggregator.SubAggCollectionMode mode : Aggregator.SubAggCollectionMode.values()) {
for (boolean withScore : new boolean[] { true, false }) {
NestedAggregationBuilder nested = new NestedAggregationBuilder("nested", "nested_object").subAggregation(new TermsAggregationBuilder("terms").userValueTypeHint(ValueType.LONG).field("nested_value").collectMode(mode).order(BucketOrder.key(true)).subAggregation(new TopHitsAggregationBuilder("top_hits").sort(withScore ? new ScoreSortBuilder() : new FieldSortBuilder("_doc")).storedField("_none_")));
MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("nested_value", NumberFieldMapper.NumberType.LONG);
try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) {
{
InternalNested result = searchAndReduce(newSearcher(indexReader, false, true), // match root document only
new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), nested, fieldType);
InternalMultiBucketAggregation<?, ?> terms = result.getAggregations().get("terms");
assertNestedTopHitsScore(terms, withScore);
}
{
FilterAggregationBuilder filter = new FilterAggregationBuilder("filter", new MatchAllQueryBuilder()).subAggregation(nested);
InternalFilter result = searchAndReduce(newSearcher(indexReader, false, true), // match root document only
new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), filter, fieldType);
InternalNested nestedResult = result.getAggregations().get("nested");
InternalMultiBucketAggregation<?, ?> terms = nestedResult.getAggregations().get("terms");
assertNestedTopHitsScore(terms, withScore);
}
}
}
}
}
}
}
use of org.opensearch.search.aggregations.InternalMultiBucketAggregation in project OpenSearch by opensearch-project.
the class MovFnPipelineAggregator method reduce.
@Override
public InternalAggregation reduce(InternalAggregation aggregation, InternalAggregation.ReduceContext reduceContext) {
InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket> histo = (InternalMultiBucketAggregation<? extends InternalMultiBucketAggregation, ? extends InternalMultiBucketAggregation.InternalBucket>) aggregation;
List<? extends InternalMultiBucketAggregation.InternalBucket> buckets = histo.getBuckets();
HistogramFactory factory = (HistogramFactory) histo;
List<MultiBucketsAggregation.Bucket> newBuckets = new ArrayList<>();
// Initialize the script
MovingFunctionScript.Factory scriptFactory = reduceContext.scriptService().compile(script, MovingFunctionScript.CONTEXT);
Map<String, Object> vars = new HashMap<>();
if (script.getParams() != null) {
vars.putAll(script.getParams());
}
MovingFunctionScript executableScript = scriptFactory.newInstance();
List<Double> values = buckets.stream().map(b -> resolveBucketValue(histo, b, bucketsPaths()[0], gapPolicy)).filter(v -> v != null && v.isNaN() == false).collect(Collectors.toList());
int index = 0;
for (InternalMultiBucketAggregation.InternalBucket bucket : buckets) {
Double thisBucketValue = resolveBucketValue(histo, bucket, bucketsPaths()[0], gapPolicy);
// Default is to reuse existing bucket. Simplifies the rest of the logic,
// since we only change newBucket if we can add to it
MultiBucketsAggregation.Bucket newBucket = bucket;
if (thisBucketValue != null && thisBucketValue.isNaN() == false) {
// The custom context mandates that the script returns a double (not Double) so we
// don't need null checks, etc.
int fromIndex = clamp(index - window + shift, values);
int toIndex = clamp(index + shift, values);
double movavg = executableScript.execute(vars, values.subList(fromIndex, toIndex).stream().mapToDouble(Double::doubleValue).toArray());
List<InternalAggregation> aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false).map(InternalAggregation.class::cast).collect(Collectors.toList());
aggs.add(new InternalSimpleValue(name(), movavg, formatter, metadata()));
newBucket = factory.createBucket(factory.getKey(bucket), bucket.getDocCount(), InternalAggregations.from(aggs));
index++;
}
newBuckets.add(newBucket);
}
return factory.createAggregation(newBuckets);
}
Aggregations