Handle utc timezone check

Signed-off-by: bowenlan-amzn <[email protected]>
bowenlan-amzn · Jan 22, 2025 · e0de5d0 · e0de5d0
1 parent 569502d
commit e0de5d0
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 17 deletions.
diff --git a/server/src/main/java/org/opensearch/common/Rounding.java b/server/src/main/java/org/opensearch/common/Rounding.java
@@ -663,6 +663,11 @@ public String toString() {
             return "Rounding[" + unit + " in " + timeZone + "]";
         }
 
+        @Override
+        public boolean isUTC() {
+            return "Z".equals(timeZone.getDisplayName(TextStyle.FULL, Locale.ENGLISH));
+        }
+
         private abstract class TimeUnitPreparedRounding extends PreparedRounding {
             @Override
             public double roundingSize(long utcMillis, DateTimeUnit timeUnit) {
@@ -1040,6 +1045,11 @@ public String toString() {
             return "Rounding[" + interval + " in " + timeZone + "]";
         }
 
+        @Override
+        public boolean isUTC() {
+            return "Z".equals(timeZone.getDisplayName(TextStyle.FULL, Locale.ENGLISH));
+        }
+
         private long roundKey(long value, long interval) {
             if (value < 0) {
                 return (value - interval + 1) / interval;
@@ -1359,6 +1369,11 @@ public boolean equals(Object obj) {
         public String toString() {
             return delegate + " offset by " + offset;
         }
+
+        @Override
+        public boolean isUTC() {
+            return delegate.isUTC();
+        }
     }
 
     public static Rounding read(StreamInput in) throws IOException {
@@ -1399,7 +1414,5 @@ public static OptionalLong getInterval(Rounding rounding) {
      * Helper function for checking if the time zone requested for date histogram
      * aggregation is utc or not
      */
-    public static boolean isUTCTimeZone(final ZoneId zoneId) {
-        return "Z".equals(zoneId.getDisplayName(TextStyle.FULL, Locale.ENGLISH));
-    }
+    public abstract boolean isUTC();
 }
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
@@ -82,7 +82,6 @@
 import org.opensearch.search.sort.SortAndFormats;
 
 import java.io.IOException;
-import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -94,7 +93,6 @@
 import java.util.function.LongUnaryOperator;
 import java.util.stream.Collectors;
 
-import static org.opensearch.common.Rounding.isUTCTimeZone;
 import static org.opensearch.search.aggregations.MultiBucketConsumerService.MAX_BUCKET_SETTING;
 import static org.opensearch.search.aggregations.bucket.filterrewrite.DateHistogramAggregatorBridge.segmentMatchAll;
 
@@ -188,12 +186,8 @@ protected boolean canOptimize() {
                      * The filter rewrite optimized path does not support bucket intervals which are not fixed.
                      * For this reason we exclude non UTC timezones.
                      */
-                    Rounding round = this.valuesSource.getRounding();
-                    if (round instanceof Rounding.TimeUnitRounding) {
-                        ZoneId tz = ((Rounding.TimeUnitRounding) round).getTimeZone();
-                        if (!isUTCTimeZone(tz)) {
-                            return false;
-                        }
+                    if (valuesSource.getRounding().isUTC() == false) {
+                        return false;
                     }
 
                     // bucketOrds is used for saving the date histogram results got from the optimization path

diff --git a/...rg/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java b/...rg/opensearch/search/aggregations/bucket/filterrewrite/DateHistogramAggregatorBridge.java
@@ -14,7 +14,6 @@
 import org.opensearch.common.Rounding;
 import org.opensearch.index.mapper.DateFieldMapper;
 import org.opensearch.index.mapper.MappedFieldType;
-import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.aggregations.bucket.histogram.LongBounds;
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
@@ -24,7 +23,6 @@
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 
-import static org.opensearch.common.Rounding.isUTCTimeZone;
 import static org.opensearch.search.aggregations.bucket.filterrewrite.PointTreeTraversal.multiRangesTraverse;
 
 /**
@@ -34,12 +32,12 @@ public abstract class DateHistogramAggregatorBridge extends AggregatorBridge {
 
     int maxRewriteFilters;
 
-    protected boolean canOptimize(ValuesSourceConfig config) {
+    protected boolean canOptimize(ValuesSourceConfig config, Rounding rounding) {
         /**
          * The filter rewrite optimized path does not support bucket intervals which are not fixed.
          * For this reason we exclude non UTC timezones.
          */
-        if (config.format() instanceof DocValueFormat.DateTime && !isUTCTimeZone(((DocValueFormat.DateTime) config.format()).getZoneId())) {
+        if (rounding.isUTC() == false) {
             return false;
         }
 

diff --git a/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/...java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java
@@ -161,7 +161,7 @@ private AutoDateHistogramAggregator(
         DateHistogramAggregatorBridge bridge = new DateHistogramAggregatorBridge() {
             @Override
             protected boolean canOptimize() {
-                return canOptimize(valuesSourceConfig);
+                return canOptimize(valuesSourceConfig, roundingInfos[0].rounding);
             }
 
             @Override

diff --git a/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/...ain/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java
@@ -122,7 +122,7 @@ class DateHistogramAggregator extends BucketsAggregator implements SizedBucketAg
         DateHistogramAggregatorBridge bridge = new DateHistogramAggregatorBridge() {
             @Override
             protected boolean canOptimize() {
-                return canOptimize(valuesSourceConfig);
+                return canOptimize(valuesSourceConfig, rounding);
             }
 
             @Override