Skip to content

Commit f0f68e0

Browse files
matrivmatriv
committed
SQL: Fix issue with date columns returned always in UTC
When selecting columns of ES type `date` (SQL's DATETIME) the `FieldHitExtractor` was not using the timezone of the client session but always resorted to UTC. The same behaviour (UTC only) was encountered also for grouping keys (`CompositeKeyExtractor`) and for First/Last functions on dates (`TopHitsAggExtractor`). Fixes: elastic#40152
1 parent 97707c7 commit f0f68e0

File tree

16 files changed

+190
-168
lines changed

16 files changed

+190
-168
lines changed

‎x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/JdbcTestUtils.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public static void logLikeCLI(ResultSet rs, Logger logger) throws SQLException {
137137
logger.info("\n" + formatter.formatWithHeader(cols, data));
138138
}
139139

140-
public static String of(long millis) {
141-
return StringUtils.toString(ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), UTC));
140+
public static String of(long millis, String zoneId) {
141+
return StringUtils.toString(ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneId.of(zoneId)));
142142
}
143-
}
143+
}

‎x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/ResultSetTestCase.java

Lines changed: 52 additions & 50 deletions
Large diffs are not rendered by default.

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/Querier.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -365,7 +365,6 @@ static class CompositeActionListener extends BaseAggActionListener {
365365
super(listener, client, cfg, output, query, request);
366366
}
367367

368-
369368
@Override
370369
protected void handleResponse(SearchResponse response, ActionListener<SchemaRowSet> listener) {
371370
// there are some results
@@ -428,7 +427,7 @@ protected List<BucketExtractor> initBucketExtractors(SearchResponse response) {
428427
private BucketExtractor createExtractor(FieldExtraction ref, BucketExtractor totalCount) {
429428
if (ref instanceof GroupByRef) {
430429
GroupByRef r = (GroupByRef) ref;
431-
return new CompositeKeyExtractor(r.key(), r.property(), r.zoneId());
430+
return new CompositeKeyExtractor(r.key(), r.property(), cfg.zoneId(), r.isDateTimeBased());
432431
}
433432

434433
if (ref instanceof MetricAggRef) {
@@ -438,7 +437,7 @@ private BucketExtractor createExtractor(FieldExtraction ref, BucketExtractor tot
438437

439438
if (ref instanceof TopHitsAggRef) {
440439
TopHitsAggRef r = (TopHitsAggRef) ref;
441-
return new TopHitsAggExtractor(r.name(), r.fieldDataType());
440+
return new TopHitsAggExtractor(r.name(), r.fieldDataType(), cfg.zoneId());
442441
}
443442

444443
if (ref == GlobalCountRef.INSTANCE) {
@@ -518,12 +517,13 @@ protected void handleResponse(SearchResponse response, ActionListener<SchemaRowS
518517
private HitExtractor createExtractor(FieldExtraction ref) {
519518
if (ref instanceof SearchHitFieldRef) {
520519
SearchHitFieldRef f = (SearchHitFieldRef) ref;
521-
return new FieldHitExtractor(f.name(), f.getDataType(), f.useDocValue(), f.hitName(), multiValueFieldLeniency);
520+
return new FieldHitExtractor(f.name(), f.getDataType(), cfg.zoneId(),
521+
f.useDocValue(), f.hitName(), multiValueFieldLeniency);
522522
}
523523

524524
if (ref instanceof ScriptFieldRef) {
525525
ScriptFieldRef f = (ScriptFieldRef) ref;
526-
return new FieldHitExtractor(f.name(), null, true, multiValueFieldLeniency);
526+
return new FieldHitExtractor(f.name(), null, cfg.zoneId(), true, multiValueFieldLeniency);
527527
}
528528

529529
if (ref instanceof ComputedRef) {

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/extractor/CompositeKeyExtractor.java

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -20,44 +20,38 @@
2020
public class CompositeKeyExtractor implements BucketExtractor {
2121

2222
/**
23-
* Key or Komposite extractor.
23+
* Key or Composite extractor.
2424
*/
2525
static final String NAME = "k";
2626

2727
private final String key;
2828
private final Property property;
2929
private final ZoneId zoneId;
30+
private final boolean isDateTimeBased;
3031

3132
/**
3233
* Constructs a new <code>CompositeKeyExtractor</code> instance.
33-
* The time-zone parameter is used to indicate a date key.
3434
*/
35-
public CompositeKeyExtractor(String key, Property property, ZoneId zoneId) {
35+
public CompositeKeyExtractor(String key, Property property, ZoneId zoneId, boolean isDateTimeBased) {
3636
this.key = key;
3737
this.property = property;
3838
this.zoneId = zoneId;
39+
this.isDateTimeBased = isDateTimeBased;
3940
}
4041

4142
CompositeKeyExtractor(StreamInput in) throws IOException {
4243
key = in.readString();
4344
property = in.readEnum(Property.class);
44-
if (in.readBoolean()) {
45-
zoneId = ZoneId.of(in.readString());
46-
} else {
47-
zoneId = null;
48-
}
45+
zoneId = ZoneId.of(in.readString());
46+
isDateTimeBased = in.readBoolean();
4947
}
5048

5149
@Override
5250
public void writeTo(StreamOutput out) throws IOException {
5351
out.writeString(key);
5452
out.writeEnum(property);
55-
if (zoneId == null) {
56-
out.writeBoolean(false);
57-
} else {
58-
out.writeBoolean(true);
59-
out.writeString(zoneId.getId());
60-
}
53+
out.writeString(zoneId.getId());
54+
out.writeBoolean(isDateTimeBased);
6155
}
6256

6357
String key() {
@@ -72,6 +66,10 @@ ZoneId zoneId() {
7266
return zoneId;
7367
}
7468

69+
public boolean isDateTimeBased() {
70+
return isDateTimeBased;
71+
}
72+
7573
@Override
7674
public String getWriteableName() {
7775
return NAME;
@@ -91,7 +89,7 @@ public Object extract(Bucket bucket) {
9189

9290
Object object = ((Map<?, ?>) m).get(key);
9391

94-
if (zoneId != null) {
92+
if (isDateTimeBased) {
9593
if (object == null) {
9694
return object;
9795
} else if (object instanceof Long) {
@@ -106,7 +104,7 @@ public Object extract(Bucket bucket) {
106104

107105
@Override
108106
public int hashCode() {
109-
return Objects.hash(key, property, zoneId);
107+
return Objects.hash(key, property, zoneId, isDateTimeBased);
110108
}
111109

112110
@Override
@@ -122,7 +120,8 @@ public boolean equals(Object obj) {
122120
CompositeKeyExtractor other = (CompositeKeyExtractor) obj;
123121
return Objects.equals(key, other.key)
124122
&& Objects.equals(property, other.property)
125-
&& Objects.equals(zoneId, other.zoneId);
123+
&& Objects.equals(zoneId, other.zoneId)
124+
&& Objects.equals(isDateTimeBased, other.isDateTimeBased);
126125
}
127126

128127
@Override

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/extractor/FieldHitExtractor.java

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.xpack.sql.util.DateUtils;
1717

1818
import java.io.IOException;
19+
import java.time.ZoneId;
1920
import java.util.ArrayDeque;
2021
import java.util.Deque;
2122
import java.util.List;
@@ -45,21 +46,23 @@ private static String[] sourcePath(String name, boolean useDocValue, String hitN
4546

4647
private final String fieldName, hitName;
4748
private final DataType dataType;
49+
private final ZoneId zoneId;
4850
private final boolean useDocValue;
4951
private final boolean arrayLeniency;
5052
private final String[] path;
5153

52-
public FieldHitExtractor(String name, DataType dataType, boolean useDocValue) {
53-
this(name, dataType, useDocValue, null, false);
54+
public FieldHitExtractor(String name, DataType dataType, ZoneId zoneId, boolean useDocValue) {
55+
this(name, dataType, zoneId, useDocValue, null, false);
5456
}
5557

56-
public FieldHitExtractor(String name, DataType dataType, boolean useDocValue, boolean arrayLeniency) {
57-
this(name, dataType, useDocValue, null, arrayLeniency);
58+
public FieldHitExtractor(String name, DataType dataType, ZoneId zoneId, boolean useDocValue, boolean arrayLeniency) {
59+
this(name, dataType, zoneId, useDocValue, null, arrayLeniency);
5860
}
5961

60-
public FieldHitExtractor(String name, DataType dataType, boolean useDocValue, String hitName, boolean arrayLeniency) {
62+
public FieldHitExtractor(String name, DataType dataType, ZoneId zoneId, boolean useDocValue, String hitName, boolean arrayLeniency) {
6163
this.fieldName = name;
6264
this.dataType = dataType;
65+
this.zoneId = zoneId;
6366
this.useDocValue = useDocValue;
6467
this.arrayLeniency = arrayLeniency;
6568
this.hitName = hitName;
@@ -77,6 +80,7 @@ public FieldHitExtractor(String name, DataType dataType, boolean useDocValue, St
7780
fieldName = in.readString();
7881
String esType = in.readOptionalString();
7982
dataType = esType != null ? DataType.fromTypeName(esType) : null;
83+
zoneId = ZoneId.of(in.readString());
8084
useDocValue = in.readBoolean();
8185
hitName = in.readOptionalString();
8286
arrayLeniency = in.readBoolean();
@@ -92,6 +96,7 @@ public String getWriteableName() {
9296
public void writeTo(StreamOutput out) throws IOException {
9397
out.writeString(fieldName);
9498
out.writeOptionalString(dataType == null ? null : dataType.typeName);
99+
out.writeString(zoneId.getId());
95100
out.writeBoolean(useDocValue);
96101
out.writeOptionalString(hitName);
97102
out.writeBoolean(arrayLeniency);
@@ -135,7 +140,7 @@ private Object unwrapMultiValue(Object values) {
135140
}
136141
if (dataType == DataType.DATETIME) {
137142
if (values instanceof String) {
138-
return DateUtils.asDateTime(Long.parseLong(values.toString()));
143+
return DateUtils.asDateTime(Long.parseLong(values.toString()), zoneId);
139144
}
140145
}
141146
if (values instanceof Long || values instanceof Double || values instanceof String || values instanceof Boolean) {
@@ -215,9 +220,13 @@ public String fieldName() {
215220
return fieldName;
216221
}
217222

223+
public ZoneId zoneId() {
224+
return zoneId;
225+
}
226+
218227
@Override
219228
public String toString() {
220-
return fieldName + "@" + hitName;
229+
return fieldName + "@" + hitName + "@" + zoneId;
221230
}
222231

223232
@Override

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/extractor/TopHitsAggExtractor.java

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.elasticsearch.xpack.sql.util.DateUtils;
1515

1616
import java.io.IOException;
17+
import java.time.ZoneId;
1718
import java.util.Objects;
1819

1920
public class TopHitsAggExtractor implements BucketExtractor {
@@ -22,21 +23,25 @@ public class TopHitsAggExtractor implements BucketExtractor {
2223

2324
private final String name;
2425
private final DataType fieldDataType;
26+
private final ZoneId zoneId;
2527

26-
public TopHitsAggExtractor(String name, DataType fieldDataType) {
28+
public TopHitsAggExtractor(String name, DataType fieldDataType, ZoneId zoneId) {
2729
this.name = name;
2830
this.fieldDataType = fieldDataType;
31+
this.zoneId = zoneId;
2932
}
3033

3134
TopHitsAggExtractor(StreamInput in) throws IOException {
3235
name = in.readString();
3336
fieldDataType = in.readEnum(DataType.class);
37+
zoneId = ZoneId.of(in.readString());
3438
}
3539

3640
@Override
3741
public void writeTo(StreamOutput out) throws IOException {
3842
out.writeString(name);
3943
out.writeEnum(fieldDataType);
44+
out.writeString(zoneId.getId());
4045
}
4146

4247
String name() {
@@ -61,15 +66,15 @@ public Object extract(Bucket bucket) {
6166

6267
Object value = agg.getHits().getAt(0).getFields().values().iterator().next().getValue();
6368
if (fieldDataType.isDateBased()) {
64-
return DateUtils.asDateTime(Long.parseLong(value.toString()));
69+
return DateUtils.asDateTime(Long.parseLong(value.toString()), zoneId);
6570
} else {
6671
return value;
6772
}
6873
}
6974

7075
@Override
7176
public int hashCode() {
72-
return Objects.hash(name, fieldDataType);
77+
return Objects.hash(name, fieldDataType, zoneId);
7378
}
7479

7580
@Override
@@ -84,11 +89,12 @@ public boolean equals(Object obj) {
8489

8590
TopHitsAggExtractor other = (TopHitsAggExtractor) obj;
8691
return Objects.equals(name, other.name)
87-
&& Objects.equals(fieldDataType, other.fieldDataType);
92+
&& Objects.equals(fieldDataType, other.fieldDataType)
93+
&& Objects.equals(zoneId, other.zoneId);
8894
}
8995

9096
@Override
9197
public String toString() {
92-
return "TopHits>" + name + "[" + fieldDataType + "]";
98+
return "TopHits>" + name + "[" + fieldDataType + "]@" + zoneId;
9399
}
94100
}

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/Planner.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,14 @@
55
*/
66
package org.elasticsearch.xpack.sql.planner;
77

8-
import java.util.List;
9-
import java.util.Map;
10-
118
import org.elasticsearch.xpack.sql.plan.logical.LogicalPlan;
129
import org.elasticsearch.xpack.sql.plan.physical.PhysicalPlan;
1310
import org.elasticsearch.xpack.sql.planner.Verifier.Failure;
1411
import org.elasticsearch.xpack.sql.tree.Node;
1512

13+
import java.util.List;
14+
import java.util.Map;
15+
1616
import static java.util.stream.Collectors.toMap;
1717

1818
public class Planner {
@@ -64,4 +64,4 @@ public Map<Node<?>, String> verifyExecutingPlanFailures(PhysicalPlan plan) {
6464
List<Failure> failures = Verifier.verifyExecutingPlan(plan);
6565
return failures.stream().collect(toMap(Failure::source, Failure::message));
6666
}
67-
}
67+
}

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.elasticsearch.xpack.sql.expression.function.grouping.GroupingFunction;
2828
import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunction;
2929
import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunctionAttribute;
30-
import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeFunction;
3130
import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeHistogramFunction;
3231
import org.elasticsearch.xpack.sql.expression.gen.pipeline.AggPathInput;
3332
import org.elasticsearch.xpack.sql.expression.gen.pipeline.Pipe;
@@ -64,9 +63,7 @@
6463
import org.elasticsearch.xpack.sql.rule.RuleExecutor;
6564
import org.elasticsearch.xpack.sql.session.EmptyExecutable;
6665
import org.elasticsearch.xpack.sql.util.Check;
67-
import org.elasticsearch.xpack.sql.util.DateUtils;
6866

69-
import java.time.ZoneId;
7067
import java.util.Arrays;
7168
import java.util.LinkedHashMap;
7269
import java.util.Map;
@@ -292,17 +289,18 @@ protected PhysicalPlan rule(AggregateExec a) {
292289
if (matchingGroup != null) {
293290
if (exp instanceof Attribute || exp instanceof ScalarFunction || exp instanceof GroupingFunction) {
294291
Processor action = null;
295-
ZoneId zi = exp.dataType().isDateBased() ? DateUtils.UTC : null;
292+
boolean isDateBased = exp.dataType().isDateBased();
296293
/*
297294
* special handling of dates since aggs return the typed Date object which needs
298295
* extraction instead of handling this in the scroller, the folder handles this
299296
* as it already got access to the extraction action
300297
*/
301298
if (exp instanceof DateTimeHistogramFunction) {
302299
action = ((UnaryPipe) p).action();
303-
zi = ((DateTimeFunction) exp).zoneId();
300+
isDateBased = true;
304301
}
305-
return new AggPathInput(exp.source(), exp, new GroupByRef(matchingGroup.id(), null, zi), action);
302+
return new AggPathInput(exp.source(), exp,
303+
new GroupByRef(matchingGroup.id(), null, isDateBased), action);
306304
}
307305
}
308306
// or found an aggregate expression (which has to work on an attribute used for grouping)
@@ -340,15 +338,12 @@ protected PhysicalPlan rule(AggregateExec a) {
340338
// attributes can only refer to declared groups
341339
if (child instanceof Attribute) {
342340
Check.notNull(matchingGroup, "Cannot find group [{}]", Expressions.name(child));
343-
// check if the field is a date - if so mark it as such to interpret the long as a date
344-
// UTC is used since that's what the server uses and there's no conversion applied
345-
// (like for date histograms)
346-
ZoneId zi = child.dataType().isDateBased() ? DateUtils.UTC : null;
347-
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, zi), ((Attribute) child));
341+
queryC = queryC.addColumn(
342+
new GroupByRef(matchingGroup.id(), null, child.dataType().isDateBased()), ((Attribute) child));
348343
}
349344
// handle histogram
350345
else if (child instanceof GroupingFunction) {
351-
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, null),
346+
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, child.dataType().isDateBased()),
352347
((GroupingFunction) child).toAttribute());
353348
}
354349
// fallback to regular agg functions
@@ -369,8 +364,8 @@ else if (child instanceof GroupingFunction) {
369364
matchingGroup = groupingContext.groupFor(ne);
370365
Check.notNull(matchingGroup, "Cannot find group [{}]", Expressions.name(ne));
371366

372-
ZoneId zi = ne.dataType().isDateBased() ? DateUtils.UTC : null;
373-
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, zi), ne.toAttribute());
367+
queryC = queryC.addColumn(
368+
new GroupByRef(matchingGroup.id(), null, ne.dataType().isDateBased()), ne.toAttribute());
374369
}
375370
}
376371
}
@@ -400,7 +395,7 @@ private Tuple<QueryContainer, AggPathInput> addAggFunction(GroupByKey groupingAg
400395
// if the count points to the total track hits, enable accurate count retrieval
401396
queryC = queryC.withTrackHits();
402397
} else {
403-
ref = new GroupByRef(groupingAgg.id(), Property.COUNT, null);
398+
ref = new GroupByRef(groupingAgg.id(), Property.COUNT, false);
404399
}
405400

406401
Map<String, GroupByKey> pseudoFunctions = new LinkedHashMap<>(queryC.pseudoFunctions());

0 commit comments

Comments
 (0)