From ccc4d7b7698275d839e62f8a028267cd14757c75 Mon Sep 17 00:00:00 2001 From: Finn Carroll Date: Fri, 9 Aug 2024 10:52:05 -0700 Subject: [PATCH] Move ordProducer delegate to agg bridge function. Ranges need to be produced dynamically to support auto date histograms. Signed-off-by: Finn Carroll --- .../bucket/composite/CompositeAggregator.java | 23 +++++++++----- .../AutoDateHistogramAggregator.java | 23 +++++++++----- .../histogram/DateHistogramAggregator.java | 23 +++++++++----- .../bucket/range/RangeAggregator.java | 5 +--- .../filterrewrite/AggregatorBridge.java | 24 ++++++--------- .../DateHistogramAggregatorBridge.java | 30 ------------------- .../filterrewrite/PackedValueRanges.java | 8 +++++ .../filterrewrite/RangeAggregatorBridge.java | 6 ---- 8 files changed, 63 insertions(+), 79 deletions(-) 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 d803bca5d47f5..c8ecb0d9dc68c 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 @@ -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(); diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java index bc0db49ba1cb0..ee036c8c837aa 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java +++ b/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(); diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java index b9877a24b7d14..99d8d44393ecf 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java +++ b/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(); 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 536c19c1e95a8..74ce313803c8b 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 @@ -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(); diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/AggregatorBridge.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/AggregatorBridge.java index c0e9d1254482f..ed8c86747698c 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/AggregatorBridge.java +++ b/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 incrementDocCount, final LeafBucketCollector sub) throws IOException { TreeTraversal.RangeAwareIntersectVisitor treeVisitor; + if (sub != null) { treeVisitor = new TreeTraversal.DocCollectRangeAwareIntersectVisitor( values.getPointTree(), diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java index 3330cb971e78d..1e7008134ad32 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/DateHistogramAggregatorBridge.java +++ b/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(); diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/PackedValueRanges.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/PackedValueRanges.java index bfee822d1ce85..aec893aff6e99 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/PackedValueRanges.java +++ b/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]); } diff --git a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/RangeAggregatorBridge.java b/server/src/main/java/org/opensearch/search/optimization/filterrewrite/RangeAggregatorBridge.java index 639f7de33857c..2adddbbf535b1 100644 --- a/server/src/main/java/org/opensearch/search/optimization/filterrewrite/RangeAggregatorBridge.java +++ b/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();