refactor the fast filter init logic in composite aggregation

Signed-off-by: bowenlan-amzn <[email protected]>
bowenlan-amzn · Jan 3, 2024 · e27d57c · e27d57c
1 parent 480eb71
commit e27d57c
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 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
@@ -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<DocValueFormat> 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/...rc/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/...rc/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
@@ -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/...java/org/opensearch/search/aggregations/bucket/composite/CompositeValuesSourceConfig.java b/...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/...c/main/java/org/opensearch/search/aggregations/bucket/composite/RoundingValuesSource.java b/...c/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;