ESQL: Timezone support in DATE_TRUNC, BUCKET and TBUCKET

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java

@@ -217,7 +217,12 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                 );
                 yield EvalMapper.toEvaluator(
                     FOLD_CONTEXT,
-                    new DateTrunc(Source.EMPTY, new Literal(Source.EMPTY, Duration.ofHours(24), DataType.TIME_DURATION), timestamp),
+                    new DateTrunc(
+                        Source.EMPTY,
+                        new Literal(Source.EMPTY, Duration.ofHours(24), DataType.TIME_DURATION),
+                        timestamp,
+                        configuration()
+                    ),
                     layout(timestamp)
                 ).get(driverContext);
             }

docs/changelog/137450.yaml

@@ -0,0 +1,5 @@
+pr: 137450
+summary: "Timezone support in DATE_TRUNC, BUCKET and TBUCKET"
+area: ES|QL
+type: feature
+issues: []

server/src/main/java/org/elasticsearch/common/Rounding.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.time.DateUtils;
+import org.elasticsearch.common.time.LocalDateTimeUtils;
 import org.elasticsearch.core.TimeValue;
 
 import java.io.IOException;
@@ -29,7 +30,6 @@
 import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.time.temporal.ChronoField;
-import java.time.temporal.ChronoUnit;
 import java.time.temporal.IsoFields;
 import java.time.temporal.TemporalField;
 import java.time.temporal.TemporalQueries;
@@ -546,16 +546,16 @@ private LocalDateTime truncateLocalDateTime(LocalDateTime localDateTime) {
                     return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonthValue(), 1, 0, 0);
 
                 case QUARTER_OF_YEAR:
-                    return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonth().firstMonthOfQuarter(), 1, 0, 0);
+                    return LocalDateTime.of(localDateTime.getYear(), (((localDateTime.getMonthValue() - 1) / 3) * 3) + 1, 1, 0, 0);
 
                 case YEAR_OF_CENTURY:
                     return LocalDateTime.of(LocalDate.of(localDateTime.getYear(), 1, 1), LocalTime.MIDNIGHT);
 
                 case YEARS_OF_CENTURY:
-                    return LocalDateTime.of(LocalDate.of(localDateTime.getYear(), 1, 1), LocalTime.MIDNIGHT);
+                    return LocalDateTimeUtils.truncateToYears(localDateTime, multiplier);
 
                 case MONTHS_OF_YEAR:
-                    return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonthValue(), 1, 0, 0);
+                    return LocalDateTimeUtils.truncateToMonths(localDateTime, multiplier);
 
                 default:
                     throw new IllegalArgumentException("NOT YET IMPLEMENTED for unit " + unit);
@@ -914,13 +914,11 @@ private LocalDateTime nextRelevantMidnight(LocalDateTime localMidnight) {
                 assert localMidnight.toLocalTime().equals(LocalTime.MIDNIGHT) : "nextRelevantMidnight should only be called at midnight";
 
                 return switch (unit) {
-                    case DAY_OF_MONTH -> localMidnight.plus(1, ChronoUnit.DAYS);
-                    case WEEK_OF_WEEKYEAR -> localMidnight.plus(7, ChronoUnit.DAYS);
-                    case MONTH_OF_YEAR -> localMidnight.plus(1, ChronoUnit.MONTHS);
-                    case QUARTER_OF_YEAR -> localMidnight.plus(3, ChronoUnit.MONTHS);
-                    case YEAR_OF_CENTURY -> localMidnight.plus(1, ChronoUnit.YEARS);
-                    case YEARS_OF_CENTURY -> localMidnight.plus(1, ChronoUnit.YEARS);
-                    case MONTHS_OF_YEAR -> localMidnight.plus(1, ChronoUnit.MONTHS);
+                    case DAY_OF_MONTH -> localMidnight.plusDays(multiplier);
+                    case WEEK_OF_WEEKYEAR -> localMidnight.plusDays(7L * multiplier);
+                    case MONTH_OF_YEAR, MONTHS_OF_YEAR -> localMidnight.plusMonths(multiplier);
+                    case QUARTER_OF_YEAR -> localMidnight.plusMonths(3L * multiplier);
+                    case YEAR_OF_CENTURY, YEARS_OF_CENTURY -> localMidnight.plusYears(multiplier);
                     default -> throw new IllegalArgumentException("Unknown round-to-midnight unit: " + unit);
                 };
             }

server/src/main/java/org/elasticsearch/common/time/LocalDateTimeUtils.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.time;
+
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+
+public class LocalDateTimeUtils {
+    public static LocalDateTime truncateToMonths(LocalDateTime dateTime, int months) {
+        int totalMonths = (dateTime.getYear() - 1) * 12 + dateTime.getMonthValue() - 1;
+        int truncatedMonths = Math.floorDiv(totalMonths, months) * months;
+        return LocalDateTime.of(LocalDate.of(truncatedMonths / 12 + 1, truncatedMonths % 12 + 1, 1), LocalTime.MIDNIGHT);
+    }
+
+    public static LocalDateTime truncateToYears(LocalDateTime dateTime, int years) {
+        int truncatedYear = Math.floorDiv(dateTime.getYear() - 1, years) * years + 1;
+        return LocalDateTime.of(LocalDate.of(truncatedYear, 1, 1), LocalTime.MIDNIGHT);
+    }
+}

server/src/test/java/org/elasticsearch/common/time/LocalDateTimeUtilsTests.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.time;
+
+import org.elasticsearch.test.ESTestCase;
+
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.Month;
+import java.time.Year;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class LocalDateTimeUtilsTests extends ESTestCase {
+    public void testTruncateToYears() {
+        int randomYear = randomIntBetween(1, 3000);
+        assertTruncateToYears(1, randomYear, randomYear);
+
+        assertTruncateToYears(10, 1, 1);
+        assertTruncateToYears(10, 11, 11);
+        assertTruncateToYears(10, 10, 1);
+        assertTruncateToYears(10, 11, 11);
+        assertTruncateToYears(10, 2015, 2011);
+
+        assertTruncateToYears(4, 2000, 1997);
+        assertTruncateToYears(4, 2003, 2001);
+        assertTruncateToYears(4, 2004, 2001);
+        assertTruncateToYears(4, 2005, 2005);
+
+        assertTruncateToYears(4, 1, 1);
+        assertTruncateToYears(4, -1, -3);
+        assertTruncateToYears(4, -3, -3);
+        assertTruncateToYears(4, -4, -7);
+    }
+
+    private void assertTruncateToYears(int interval, int year, int expectedYear) {
+        int inputMonth = randomIntBetween(1, 12);
+        LocalDateTime inputDate = LocalDateTime.of(
+            year,
+            inputMonth,
+            Month.of(inputMonth).length(Year.isLeap(year)),
+            randomIntBetween(0, 23),
+            randomIntBetween(0, 59),
+            randomIntBetween(0, 59)
+        );
+
+        LocalDateTime expectedResult = LocalDateTime.of(LocalDate.of(expectedYear, 1, 1), LocalTime.MIDNIGHT);
+
+        LocalDateTime resultYears = LocalDateTimeUtils.truncateToYears(inputDate, interval);
+        assertThat(resultYears, equalTo(expectedResult));
+
+        // Also tests the same with months
+        LocalDateTime resultMonths = LocalDateTimeUtils.truncateToMonths(inputDate, interval * 12);
+        assertThat(resultMonths, equalTo(expectedResult));
+    }
+
+}

test/framework/src/main/java/org/elasticsearch/test/ReadableMatchers.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test;
+
+import org.elasticsearch.common.time.DateFormatter;
+import org.elasticsearch.common.time.DateUtils;
+import org.hamcrest.Description;
+import org.hamcrest.TypeSafeMatcher;
+
+import java.time.Instant;
+
+public class ReadableMatchers {
+    private static final DateFormatter dateFormatter = DateFormatter.forPattern("strict_date_optional_time");
+
+    /**
+     * Test matcher for millis dates that expects longs, but describes the errors as dates, for better readability.
+     * <p>
+     *     See {@link #matchesDateNanos} for the nanos counterpart.
+     * </p>
+     */
+    public static DateMillisMatcher matchesDateMillis(String date) {
+        return new DateMillisMatcher(date);
+    }
+
+    /**
+     * Test matcher for nanos dates that expects longs, but describes the errors as dates, for better readability.
+     * <p>
+     *     See {@link DateMillisMatcher} for the millis counterpart.
+     * </p>
+     */
+    public static DateNanosMatcher matchesDateNanos(String date) {
+        return new DateNanosMatcher(date);
+    }
+
+    public static class DateMillisMatcher extends TypeSafeMatcher<Long> {
+        private final long timeMillis;
+
+        public DateMillisMatcher(String date) {
+            this.timeMillis = Instant.parse(date).toEpochMilli();
+        }
+
+        @Override
+        public boolean matchesSafely(Long item) {
+            return timeMillis == item;
+        }
+
+        @Override
+        public void describeMismatchSafely(Long item, Description description) {
+            description.appendText("was ").appendValue(dateFormatter.formatMillis(item));
+        }
+
+        @Override
+        public void describeTo(Description description) {
+            description.appendText(dateFormatter.formatMillis(timeMillis));
+        }
+    }
+
+    public static class DateNanosMatcher extends TypeSafeMatcher<Long> {
+        private final long timeNanos;
+
+        public DateNanosMatcher(String date) {
+            this.timeNanos = DateUtils.toLong(Instant.parse(date));
+        }
+
+        @Override
+        public boolean matchesSafely(Long item) {
+            return timeNanos == item;
+        }
+
+        @Override
+        public void describeMismatchSafely(Long item, Description description) {
+            description.appendText("was ").appendValue(dateFormatter.formatNanos(item));
+        }
+
+        @Override
+        public void describeTo(Description description) {
+            description.appendText(dateFormatter.formatNanos(timeNanos));
+        }
+    }
+}

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

@@ -56,6 +56,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.IntFunction;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptySet;
 import static java.util.Map.entry;
@@ -74,6 +75,7 @@
 import static org.hamcrest.Matchers.emptyOrNullString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.nullValue;
@@ -172,7 +174,7 @@ public RequestObjectBuilder params(String rawParams) throws IOException {
         }
 
         public RequestObjectBuilder timeZone(ZoneId zoneId) throws IOException {
-            builder.field("time_zone", zoneId);
+            builder.field("time_zone", zoneId.toString());
             return this;
         }
 
@@ -1852,6 +1854,73 @@ public void testMatchFunctionAcrossMultipleIndicesWithMissingField() throws IOEx
         }
     }
 
+    @SuppressWarnings("fallthrough")
+    public void testRandomTimezoneBuckets() throws IOException {
+        assumeTrue("timezone support for date_trunc is required", EsqlCapabilities.Cap.DATE_TRUNC_TIMEZONE_SUPPORT.isEnabled());
+
+        createIndex(testIndexName(), Settings.EMPTY, """
+            {
+                "properties": {
+                    "@timestamp": {
+                        "type": "date"
+                    }
+                }
+            }
+            """);
+        bulkLoadTestData(randomIntBetween(1, 10), 0, false, i -> """
+            {"index":{"_id":"%s"}}
+            {"@timestamp": %s}
+            """.formatted(testIndexName(), randomLong()));
+
+        var timeZone = randomZone();
+        var interval = randomFrom("1 hour", "1 day", "1 month");
+        var functions = Stream.of("DATE_TRUNC(\"%s\", @timestamp)", "BUCKET(@timestamp, \"%s\")", "TBUCKET(\"%s\")")
+            .map(f -> f.formatted(interval))
+            .toList();
+
+        Object firstResultValues = null;
+
+        for (int timezoneSettingMethod = 0; timezoneSettingMethod < 3; timezoneSettingMethod++) {
+            for (var function : functions) {
+                var query = "FROM %s | STATS BY bucket=%s | SORT bucket".formatted(testIndexName(), function);
+                var builder = requestObjectBuilder();
+
+                switch (timezoneSettingMethod) {
+                    case 0: // "time_zone" request param
+                        builder.query(query).timeZone(timeZone);
+                        break;
+                    case 1: // Random "time_zone" request param with a SET overriding it
+                        builder.timeZone(randomZone());
+                        // Fall-through and set the actual timezone. This case only sets a random, to-be-overridden request timezone
+                    case 2: // SET "time_zone" param
+                        builder.query("SET time_zone=\"%s\"; %s".formatted(timeZone, query));
+                        break;
+                }
+
+                var result = runEsql(requestObjectBuilder().query(query).timeZone(timeZone));
+
+                assertResultMap(
+                    result,
+                    getResultMatcher(result),
+                    matchesList().item(matchesMap().entry("name", "bucket").entry("type", "date")),
+                    hasSize(greaterThanOrEqualTo(1))
+                );
+
+                Object values = result.get("values");
+
+                if (firstResultValues == null) {
+                    firstResultValues = values;
+                } else {
+                    assertThat(
+                        "function %s for timezone %s didn't return the same values".formatted(function, timeZone),
+                        values,
+                        equalTo(firstResultValues)
+                    );
+                }
+            }
+        }
+    }
+
     protected static Request prepareRequestWithOptions(RequestObjectBuilder requestObject, Mode mode) throws IOException {
         requestObject.build();
         Request request = prepareRequest(mode);