Move ordProducer delegate to agg bridge function. Ranges need to be p…

…roduced dynamically to support auto date histograms.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -32,6 +32,7 @@
 
 package org.opensearch.search.aggregations.bucket.composite;
 
+import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
@@ -193,14 +194,7 @@ public boolean canOptimize() {
  }
 
  @Override
- public void prepare() throws IOException {
- buildRanges(context);
- this.ordProducer = new DateHistogramAggregatorBridge.DateHistoOrdProducer(
- getFieldType(),
- optimizationContext.getRanges(),
- bucketOrds,
- getRoundingPrepared());
- }
+ public void prepare() throws IOException { buildRanges(context); }
 
  protected Rounding getRounding(final long low, final long high) {
  return valuesSource.getRounding();
@@ -224,6 +218,19 @@ protected long[] processAfterKey(long[] bounds, long interval) {
  protected int rangeMax() {
  return size;
  }
+
+ @Override
+ protected long getOrd(int rangeIdx){
+ long rangeStart = LongPoint.decodeDimension(optimizationContext.getRanges().getLower(rangeIdx), 0);
+ rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
+ long ord = bucketOrds.add(0, getRoundingPrepared().round(rangeStart));
+
+ if (ord < 0) { // already seen
+ ord = -1 - ord;
+ }
+
+ return ord;
+ }
  });
  if (optimizationContext.canOptimize(parent, context)) {
  optimizationContext.prepare();

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java

@@ -31,6 +31,7 @@
 
 package org.opensearch.search.aggregations.bucket.histogram;
 
+import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.CollectionTerminatedException;
@@ -167,14 +168,7 @@ public boolean canOptimize() {
  }
 
  @Override
- public void prepare() throws IOException {
- buildRanges(context);
- this.ordProducer = new DateHistogramAggregatorBridge.DateHistoOrdProducer(
- getFieldType(),
- optimizationContext.getRanges(),
- getBucketOrds(),
- getRoundingPrepared());
- }
+ public void prepare() throws IOException { buildRanges(context); }
 
  @Override
  protected Rounding getRounding(final long low, final long high) {
@@ -204,6 +198,19 @@ protected Rounding getRounding(final long low, final long high) {
  protected Prepared getRoundingPrepared() {
  return preparedRounding;
  }
+
+ @Override
+ protected long getOrd(int rangeIdx){
+ long rangeStart = LongPoint.decodeDimension(optimizationContext.getRanges().getLower(rangeIdx), 0);
+ rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
+ long ord = getBucketOrds().add(0, getRoundingPrepared().round(rangeStart));
+
+ if (ord < 0) { // already seen
+ ord = -1 - ord;
+ }
+
+ return ord;
+ }
  });
  if (optimizationContext.canOptimize(parent, context)) {
  optimizationContext.prepare();

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -31,6 +31,7 @@
 
 package org.opensearch.search.aggregations.bucket.histogram;
 
+import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.CollectionTerminatedException;
@@ -130,14 +131,7 @@ public boolean canOptimize() {
  }
 
  @Override
- public void prepare() throws IOException {
- buildRanges(context);
- this.ordProducer = new DateHistogramAggregatorBridge.DateHistoOrdProducer(
- getFieldType(),
- optimizationContext.getRanges(),
- bucketOrds,
- preparedRounding);
- }
+ public void prepare() throws IOException { buildRanges(context); }
 
  @Override
  protected Rounding getRounding(long low, long high) {
@@ -153,6 +147,19 @@ protected Rounding.Prepared getRoundingPrepared() {
  protected long[] processHardBounds(long[] bounds) {
  return super.processHardBounds(bounds, hardBounds);
  }
+
+ @Override
+ protected long getOrd(int rangeIdx){
+ long rangeStart = LongPoint.decodeDimension(optimizationContext.getRanges().getLower(rangeIdx), 0);
+ rangeStart = this.getFieldType().convertNanosToMillis(rangeStart);
+ long ord = bucketOrds.add(0, getRoundingPrepared().round(rangeStart));
+
+ if (ord < 0) { // already seen
+ ord = -1 - ord;
+ }
+
+ return ord;
+ }
  });
  if (optimizationContext.canOptimize(parent, context)) {
  optimizationContext.prepare();

server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -289,10 +289,7 @@ public boolean canOptimize() {
  }
 
  @Override
- public void prepare() {
- buildRanges(ranges);
- this.ordProducer = new RangeOrdProducer();
- }
+ public void prepare() { buildRanges(ranges); }
  });
  if (optimizationContext.canOptimize(parent, context)) {
  optimizationContext.prepare();

server/src/main/java/org/opensearch/search/optimization/filterrewrite/AggregatorBridge.java

@@ -41,20 +41,6 @@ public abstract class AggregatorBridge {
  */
  OptimizationContext optimizationContext;
 
- /**
- * Produce bucket ordinals from index of the corresponding range in the range array
- */
- public abstract static class OrdProducer {
- abstract long get(int idx);
- }
-
- protected OrdProducer ordProducer;
-
- public long getOrd(int rangeIdx) {
- return ordProducer.get(rangeIdx);
- }
-
-
  /**
  * The field type associated with this aggregator bridge.
  */
@@ -87,12 +73,19 @@ void setOptimizationContext(OptimizationContext context) {
 
  /**
  * @return max range to stop collecting at.
- * Utilized by aggs which stop early
+ * Utilized by aggs which stop early.
  */
  protected int rangeMax() {
  return Integer.MAX_VALUE;
  }
 
+ /**
+ * Translate an index of the packed value range array to an agg bucket ordinal.
+ */
+ protected long getOrd(int rangeIdx){
+ return rangeIdx;
+ }
+
  /**
  * Attempts to build aggregation results for a segment.
  * With no sub agg count docs and avoid iterating docIds.
@@ -104,6 +97,7 @@ protected int rangeMax() {
  public final void tryOptimize(PointValues values, BiConsumer<Long, Long> incrementDocCount, final LeafBucketCollector sub)
  throws IOException {
  TreeTraversal.RangeAwareIntersectVisitor treeVisitor;
+
  if (sub != null) {
  treeVisitor = new TreeTraversal.DocCollectRangeAwareIntersectVisitor(
  values.getPointTree(),

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

@@ -33,36 +33,6 @@
  * For date histogram aggregation
  */
 public abstract class DateHistogramAggregatorBridge extends AggregatorBridge {
-
- public static class DateHistoOrdProducer extends OrdProducer {
- DateFieldMapper.DateFieldType fieldType;
- PackedValueRanges ranges;
- LongKeyedBucketOrds bucketOrds;
- Rounding.Prepared rounding;
-
- public DateHistoOrdProducer(DateFieldMapper.DateFieldType fieldType,
- PackedValueRanges ranges,
- LongKeyedBucketOrds bucketOrds,
- Rounding.Prepared rounding) {
- this.fieldType = fieldType;
- this.ranges = ranges;
- this.bucketOrds = bucketOrds;
- this.rounding = rounding;
- }
-
- long get(int idx) {
- long rangeStart = LongPoint.decodeDimension(ranges.lowers[idx], 0);
- rangeStart = fieldType.convertNanosToMillis(rangeStart);
- long ord = bucketOrds.add(0, rounding.round((long) rangeStart));
-
- if (ord < 0) { // already seen
- ord = -1 - ord;
- }
-
- return ord;
- }
- }
-
  protected boolean canOptimize(ValuesSourceConfig config) {
  if (config.script() == null && config.missing() == null) {
  MappedFieldType fieldType = config.fieldType();

server/src/main/java/org/opensearch/search/optimization/filterrewrite/PackedValueRanges.java

@@ -42,6 +42,14 @@ public static boolean withinUpperBound(byte[] value, byte[] upperBound) {
  return compareByteValue(value, upperBound) < 0;
  }
 
+ public byte[] getLower(int idx){
+ return lowers[idx];
+ }
+
+ public byte[] getUpper(int idx){
+ return uppers[idx];
+ }
+
  public boolean withinLowerBound(byte[] value, int idx) {
  return PackedValueRanges.withinLowerBound(value, lowers[idx]);
  }

server/src/main/java/org/opensearch/search/optimization/filterrewrite/RangeAggregatorBridge.java

@@ -28,12 +28,6 @@
  */
 public abstract class RangeAggregatorBridge extends AggregatorBridge {
 
- public static class RangeOrdProducer extends OrdProducer {
- long get(int idx) {
- return idx;
- }
- }
-
  protected boolean canOptimize(ValuesSourceConfig config, RangeAggregator.Range[] ranges) {
  if (config.fieldType() == null) return false;
  MappedFieldType fieldType = config.fieldType();