From 53cb70f70aff4c56e6b59e1f55198a9a4ddb8c48 Mon Sep 17 00:00:00 2001
From: bowenlan-amzn <bowenlan23@gmail.com>
Date: Tue, 21 May 2024 11:10:06 -0700
Subject: [PATCH] Refactor the ranges representation

Signed-off-by: bowenlan-amzn <bowenlan23@gmail.com>
---
 .idea/runConfigurations/Debug_OpenSearch.xml  |   6 +-
 .../bucket/FastFilterRewriteHelper.java       | 188 ++++++++++++++----
 .../bucket/composite/CompositeAggregator.java |   3 +-
 .../bucket/range/RangeAggregator.java         |  22 +-
 4 files changed, 172 insertions(+), 47 deletions(-)

diff --git a/.idea/runConfigurations/Debug_OpenSearch.xml b/.idea/runConfigurations/Debug_OpenSearch.xml
index 0d8bf59823acf..c18046f873477 100644
--- a/.idea/runConfigurations/Debug_OpenSearch.xml
+++ b/.idea/runConfigurations/Debug_OpenSearch.xml
@@ -6,6 +6,10 @@
     <option name="HOST" value="localhost" />
     <option name="PORT" value="5005" />
     <option name="AUTO_RESTART" value="true" />
+    <RunnerSettings RunnerId="Debug">
+      <option name="DEBUG_PORT" value="5005" />
+      <option name="LOCAL" value="false" />
+    </RunnerSettings>
     <method v="2" />
   </configuration>
-</component>
+</component>
\ No newline at end of file
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 c8ce39a52f869..6efab16711def 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
@@ -10,6 +10,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
@@ -37,12 +38,12 @@
 import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceConfig;
 import org.opensearch.search.aggregations.bucket.composite.RoundingValuesSource;
 import org.opensearch.search.aggregations.bucket.histogram.LongBounds;
+import org.opensearch.search.aggregations.bucket.range.RangeAggregator.Range;
+import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.OptionalLong;
@@ -188,7 +189,7 @@ public static class FastFilterContext {
         private final SearchContext context;
 
         private String fieldName;
-        private long[][] ranges;
+        private Ranges ranges;
 
         // debug info related fields
         public int leaf;
@@ -230,8 +231,8 @@ public void buildRanges() throws IOException {
             }
         }
 
-        public long[][] buildRanges(LeafReaderContext leaf) throws IOException {
-            long[][] ranges = this.aggregationType.buildRanges(leaf, context);
+        public Ranges buildRanges(LeafReaderContext leaf) throws IOException {
+            Ranges ranges = this.aggregationType.buildRanges(leaf, context);
             if (ranges != null) {
                 logger.debug("Ranges built for shard {} segment {}", context.indexShard().shardId(), leaf.ord);
             }
@@ -250,11 +251,13 @@ private void consumeDebugInfo(DebugInfo debug) {
     interface AggregationType {
         boolean isRewriteable(Object parent, int subAggLength);
 
-        long[][] buildRanges(SearchContext ctx) throws IOException;
+        Ranges buildRanges(SearchContext ctx) throws IOException;
 
-        long[][] buildRanges(LeafReaderContext leaf, SearchContext ctx) throws IOException;
+        Ranges buildRanges(LeafReaderContext leaf, SearchContext ctx) throws IOException;
     }
 
+
+
     /**
      * For date histogram aggregation
      */
@@ -286,13 +289,13 @@ public boolean isRewriteable(Object parent, int subAggLength) {
         }
 
         @Override
-        public long[][] buildRanges(SearchContext context) throws IOException {
+        public Ranges buildRanges(SearchContext context) throws IOException {
             long[] bounds = getDateHistoAggBounds(context, fieldType.name());
             logger.debug("Bounds are {} for shard {}", bounds, context.indexShard().shardId());
             return buildRanges(context, bounds);
         }
 
-        private long[][] buildRanges(SearchContext context, long[] bounds) throws IOException {
+        private Ranges buildRanges(SearchContext context, long[] bounds) throws IOException {
             bounds = processHardBounds(bounds);
             if (bounds == null) {
                 return null;
@@ -320,7 +323,7 @@ private long[][] buildRanges(SearchContext context, long[] bounds) throws IOExce
         }
 
         @Override
-        public long[][] buildRanges(LeafReaderContext leaf, SearchContext context) throws IOException {
+        public Ranges buildRanges(LeafReaderContext leaf, SearchContext context) throws IOException {
             long[] bounds = getSegmentBounds(leaf, fieldType.name());
             logger.debug("Bounds are {} for shard {} segment {}", bounds, context.indexShard().shardId(), leaf.ord);
             return buildRanges(context, bounds);
@@ -356,6 +359,47 @@ public DateFieldMapper.DateFieldType getFieldType() {
         }
     }
 
+    public static class RangeAggregationType implements FastFilterRewriteHelper.AggregationType {
+
+        private final ValuesSource.Numeric source;
+        private final Range[] ranges;
+
+        public RangeAggregationType(ValuesSource.Numeric source, Range[] ranges) {
+            this.source = source;
+            this.ranges = ranges;
+        }
+
+        @Override
+        public boolean isRewriteable(Object parent, int subAggLength) {
+            if (parent == null && subAggLength == 0) {
+                // don't accept values source with scripts
+                if (source instanceof ValuesSource.Numeric.FieldData) {
+                    // ranges are already sorted by from and then to
+                    // we want ranges not overlapping with each other
+                    double prevTo = ranges[0].getTo();
+                    for (int i = 1; i < ranges.length; i++) {
+                        if (prevTo > ranges[i].getFrom()) {
+                            return false;
+                        }
+                        prevTo = ranges[i].getTo();
+                    }
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        @Override
+        public Ranges buildRanges(SearchContext ctx) throws IOException {
+            return new Ranges(new byte[][]{}, new byte[][]{});
+        }
+
+        @Override
+        public Ranges buildRanges(LeafReaderContext leaf, SearchContext ctx) throws IOException {
+            return new Ranges(new byte[][]{}, new byte[][]{});
+        }
+    }
+
     public static boolean isCompositeAggRewriteable(CompositeValuesSourceConfig[] sourceConfigs) {
         return sourceConfigs.length == 1 && sourceConfigs[0].valuesSource() instanceof RoundingValuesSource;
     }
@@ -407,7 +451,7 @@ public static boolean tryFastFilterAggregation(
         if (!fastFilterContext.rangesBuiltAtShardLevel && !segmentMatchAll(fastFilterContext.context, ctx)) {
             return false;
         }
-        long[][] ranges = fastFilterContext.ranges;
+        Ranges ranges = fastFilterContext.ranges;
         if (ranges == null) {
             logger.debug(
                 "Shard {} segment {} functionally match all documents. Build the fast filter",
@@ -445,7 +489,7 @@ private static boolean segmentMatchAll(SearchContext ctx, LeafReaderContext leaf
      * Creates the date ranges from date histo aggregations using its interval,
      * and min/max boundaries
      */
-    private static long[][] createRangesFromAgg(
+    private static Ranges createRangesFromAgg(
         final SearchContext context,
         final DateFieldMapper.DateFieldType fieldType,
         final long interval,
@@ -491,7 +535,18 @@ private static long[][] createRangesFromAgg(
             }
         }
 
-        return ranges;
+        byte[][] mins = new byte[ranges.length][];
+        byte[][] maxs = new byte[ranges.length][];
+        for (int i = 0; i < ranges.length; i++) {
+            byte[] min = new byte[8];
+            byte[] max = new byte[8];
+            LongPoint.encodeDimension(ranges[i][0], min, 0);
+            LongPoint.encodeDimension(ranges[i][1], max, 0);
+            mins[i] = min;
+            maxs[i] = max;
+        }
+
+        return new Ranges(mins, maxs);
     }
 
     /**
@@ -499,35 +554,42 @@ private static long[][] createRangesFromAgg(
      */
     private static DebugInfo multiRangesTraverse(
         final PointValues.PointTree tree,
-        final long[][] ranges,
+        final Ranges ranges,
         final BiConsumer<Long, Integer> incrementDocCount,
         final DateFieldMapper.DateFieldType fieldType,
         final int maxNumNonZeroRanges
     ) throws IOException {
         // ranges are connected and in ascending order
-        Iterator<long[]> rangeIter = Arrays.stream(ranges).iterator();
-        long[] activeRange = rangeIter.next();
+        // Iterator<long[]> rangeIter = Arrays.stream(ranges).iterator();
+        // long[] activeRange = rangeIter.next();
+        //
+        // // make sure the first range at least crosses the min value of the tree
+        // DebugInfo debugInfo = new DebugInfo();
+        // if (activeRange[0] > NumericUtils.sortableBytesToLong(tree.getMaxPackedValue(), 0)) {
+        //     logger.debug("No ranges match the query, skip the fast filter optimization");
+        //     return debugInfo;
+        // }
+        // while (activeRange[1] < NumericUtils.sortableBytesToLong(tree.getMinPackedValue(), 0)) {
+        //     if (!rangeIter.hasNext()) {
+        //         logger.debug("No ranges match the query, skip the fast filter optimization");
+        //         return debugInfo;
+        //     }
+        //     activeRange = rangeIter.next();
+        // }
 
-        // make sure the first range at least crosses the min value of the tree
         DebugInfo debugInfo = new DebugInfo();
-        if (activeRange[0] > NumericUtils.sortableBytesToLong(tree.getMaxPackedValue(), 0)) {
+        int activeIndex = ranges.firstRange(tree.getMinPackedValue(), tree.getMaxPackedValue());
+        if (activeIndex < 0) {
             logger.debug("No ranges match the query, skip the fast filter optimization");
             return debugInfo;
         }
-        while (activeRange[1] < NumericUtils.sortableBytesToLong(tree.getMinPackedValue(), 0)) {
-            if (!rangeIter.hasNext()) {
-                logger.debug("No ranges match the query, skip the fast filter optimization");
-                return debugInfo;
-            }
-            activeRange = rangeIter.next();
-        }
 
         RangeCollectorForPointTree collector = new RangeCollectorForPointTree(
             incrementDocCount,
             fieldType,
-            rangeIter,
             maxNumNonZeroRanges,
-            activeRange
+            ranges,
+            activeIndex
         );
 
         final ArrayUtil.ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(8);
@@ -542,6 +604,40 @@ private static DebugInfo multiRangesTraverse(
         return debugInfo;
     }
 
+    private static final ArrayUtil.ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(8);
+
+    private static int compareByteValue(byte[] value1, byte[] value2) {
+        return comparator.compare(value1, 0, value2, 0);
+    }
+
+    public static class Ranges {
+        byte[][] min;
+        byte[][] max;
+        int size;
+
+        Ranges(byte[][] min, byte[][] max) {
+            this.min = min;
+            this.max = max;
+            assert min.length == max.length;
+            this.size =  min.length;
+        }
+
+        // method to return the first range
+        public int firstRange(byte[] globalMin, byte[] globalMax) {
+            if (compareByteValue(min[0], globalMax) > 0) {
+                return -1;
+            }
+            int i = 0;
+            while (compareByteValue(max[i], globalMin) < 0) {
+                i++;
+                if (i >= size) {
+                    return -1;
+                }
+            }
+            return i;
+        }
+    }
+
     private static void intersectWithRanges(
         PointValues.IntersectVisitor visitor,
         PointValues.PointTree pointTree,
@@ -651,9 +747,12 @@ private static class RangeCollectorForPointTree {
         private final DateFieldMapper.DateFieldType fieldType;
         private int counter = 0;
 
-        private long[] activeRange;
+        // private long[] activeRange;
         private byte[][] activeRangeAsByteArray;
-        private final Iterator<long[]> rangeIter;
+        // private final Iterator<long[]> rangeIter;
+
+        private Ranges ranges;
+        private int activeIndex;
 
         private int visitedRange = 0;
         private final int maxNumNonZeroRange;
@@ -661,15 +760,15 @@ private static class RangeCollectorForPointTree {
         public RangeCollectorForPointTree(
             BiConsumer<Long, Integer> incrementDocCount,
             DateFieldMapper.DateFieldType fieldType,
-            Iterator<long[]> rangeIter,
             int maxNumNonZeroRange,
-            long[] activeRange
+            Ranges ranges,
+            int activeIndex
         ) {
             this.incrementDocCount = incrementDocCount;
             this.fieldType = fieldType;
-            this.rangeIter = rangeIter;
             this.maxNumNonZeroRange = maxNumNonZeroRange;
-            this.activeRange = activeRange;
+            this.ranges = ranges;
+            this.activeIndex = activeIndex;
             this.activeRangeAsByteArray = activeRangeAsByteArray();
         }
 
@@ -683,9 +782,10 @@ private void countNode(int count) {
 
         private void finalizePreviousRange() {
             if (counter > 0) {
-                logger.debug("finalize previous range: {}", activeRange[0]);
+                long range = LongPoint.decodeDimension(ranges.min[activeIndex],0);
+                logger.debug("finalize previous range: {}", range);
                 logger.debug("counter: {}", counter);
-                incrementDocCount.accept(fieldType.convertNanosToMillis(activeRange[0]), counter);
+                incrementDocCount.accept(fieldType.convertNanosToMillis(range), counter);
                 counter = 0;
             }
         }
@@ -697,10 +797,13 @@ private boolean iterateRangeEnd(byte[] value, BiFunction<byte[], byte[], Integer
             // the new value may not be contiguous to the previous one
             // so try to find the first next range that cross the new value
             while (comparator.apply(activeRangeAsByteArray[1], value) < 0) {
-                if (!rangeIter.hasNext()) {
+                // if (!rangeIter.hasNext()) {
+                //     return true;
+                // }
+                // activeRange = rangeIter.next();
+                if (++activeIndex >= ranges.size) {
                     return true;
                 }
-                activeRange = rangeIter.next();
                 activeRangeAsByteArray = activeRangeAsByteArray();
             }
             visitedRange++;
@@ -708,11 +811,12 @@ private boolean iterateRangeEnd(byte[] value, BiFunction<byte[], byte[], Integer
         }
 
         private byte[][] activeRangeAsByteArray() {
-            byte[] lower = new byte[8];
-            byte[] upper = new byte[8];
-            NumericUtils.longToSortableBytes(activeRange[0], lower, 0);
-            NumericUtils.longToSortableBytes(activeRange[1], upper, 0);
-            return new byte[][] { lower, upper };
+            // byte[] lower = new byte[8];
+            // byte[] upper = new byte[8];
+            // NumericUtils.longToSortableBytes(activeRange[0], lower, 0);
+            // NumericUtils.longToSortableBytes(activeRange[1], upper, 0);
+            // return new byte[][] { lower, upper };
+            return new byte[][] { ranges.min[activeIndex], ranges.max[activeIndex] };
         }
     }
 
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 3713d8f83990d..5cc16c1dc89b3 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
@@ -74,6 +74,7 @@
 import org.opensearch.search.aggregations.MultiBucketConsumerService;
 import org.opensearch.search.aggregations.bucket.BucketsAggregator;
 import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper;
+import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper.AbstractDateHistogramAggregationType;
 import org.opensearch.search.aggregations.bucket.missing.MissingOrder;
 import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
 import org.opensearch.search.internal.SearchContext;
@@ -180,7 +181,7 @@ public final class CompositeAggregator extends BucketsAggregator {
     /**
      * Currently the filter rewrite is only supported for date histograms
      */
-    public class CompositeAggregationType extends FastFilterRewriteHelper.AbstractDateHistogramAggregationType {
+    public class CompositeAggregationType extends AbstractDateHistogramAggregationType {
         private final RoundingValuesSource valuesSource;
         private long afterKey = -1L;
 
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java
index b56b817b8177b..4a6faf4b4bd21 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java
@@ -54,6 +54,7 @@
 import org.opensearch.search.aggregations.LeafBucketCollectorBase;
 import org.opensearch.search.aggregations.NonCollectingAggregator;
 import org.opensearch.search.aggregations.bucket.BucketsAggregator;
+import org.opensearch.search.aggregations.bucket.FastFilterRewriteHelper;
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
 
@@ -245,6 +246,8 @@ public boolean equals(Object obj) {
 
     final double[] maxTo;
 
+    private final FastFilterRewriteHelper.FastFilterContext fastFilterContext;
+
     public RangeAggregator(
         String name,
         AggregatorFactories factories,
@@ -258,7 +261,6 @@ public RangeAggregator(
         CardinalityUpperBound cardinality,
         Map<String, Object> metadata
     ) throws IOException {
-
         super(name, factories, context, parent, cardinality.multiply(ranges.length), metadata);
         assert valuesSource != null;
         this.valuesSource = valuesSource;
@@ -266,16 +268,29 @@ public RangeAggregator(
         this.keyed = keyed;
         this.rangeFactory = rangeFactory;
 
-        this.ranges = ranges;
+        this.ranges = ranges; // sorted by the from then to
+
+        // if the range connect with each other, we can try apply the optimization
+        fastFilterContext = new FastFilterRewriteHelper.FastFilterContext(context);
+        fastFilterContext.setAggregationType(
+            new FastFilterRewriteHelper.RangeAggregationType(
+                valuesSource,
+                ranges
+            )
+        );
+        if (fastFilterContext.isRewriteable(parent, subAggregators.length)) {
+            fastFilterContext.buildRanges();
+        }
 
         maxTo = new double[this.ranges.length];
         maxTo[0] = this.ranges[0].to;
         for (int i = 1; i < this.ranges.length; ++i) {
             maxTo[i] = Math.max(this.ranges[i].to, maxTo[i - 1]);
         }
-
     }
 
+
+
     @Override
     public ScoreMode scoreMode() {
         if (valuesSource != null && valuesSource.needsScores()) {
@@ -293,6 +308,7 @@ public void collect(int doc, long bucket) throws IOException {
                 if (values.advanceExact(doc)) {
                     final int valuesCount = values.docValueCount();
                     for (int i = 0, lo = 0; i < valuesCount; ++i) {
+                        // for each value of the document
                         final double value = values.nextValue();
                         lo = collect(doc, value, bucket, lo);
                     }