Skip to content

Commit 8a7f820

Browse files
ivanceaivancea
authored
ESQL: Timezone support in DATE_TRUNC, BUCKET and TBUCKET (#137450)
Add timezone support to DATE_TRUNC, BUCKET and TBUCKET. The current timezone is already in the Configuration. This PR: - Uses the configuration ZoneId from those functions - Adds unit tests for edge cases - As BUCKET and TBUCKET depend on DATE_TRUNC, DATE_TRUNC tests are reused - Fixed Rounding for date intervals with variable timezones not taking the amount of units (e.g. 5 days) into account - Adds CSV tests for common cases
1 parent a5e5936 commit 8a7f820

File tree

35 files changed

+1911
-770
lines changed

35 files changed

+1911
-770
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,12 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
217217
);
218218
yield EvalMapper.toEvaluator(
219219
FOLD_CONTEXT,
220-
new DateTrunc(Source.EMPTY, new Literal(Source.EMPTY, Duration.ofHours(24), DataType.TIME_DURATION), timestamp),
220+
new DateTrunc(
221+
Source.EMPTY,
222+
new Literal(Source.EMPTY, Duration.ofHours(24), DataType.TIME_DURATION),
223+
timestamp,
224+
configuration()
225+
),
221226
layout(timestamp)
222227
).get(driverContext);
223228
}

‎docs/changelog/137450.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 137450
2+
summary: "Timezone support in DATE_TRUNC, BUCKET and TBUCKET"
3+
area: ES|QL
4+
type: feature
5+
issues: []

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

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.common.io.stream.StreamOutput;
1919
import org.elasticsearch.common.io.stream.Writeable;
2020
import org.elasticsearch.common.time.DateUtils;
21+
import org.elasticsearch.common.time.LocalDateTimeUtils;
2122
import org.elasticsearch.core.TimeValue;
2223

2324
import java.io.IOException;
@@ -29,7 +30,6 @@
2930
import java.time.ZoneId;
3031
import java.time.ZoneOffset;
3132
import java.time.temporal.ChronoField;
32-
import java.time.temporal.ChronoUnit;
3333
import java.time.temporal.IsoFields;
3434
import java.time.temporal.TemporalField;
3535
import java.time.temporal.TemporalQueries;
@@ -546,16 +546,16 @@ private LocalDateTime truncateLocalDateTime(LocalDateTime localDateTime) {
546546
return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonthValue(), 1, 0, 0);
547547

548548
case QUARTER_OF_YEAR:
549-
return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonth().firstMonthOfQuarter(), 1, 0, 0);
549+
return LocalDateTime.of(localDateTime.getYear(), (((localDateTime.getMonthValue() - 1) / 3) * 3) + 1, 1, 0, 0);
550550

551551
case YEAR_OF_CENTURY:
552552
return LocalDateTime.of(LocalDate.of(localDateTime.getYear(), 1, 1), LocalTime.MIDNIGHT);
553553

554554
case YEARS_OF_CENTURY:
555-
return LocalDateTime.of(LocalDate.of(localDateTime.getYear(), 1, 1), LocalTime.MIDNIGHT);
555+
return LocalDateTimeUtils.truncateToYears(localDateTime, multiplier);
556556

557557
case MONTHS_OF_YEAR:
558-
return LocalDateTime.of(localDateTime.getYear(), localDateTime.getMonthValue(), 1, 0, 0);
558+
return LocalDateTimeUtils.truncateToMonths(localDateTime, multiplier);
559559

560560
default:
561561
throw new IllegalArgumentException("NOT YET IMPLEMENTED for unit " + unit);
@@ -914,13 +914,11 @@ private LocalDateTime nextRelevantMidnight(LocalDateTime localMidnight) {
914914
assert localMidnight.toLocalTime().equals(LocalTime.MIDNIGHT) : "nextRelevantMidnight should only be called at midnight";
915915

916916
return switch (unit) {
917-
case DAY_OF_MONTH -> localMidnight.plus(1, ChronoUnit.DAYS);
918-
case WEEK_OF_WEEKYEAR -> localMidnight.plus(7, ChronoUnit.DAYS);
919-
case MONTH_OF_YEAR -> localMidnight.plus(1, ChronoUnit.MONTHS);
920-
case QUARTER_OF_YEAR -> localMidnight.plus(3, ChronoUnit.MONTHS);
921-
case YEAR_OF_CENTURY -> localMidnight.plus(1, ChronoUnit.YEARS);
922-
case YEARS_OF_CENTURY -> localMidnight.plus(1, ChronoUnit.YEARS);
923-
case MONTHS_OF_YEAR -> localMidnight.plus(1, ChronoUnit.MONTHS);
917+
case DAY_OF_MONTH -> localMidnight.plusDays(multiplier);
918+
case WEEK_OF_WEEKYEAR -> localMidnight.plusDays(7L * multiplier);
919+
case MONTH_OF_YEAR, MONTHS_OF_YEAR -> localMidnight.plusMonths(multiplier);
920+
case QUARTER_OF_YEAR -> localMidnight.plusMonths(3L * multiplier);
921+
case YEAR_OF_CENTURY, YEARS_OF_CENTURY -> localMidnight.plusYears(multiplier);
924922
default -> throw new IllegalArgumentException("Unknown round-to-midnight unit: " + unit);
925923
};
926924
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.common.time;
11+
12+
import java.time.LocalDate;
13+
import java.time.LocalDateTime;
14+
import java.time.LocalTime;
15+
16+
public class LocalDateTimeUtils {
17+
public static LocalDateTime truncateToMonths(LocalDateTime dateTime, int months) {
18+
int totalMonths = (dateTime.getYear() - 1) * 12 + dateTime.getMonthValue() - 1;
19+
int truncatedMonths = Math.floorDiv(totalMonths, months) * months;
20+
return LocalDateTime.of(LocalDate.of(truncatedMonths / 12 + 1, truncatedMonths % 12 + 1, 1), LocalTime.MIDNIGHT);
21+
}
22+
23+
public static LocalDateTime truncateToYears(LocalDateTime dateTime, int years) {
24+
int truncatedYear = Math.floorDiv(dateTime.getYear() - 1, years) * years + 1;
25+
return LocalDateTime.of(LocalDate.of(truncatedYear, 1, 1), LocalTime.MIDNIGHT);
26+
}
27+
}

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.common.time;
11+
12+
import org.elasticsearch.test.ESTestCase;
13+
14+
import java.time.LocalDate;
15+
import java.time.LocalDateTime;
16+
import java.time.LocalTime;
17+
import java.time.Month;
18+
import java.time.Year;
19+
20+
import static org.hamcrest.Matchers.equalTo;
21+
22+
public class LocalDateTimeUtilsTests extends ESTestCase {
23+
public void testTruncateToYears() {
24+
int randomYear = randomIntBetween(1, 3000);
25+
assertTruncateToYears(1, randomYear, randomYear);
26+
27+
assertTruncateToYears(10, 1, 1);
28+
assertTruncateToYears(10, 11, 11);
29+
assertTruncateToYears(10, 10, 1);
30+
assertTruncateToYears(10, 11, 11);
31+
assertTruncateToYears(10, 2015, 2011);
32+
33+
assertTruncateToYears(4, 2000, 1997);
34+
assertTruncateToYears(4, 2003, 2001);
35+
assertTruncateToYears(4, 2004, 2001);
36+
assertTruncateToYears(4, 2005, 2005);
37+
38+
assertTruncateToYears(4, 1, 1);
39+
assertTruncateToYears(4, -1, -3);
40+
assertTruncateToYears(4, -3, -3);
41+
assertTruncateToYears(4, -4, -7);
42+
}
43+
44+
private void assertTruncateToYears(int interval, int year, int expectedYear) {
45+
int inputMonth = randomIntBetween(1, 12);
46+
LocalDateTime inputDate = LocalDateTime.of(
47+
year,
48+
inputMonth,
49+
Month.of(inputMonth).length(Year.isLeap(year)),
50+
randomIntBetween(0, 23),
51+
randomIntBetween(0, 59),
52+
randomIntBetween(0, 59)
53+
);
54+
55+
LocalDateTime expectedResult = LocalDateTime.of(LocalDate.of(expectedYear, 1, 1), LocalTime.MIDNIGHT);
56+
57+
LocalDateTime resultYears = LocalDateTimeUtils.truncateToYears(inputDate, interval);
58+
assertThat(resultYears, equalTo(expectedResult));
59+
60+
// Also tests the same with months
61+
LocalDateTime resultMonths = LocalDateTimeUtils.truncateToMonths(inputDate, interval * 12);
62+
assertThat(resultMonths, equalTo(expectedResult));
63+
}
64+
65+
}

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

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.test;
11+
12+
import org.elasticsearch.common.time.DateFormatter;
13+
import org.elasticsearch.common.time.DateUtils;
14+
import org.hamcrest.Description;
15+
import org.hamcrest.TypeSafeMatcher;
16+
17+
import java.time.Instant;
18+
19+
public class ReadableMatchers {
20+
private static final DateFormatter dateFormatter = DateFormatter.forPattern("strict_date_optional_time");
21+
22+
/**
23+
* Test matcher for millis dates that expects longs, but describes the errors as dates, for better readability.
24+
* <p>
25+
* See {@link #matchesDateNanos} for the nanos counterpart.
26+
* </p>
27+
*/
28+
public static DateMillisMatcher matchesDateMillis(String date) {
29+
return new DateMillisMatcher(date);
30+
}
31+
32+
/**
33+
* Test matcher for nanos dates that expects longs, but describes the errors as dates, for better readability.
34+
* <p>
35+
* See {@link DateMillisMatcher} for the millis counterpart.
36+
* </p>
37+
*/
38+
public static DateNanosMatcher matchesDateNanos(String date) {
39+
return new DateNanosMatcher(date);
40+
}
41+
42+
public static class DateMillisMatcher extends TypeSafeMatcher<Long> {
43+
private final long timeMillis;
44+
45+
public DateMillisMatcher(String date) {
46+
this.timeMillis = Instant.parse(date).toEpochMilli();
47+
}
48+
49+
@Override
50+
public boolean matchesSafely(Long item) {
51+
return timeMillis == item;
52+
}
53+
54+
@Override
55+
public void describeMismatchSafely(Long item, Description description) {
56+
description.appendText("was ").appendValue(dateFormatter.formatMillis(item));
57+
}
58+
59+
@Override
60+
public void describeTo(Description description) {
61+
description.appendText(dateFormatter.formatMillis(timeMillis));
62+
}
63+
}
64+
65+
public static class DateNanosMatcher extends TypeSafeMatcher<Long> {
66+
private final long timeNanos;
67+
68+
public DateNanosMatcher(String date) {
69+
this.timeNanos = DateUtils.toLong(Instant.parse(date));
70+
}
71+
72+
@Override
73+
public boolean matchesSafely(Long item) {
74+
return timeNanos == item;
75+
}
76+
77+
@Override
78+
public void describeMismatchSafely(Long item, Description description) {
79+
description.appendText("was ").appendValue(dateFormatter.formatNanos(item));
80+
}
81+
82+
@Override
83+
public void describeTo(Description description) {
84+
description.appendText(dateFormatter.formatNanos(timeNanos));
85+
}
86+
}
87+
}

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

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import java.util.concurrent.ConcurrentHashMap;
5757
import java.util.concurrent.ConcurrentMap;
5858
import java.util.function.IntFunction;
59+
import java.util.stream.Stream;
5960

6061
import static java.util.Collections.emptySet;
6162
import static java.util.Map.entry;
@@ -74,6 +75,7 @@
7475
import static org.hamcrest.Matchers.emptyOrNullString;
7576
import static org.hamcrest.Matchers.equalTo;
7677
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
78+
import static org.hamcrest.Matchers.hasSize;
7779
import static org.hamcrest.Matchers.is;
7880
import static org.hamcrest.Matchers.not;
7981
import static org.hamcrest.Matchers.nullValue;
@@ -172,7 +174,7 @@ public RequestObjectBuilder params(String rawParams) throws IOException {
172174
}
173175

174176
public RequestObjectBuilder timeZone(ZoneId zoneId) throws IOException {
175-
builder.field("time_zone", zoneId);
177+
builder.field("time_zone", zoneId.toString());
176178
return this;
177179
}
178180

@@ -1852,6 +1854,73 @@ public void testMatchFunctionAcrossMultipleIndicesWithMissingField() throws IOEx
18521854
}
18531855
}
18541856

1857+
@SuppressWarnings("fallthrough")
1858+
public void testRandomTimezoneBuckets() throws IOException {
1859+
assumeTrue("timezone support for date_trunc is required", EsqlCapabilities.Cap.DATE_TRUNC_TIMEZONE_SUPPORT.isEnabled());
1860+
1861+
createIndex(testIndexName(), Settings.EMPTY, """
1862+
{
1863+
"properties": {
1864+
"@timestamp": {
1865+
"type": "date"
1866+
}
1867+
}
1868+
}
1869+
""");
1870+
bulkLoadTestData(randomIntBetween(1, 10), 0, false, i -> """
1871+
{"index":{"_id":"%s"}}
1872+
{"@timestamp": %s}
1873+
""".formatted(testIndexName(), randomLong()));
1874+
1875+
var timeZone = randomZone();
1876+
var interval = randomFrom("1 hour", "1 day", "1 month");
1877+
var functions = Stream.of("DATE_TRUNC(\"%s\", @timestamp)", "BUCKET(@timestamp, \"%s\")", "TBUCKET(\"%s\")")
1878+
.map(f -> f.formatted(interval))
1879+
.toList();
1880+
1881+
Object firstResultValues = null;
1882+
1883+
for (int timezoneSettingMethod = 0; timezoneSettingMethod < 3; timezoneSettingMethod++) {
1884+
for (var function : functions) {
1885+
var query = "FROM %s | STATS BY bucket=%s | SORT bucket".formatted(testIndexName(), function);
1886+
var builder = requestObjectBuilder();
1887+
1888+
switch (timezoneSettingMethod) {
1889+
case 0: // "time_zone" request param
1890+
builder.query(query).timeZone(timeZone);
1891+
break;
1892+
case 1: // Random "time_zone" request param with a SET overriding it
1893+
builder.timeZone(randomZone());
1894+
// Fall-through and set the actual timezone. This case only sets a random, to-be-overridden request timezone
1895+
case 2: // SET "time_zone" param
1896+
builder.query("SET time_zone=\"%s\"; %s".formatted(timeZone, query));
1897+
break;
1898+
}
1899+
1900+
var result = runEsql(requestObjectBuilder().query(query).timeZone(timeZone));
1901+
1902+
assertResultMap(
1903+
result,
1904+
getResultMatcher(result),
1905+
matchesList().item(matchesMap().entry("name", "bucket").entry("type", "date")),
1906+
hasSize(greaterThanOrEqualTo(1))
1907+
);
1908+
1909+
Object values = result.get("values");
1910+
1911+
if (firstResultValues == null) {
1912+
firstResultValues = values;
1913+
} else {
1914+
assertThat(
1915+
"function %s for timezone %s didn't return the same values".formatted(function, timeZone),
1916+
values,
1917+
equalTo(firstResultValues)
1918+
);
1919+
}
1920+
}
1921+
}
1922+
}
1923+
18551924
protected static Request prepareRequestWithOptions(RequestObjectBuilder requestObject, Mode mode) throws IOException {
18561925
requestObject.build();
18571926
Request request = prepareRequest(mode);

0 commit comments

Comments
 (0)