move multirangetraversal logic to helper

Signed-off-by: bowenlan-amzn <[email protected]>

server/src/main/java/org/opensearch/search/optimization/ranges/DateHistogramAggregatorBridge.java

@@ -26,7 +26,7 @@
 import java.util.OptionalLong;
 import java.util.function.BiConsumer;
 
-import static org.opensearch.search.optimization.ranges.OptimizationContext.multiRangesTraverse;
+import static org.opensearch.search.optimization.ranges.Helper.multiRangesTraverse;
 
 /**
  * For date histogram aggregation
@@ -74,7 +74,7 @@ public void prepareFromSegment(LeafReaderContext leaf) throws IOException {
         optimizationContext.setRangesFromSegment(buildRanges(bounds));
     }
 
-    private OptimizationContext.Ranges buildRanges(long[] bounds) {
+    private Ranges buildRanges(long[] bounds) {
         bounds = processHardBounds(bounds);
         if (bounds == null) {
             return null;

server/src/main/java/org/opensearch/search/optimization/ranges/Helper.java

@@ -12,33 +12,39 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.PointValues;
+import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.FieldExistsQuery;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.PointRangeQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.NumericUtils;
+import org.opensearch.common.CheckedRunnable;
 import org.opensearch.common.Rounding;
 import org.opensearch.common.lucene.search.function.FunctionScoreQuery;
 import org.opensearch.index.mapper.DateFieldMapper;
 import org.opensearch.index.query.DateRangeIncludingNowQuery;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.optimization.ranges.OptimizationContext.DebugInfo;
 
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiConsumer;
 import java.util.function.Function;
 
 import static org.opensearch.index.mapper.NumberFieldMapper.NumberType.LONG;
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
 
 /**
  * Utility class to help range filters rewrite optimization
  *
  * @opensearch.internal
  */
-public final class Helper {
+final class Helper {
 
     static final String loggerName = Helper.class.getPackageName();
     private static final Logger logger = LogManager.getLogger(loggerName);
@@ -155,7 +161,7 @@ private static long[] getBoundsWithRangeQuery(PointRangeQuery prq, String fieldN
      * Creates the date ranges from date histo aggregations using its interval,
      * and min/max boundaries
      */
-    static OptimizationContext.Ranges createRangesFromAgg(
+    static Ranges createRangesFromAgg(
         final DateFieldMapper.DateFieldType fieldType,
         final long interval,
         final Rounding.Prepared preparedRounding,
@@ -208,6 +214,185 @@ static OptimizationContext.Ranges createRangesFromAgg(
             uppers[i] = max;
         }
 
-        return new OptimizationContext.Ranges(lowers, uppers);
+        return new Ranges(lowers, uppers);
+    }
+
+    /**
+     * @param maxNumNonZeroRanges the number of non-zero ranges to collect
+     */
+    static DebugInfo multiRangesTraverse(
+        final PointValues.PointTree tree,
+        final Ranges ranges,
+        final BiConsumer<Integer, Integer> incrementDocCount,
+        final int maxNumNonZeroRanges
+    ) throws IOException {
+        DebugInfo debugInfo = new DebugInfo();
+        int activeIndex = ranges.firstRangeIndex(tree.getMinPackedValue(), tree.getMaxPackedValue());
+        if (activeIndex < 0) {
+            logger.debug("No ranges match the query, skip the fast filter optimization");
+            return debugInfo;
+        }
+        RangeCollectorForPointTree collector = new RangeCollectorForPointTree(incrementDocCount, maxNumNonZeroRanges, ranges, activeIndex);
+        PointValues.IntersectVisitor visitor = getIntersectVisitor(collector);
+        try {
+            intersectWithRanges(visitor, tree, collector, debugInfo);
+        } catch (CollectionTerminatedException e) {
+            logger.debug("Early terminate since no more range to collect");
+        }
+        collector.finalizePreviousRange();
+
+        return debugInfo;
+    }
+
+    private static void intersectWithRanges(
+        PointValues.IntersectVisitor visitor,
+        PointValues.PointTree pointTree,
+        RangeCollectorForPointTree collector,
+        DebugInfo debug
+    ) throws IOException {
+        PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
+
+        switch (r) {
+            case CELL_INSIDE_QUERY:
+                collector.countNode((int) pointTree.size());
+                debug.visitInner();
+                break;
+            case CELL_CROSSES_QUERY:
+                if (pointTree.moveToChild()) {
+                    do {
+                        intersectWithRanges(visitor, pointTree, collector, debug);
+                    } while (pointTree.moveToSibling());
+                    pointTree.moveToParent();
+                } else {
+                    pointTree.visitDocValues(visitor);
+                    debug.visitLeaf();
+                }
+                break;
+            case CELL_OUTSIDE_QUERY:
+        }
+    }
+
+    private static PointValues.IntersectVisitor getIntersectVisitor(RangeCollectorForPointTree collector) {
+        return new PointValues.IntersectVisitor() {
+            @Override
+            public void visit(int docID) {
+                // this branch should be unreachable
+                throw new UnsupportedOperationException(
+                    "This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited"
+                );
+            }
+
+            @Override
+            public void visit(int docID, byte[] packedValue) throws IOException {
+                visitPoints(packedValue, collector::count);
+            }
+
+            @Override
+            public void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException {
+                visitPoints(packedValue, () -> {
+                    for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) {
+                        collector.count();
+                    }
+                });
+            }
+
+            private void visitPoints(byte[] packedValue, CheckedRunnable<IOException> collect) throws IOException {
+                if (!collector.withinUpperBound(packedValue)) {
+                    collector.finalizePreviousRange();
+                    if (collector.iterateRangeEnd(packedValue)) {
+                        throw new CollectionTerminatedException();
+                    }
+                }
+
+                if (collector.withinRange(packedValue)) {
+                    collect.run();
+                }
+            }
+
+            @Override
+            public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
+                // try to find the first range that may collect values from this cell
+                if (!collector.withinUpperBound(minPackedValue)) {
+                    collector.finalizePreviousRange();
+                    if (collector.iterateRangeEnd(minPackedValue)) {
+                        throw new CollectionTerminatedException();
+                    }
+                }
+                // after the loop, min < upper
+                // cell could be outside [min max] lower
+                if (!collector.withinLowerBound(maxPackedValue)) {
+                    return PointValues.Relation.CELL_OUTSIDE_QUERY;
+                }
+                if (collector.withinRange(minPackedValue) && collector.withinRange(maxPackedValue)) {
+                    return PointValues.Relation.CELL_INSIDE_QUERY;
+                }
+                return PointValues.Relation.CELL_CROSSES_QUERY;
+            }
+        };
+    }
+
+    private static class RangeCollectorForPointTree {
+        private final BiConsumer<Integer, Integer> incrementRangeDocCount;
+        private int counter = 0;
+
+        private final Ranges ranges;
+        private int activeIndex;
+
+        private int visitedRange = 0;
+        private final int maxNumNonZeroRange;
+
+        public RangeCollectorForPointTree(
+            BiConsumer<Integer, Integer> incrementRangeDocCount,
+            int maxNumNonZeroRange,
+            Ranges ranges,
+            int activeIndex
+        ) {
+            this.incrementRangeDocCount = incrementRangeDocCount;
+            this.maxNumNonZeroRange = maxNumNonZeroRange;
+            this.ranges = ranges;
+            this.activeIndex = activeIndex;
+        }
+
+        private void count() {
+            counter++;
+        }
+
+        private void countNode(int count) {
+            counter += count;
+        }
+
+        private void finalizePreviousRange() {
+            if (counter > 0) {
+                incrementRangeDocCount.accept(activeIndex, counter);
+                counter = 0;
+            }
+        }
+
+        /**
+         * @return true when iterator exhausted or collect enough non-zero ranges
+         */
+        private boolean iterateRangeEnd(byte[] value) {
+            // 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 (!withinUpperBound(value)) {
+                if (++activeIndex >= ranges.size) {
+                    return true;
+                }
+            }
+            visitedRange++;
+            return visitedRange > maxNumNonZeroRange;
+        }
+
+        private boolean withinLowerBound(byte[] value) {
+            return Ranges.withinLowerBound(value, ranges.lowers[activeIndex]);
+        }
+
+        private boolean withinUpperBound(byte[] value) {
+            return Ranges.withinUpperBound(value, ranges.uppers[activeIndex]);
+        }
+
+        private boolean withinRange(byte[] value) {
+            return withinLowerBound(value) && withinUpperBound(value);
+        }
     }
 }