Capture local ranges. Context getter is volatile

Signed-off-by: Finn Carroll <[email protected]>
finnegancarroll · Aug 18, 2024 · d05b2e3 · d05b2e3
1 parent 95dc749
commit d05b2e3
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 20 deletions.
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
@@ -76,6 +76,7 @@
 import org.opensearch.search.aggregations.bucket.BucketsAggregator;
 import org.opensearch.search.aggregations.bucket.filterrewrite.CompositeAggregatorBridge;
 import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext;
+import org.opensearch.search.aggregations.bucket.filterrewrite.PackedValueRanges;
 import org.opensearch.search.aggregations.bucket.missing.MissingOrder;
 import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
 import org.opensearch.search.internal.SearchContext;
@@ -218,8 +219,9 @@ protected int rangeMax() {
             }
 
             @Override
-            protected long getOrd(int rangeIdx) {
-                long rangeStart = LongPoint.decodeDimension(filterRewriteOptimizationContext.getRanges().getLower(rangeIdx), 0);
+            protected long getOrd(int rangeIdx, PackedValueRanges ranges) {
+                assert(ranges != null);
+                long rangeStart = LongPoint.decodeDimension(ranges.getLower(rangeIdx), 0);
                 rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
                 long ord = bucketOrds.add(0, getRoundingPrepared().round(rangeStart));
 

diff --git a/...c/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java b/...c/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/AggregatorBridge.java
@@ -83,7 +83,7 @@ protected int rangeMax() {
     /**
      * Translate an index of the packed value range array to an agg bucket ordinal.
      */
-    protected long getOrd(int rangeIdx) {
+    protected long getOrd(int rangeIdx, PackedValueRanges ranges) {
         return rangeIdx;
     }
 
@@ -109,7 +109,7 @@ public final FilterRewriteOptimizationContext.DebugInfo tryOptimize(
                 ranges,
                 rangeMax(),
                 (activeIndex, docID) -> {
-                    long ord = this.getOrd(activeIndex);
+                    long ord = this.getOrd(activeIndex, ranges);
                     try {
                         incrementDocCount.accept(ord, (long) 1);
                         sub.collect(docID, ord);
@@ -124,7 +124,7 @@ public final FilterRewriteOptimizationContext.DebugInfo tryOptimize(
                 ranges,
                 rangeMax(),
                 (activeIndex, docCount) -> {
-                    long ord = this.getOrd(activeIndex);
+                    long ord = this.getOrd(activeIndex, ranges);
                     incrementDocCount.accept(ord, (long) docCount);
                 }
             );

diff --git a/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java b/...opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteOptimizationContext.java
@@ -66,7 +66,9 @@ private boolean canOptimize(final Object parent, SearchContext context) throws I
 
         boolean canOptimize = aggregatorBridge.canOptimize();
         if (canOptimize) {
-            aggregatorBridge.setRangesConsumer(this::setRanges);
+            aggregatorBridge.setRangesConsumer((PackedValueRanges ranges) -> {
+                this.ranges = ranges;
+            });
 
             this.shardId = context.indexShard().shardId().toString();
 
@@ -81,14 +83,6 @@ private boolean canOptimize(final Object parent, SearchContext context) throws I
         return canOptimize;
     }
 
-    public void setRanges(PackedValueRanges ranges) {
-        this.ranges = ranges;
-    }
-
-    public PackedValueRanges getRanges() {
-        return this.ranges;
-    }
-
     /**
      * Try to populate the bucket doc counts for aggregation
      * <p>

diff --git a/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java
@@ -202,10 +202,8 @@ protected Prepared getRoundingPrepared() {
             }
 
             @Override
-            protected long getOrd(int rangeIdx) {
-                PackedValueRanges ranges = filterRewriteOptimizationContext.getRanges();
+            protected long getOrd(int rangeIdx, PackedValueRanges ranges) {
                 assert(ranges != null);
-
                 long rangeStart = LongPoint.decodeDimension(ranges.getLower(rangeIdx), 0);
                 rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
                 long ord = getBucketOrds().add(0, getRoundingPrepared().round(rangeStart));

diff --git a/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java
@@ -147,10 +147,8 @@ protected long[] processHardBounds(long[] bounds) {
             }
 
             @Override
-            protected long getOrd(int rangeIdx) {
-                PackedValueRanges ranges = filterRewriteOptimizationContext.getRanges();
+            protected long getOrd(int rangeIdx, PackedValueRanges ranges) {
                 assert(ranges != null);
-
                 long rangeStart = LongPoint.decodeDimension(ranges.getLower(rangeIdx), 0);
                 rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
                 long ord = bucketOrds.add(0, getRoundingPrepared().round(rangeStart));