From e27d57c769399408b08ec0476099d81c90fe1055 Mon Sep 17 00:00:00 2001 From: bowenlan-amzn Date: Wed, 3 Jan 2024 10:48:15 -0800 Subject: [PATCH] refactor the fast filter init logic in composite aggregation Signed-off-by: bowenlan-amzn --- .../bucket/FastFilterRewriteHelper.java | 35 +++++++++++++++-- .../bucket/composite/CompositeAggregator.java | 38 +++++++------------ .../bucket/composite/CompositeKey.java | 6 +-- .../CompositeValuesSourceConfig.java | 2 +- .../composite/RoundingValuesSource.java | 2 +- 5 files changed, 49 insertions(+), 34 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java index 0058a875c4032..083e4702e02f0 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/FastFilterRewriteHelper.java @@ -24,6 +24,10 @@ import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.query.DateRangeIncludingNowQuery; +import org.opensearch.search.DocValueFormat; +import org.opensearch.search.aggregations.bucket.composite.CompositeKey; +import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig; +import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource; import org.opensearch.search.internal.SearchContext; import java.io.IOException; @@ -233,11 +237,31 @@ public static class FastFilterContext { private final Type type; + private RoundingValuesSource valuesSource = null; + public FastFilterContext(MappedFieldType fieldType) { this.fieldType = fieldType; this.type = Type.DATE_HISTO; } + public FastFilterContext(CompositeValuesSourceConfig[] sourceConfigs, CompositeKey rawAfterKey, List formats) { + if (sourceConfigs.length == 1 && sourceConfigs[0].valuesSource() instanceof RoundingValuesSource) { + this.fieldType = sourceConfigs[0].fieldType(); + this.valuesSource = (RoundingValuesSource) sourceConfigs[0].valuesSource(); + this.missing = sourceConfigs[0].missingBucket(); + this.hasScript = sourceConfigs[0].hasScript(); + if (rawAfterKey != null) { + assert rawAfterKey.size() == 1 && formats.size() == 1; + this.afterKey = formats.get(0).parseLong(rawAfterKey.get(0).toString(), false, () -> { + throw new IllegalArgumentException("now() is not supported in [after] key"); + }); + } + } else { + this.fieldType = null; + } + this.type = Type.DATE_HISTO; + } + public FastFilterContext(Type type) { this.fieldType = null; this.type = type; @@ -248,6 +272,10 @@ public DateFieldMapper.DateFieldType getFieldType() { return (DateFieldMapper.DateFieldType) fieldType; } + public RoundingValuesSource getDateHistogramSource() { + return valuesSource; + } + public void setSize(int size) { this.size = size; } @@ -256,10 +284,6 @@ public void setFilters(Weight[] filters) { this.filters = filters; } - public void setAfterKey(long afterKey) { - this.afterKey = afterKey; - } - public void setMissing(boolean missing) { this.missing = missing; } @@ -283,6 +307,9 @@ public boolean isRewriteable(Object parent, int subAggLength) { return false; } + /** + * Different types have different pre-conditions, filter building logic, etc. + */ public enum Type { FILTERS, DATE_HISTO diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java index 89628de6c5812..71249266c6bb3 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java @@ -115,9 +115,9 @@ final class CompositeAggregator extends BucketsAggregator { private boolean earlyTerminated; + private final FastFilterRewriteHelper.FastFilterContext fastFilterContext; private LongKeyedBucketOrds bucketOrds = null; private Rounding.Prepared preparedRounding = null; - private FastFilterRewriteHelper.FastFilterContext fastFilterContext = null; CompositeAggregator( String name, @@ -163,31 +163,19 @@ final class CompositeAggregator extends BucketsAggregator { this.queue = new CompositeValuesCollectorQueue(context.bigArrays(), sources, size, rawAfterKey); this.rawAfterKey = rawAfterKey; - // Try fast filter optimization when the only source is date histogram - if (sourceConfigs.length == 1 && sourceConfigs[0].valuesSource() instanceof RoundingValuesSource) { - RoundingValuesSource dateHistogramSource = (RoundingValuesSource) sourceConfigs[0].valuesSource(); - bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), CardinalityUpperBound.ONE); + fastFilterContext = new FastFilterRewriteHelper.FastFilterContext(sourceConfigs, rawAfterKey, formats); + if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { + RoundingValuesSource dateHistogramSource = fastFilterContext.getDateHistogramSource(); preparedRounding = dateHistogramSource.getPreparedRounding(); - fastFilterContext = new FastFilterRewriteHelper.FastFilterContext(sourceConfigs[0].fieldType()); - fastFilterContext.setMissing(sourceConfigs[0].missingBucket()); - fastFilterContext.setHasScript(sourceConfigs[0].hasScript()); - if (rawAfterKey != null) { - assert rawAfterKey.size() == 1 && formats.size() == 1; - long afterValue = formats.get(0).parseLong(rawAfterKey.get(0).toString(), false, () -> { - throw new IllegalArgumentException("now() is not supported in [after] key"); - }); - fastFilterContext.setAfterKey(afterValue); - } - if (fastFilterContext.isRewriteable(parent, subAggregators.length)) { - fastFilterContext.setSize(size); - FastFilterRewriteHelper.buildFastFilter( - context, - fc -> FastFilterRewriteHelper.getAggregationBounds(context, fc.getFieldType().name()), - x -> dateHistogramSource.getRounding(), - () -> preparedRounding, - fastFilterContext - ); - } + bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), CardinalityUpperBound.ONE); + fastFilterContext.setSize(size); + FastFilterRewriteHelper.buildFastFilter( + context, + fc -> FastFilterRewriteHelper.getAggregationBounds(context, fc.getFieldType().name()), + x -> dateHistogramSource.getRounding(), + () -> preparedRounding, + fastFilterContext + ); } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeKey.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeKey.java index 5ddeb22d33a6f..338ebdc66eef7 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeKey.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeKey.java @@ -44,7 +44,7 @@ * * @opensearch.internal */ -class CompositeKey implements Writeable { +public class CompositeKey implements Writeable { private final Comparable[] values; CompositeKey(Comparable... values) { @@ -64,11 +64,11 @@ Comparable[] values() { return values; } - int size() { + public int size() { return values.length; } - Comparable get(int pos) { + public Comparable get(int pos) { assert pos < values.length; return values[pos]; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeValuesSourceConfig.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeValuesSourceConfig.java index 788a4ddc15374..5289b3a34ab34 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeValuesSourceConfig.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeValuesSourceConfig.java @@ -156,7 +156,7 @@ public MissingOrder missingOrder() { /** * Returns true if the source contains a script that can change the value. */ - protected boolean hasScript() { + public boolean hasScript() { return hasScript; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/RoundingValuesSource.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/RoundingValuesSource.java index 37fbe8cb03fd9..3f5cf919f1755 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/RoundingValuesSource.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/RoundingValuesSource.java @@ -47,7 +47,7 @@ * * @opensearch.internal */ -class RoundingValuesSource extends ValuesSource.Numeric { +public class RoundingValuesSource extends ValuesSource.Numeric { private final ValuesSource.Numeric vs; private final Rounding.Prepared preparedRounding; private final Rounding rounding;