Skip to content

Commit 91179dd

Browse files
update according to review comment
1 parent 41a51f5 commit 91179dd

File tree

9 files changed

+52
-74
lines changed

9 files changed

+52
-74
lines changed

x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundToDouble.java

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundToInt.java

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundToLong.java

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTo.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import java.io.IOException;
3131
import java.util.List;
3232
import java.util.Map;
33+
import java.util.Objects;
3334

3435
import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
3536
import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
@@ -39,6 +40,7 @@
3940
import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
4041
import static org.elasticsearch.xpack.esql.core.type.DataType.LONG;
4142
import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
43+
import static org.elasticsearch.xpack.esql.core.type.DataTypeConverter.safeToLong;
4244
import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.commonType;
4345

4446
/**
@@ -181,7 +183,8 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
181183
ExpressionEvaluator.Factory field = toEvaluator.apply(field());
182184
field = Cast.cast(source(), field().dataType(), dataType, field);
183185
List<Object> points = Iterators.toList(Iterators.map(points().iterator(), p -> Foldables.valueOf(toEvaluator.foldCtx(), p)));
184-
return build.build(source(), field, points);
186+
List<Object> sortedPoints = sortedRoundingPoints(points, dataType); // provide sorted points to the evaluator
187+
return build.build(source(), field, sortedPoints);
185188
}
186189

187190
interface Build {
@@ -195,4 +198,19 @@ interface Build {
195198
Map.entry(LONG, RoundToLong.BUILD),
196199
Map.entry(DOUBLE, RoundToDouble.BUILD)
197200
);
201+
202+
public static List<Object> sortedRoundingPoints(List<Object> points, DataType dataType) {
203+
return points.stream().filter(Objects::nonNull).map(p -> switch (dataType) { // the types are in sync with SIGNATURES
204+
case INTEGER -> ((Number) p).intValue();
205+
case DOUBLE -> ((Number) p).doubleValue();
206+
case LONG, DATETIME, DATE_NANOS -> safeToLong((Number) p);
207+
default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
208+
}).sorted((a, b) -> {
209+
if (a instanceof Double || b instanceof Double) {
210+
return Double.compare(a.doubleValue(), b.doubleValue());
211+
} else {
212+
return Long.compare(a.longValue(), b.longValue());
213+
}
214+
}).collect(java.util.stream.Collectors.toList());
215+
}
198216
}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/X-RoundTo.java.st

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ import java.util.Arrays;
2525
class RoundTo$Type$ {
2626
static final RoundTo.Build BUILD = (source, field, points) -> {
2727
$type$[] f = points.stream().mapTo$Type$(p -> ((Number) p).$type$Value()).toArray();
28-
Arrays.sort(f);
2928
return switch (f.length) {
3029
// TODO should be a consistent way to do the 0 version - is CASE(MV_COUNT(f) == 1, f[0])
3130
case 1 -> new RoundTo$Type$1Evaluator.Factory(source, field, f[0]);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java

Lines changed: 14 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131

3232
import java.time.ZoneId;
3333
import java.util.ArrayList;
34-
import java.util.Collections;
3534
import java.util.List;
3635
import java.util.Map;
3736

@@ -339,20 +338,20 @@ private static List<EsQueryExec.QueryBuilderAndTags> queryBuilderAndTags(
339338
int count = roundingPoints.size();
340339
DataType dataType = roundTo.dataType();
341340
// sort rounding points
342-
List<? extends Number> points = resolveRoundingPoints(dataType, roundingPoints);
341+
List<Object> points = resolveRoundingPoints(roundingPoints, dataType);
343342
if (points.size() != count || points.isEmpty()) {
344343
return null;
345344
}
346345
List<EsQueryExec.QueryBuilderAndTags> queries = new ArrayList<>(count);
347346

348-
Number tag = points.get(0);
347+
Object tag = points.get(0);
349348
if (points.size() == 1) { // if there is only one rounding point, just tag the main query
350349
EsQueryExec.QueryBuilderAndTags queryBuilderAndTags = tagOnlyBucket(queryExec, tag);
351350
queries.add(queryBuilderAndTags);
352351
} else {
353352
Source source = roundTo.source();
354-
Number lower = null;
355-
Number upper = null;
353+
Object lower = null;
354+
Object upper = null;
356355
Queries.Clause clause = queryExec.hasScoring() ? Queries.Clause.MUST : Queries.Clause.FILTER;
357356
ZoneId zoneId = ctx.configuration().zoneId();
358357
for (int i = 1; i < count; i++) {
@@ -370,46 +369,20 @@ private static List<EsQueryExec.QueryBuilderAndTags> queryBuilderAndTags(
370369
return queries;
371370
}
372371

373-
private static List<? extends Number> resolveRoundingPoints(DataType dataType, List<Expression> roundingPoints) {
374-
return switch (dataType) {
375-
case LONG, DATETIME, DATE_NANOS -> sortedLongRoundingPoints(roundingPoints);
376-
case INTEGER -> sortedIntRoundingPoints(roundingPoints);
377-
case DOUBLE -> sortedDoubleRoundingPoints(roundingPoints);
378-
default -> List.of();
379-
};
380-
}
381-
382-
private static List<Long> sortedLongRoundingPoints(List<Expression> roundingPoints) {
383-
List<Long> points = new ArrayList<>(roundingPoints.size());
372+
private static List<Object> resolveRoundingPoints(List<Expression> roundingPoints, DataType dataType) {
373+
List<Object> points = new ArrayList<>(roundingPoints.size());
384374
for (Expression e : roundingPoints) {
385375
if (e instanceof Literal l && l.value() instanceof Number n) {
386-
points.add(safeToLong(n));
387-
}
388-
}
389-
Collections.sort(points);
390-
return points;
391-
}
392-
393-
private static List<Integer> sortedIntRoundingPoints(List<Expression> roundingPoints) {
394-
List<Integer> points = new ArrayList<>(roundingPoints.size());
395-
for (Expression e : roundingPoints) {
396-
if (e instanceof Literal l) {
397-
points.add((Integer) l.value());
398-
}
399-
}
400-
Collections.sort(points);
401-
return points;
402-
}
403-
404-
private static List<Double> sortedDoubleRoundingPoints(List<Expression> roundingPoints) {
405-
List<Double> points = new ArrayList<>(roundingPoints.size());
406-
for (Expression e : roundingPoints) {
407-
if (e instanceof Literal l) {
408-
points.add((Double) l.value());
376+
switch (dataType) {
377+
case INTEGER -> points.add(n.intValue());
378+
case LONG, DATETIME, DATE_NANOS -> points.add(safeToLong(n));
379+
case DOUBLE -> points.add(n.doubleValue());
380+
// this should not happen, as RoundTo type resolution will fail with the other data types
381+
default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
382+
}
409383
}
410384
}
411-
Collections.sort(points);
412-
return points;
385+
return RoundTo.sortedRoundingPoints(points, dataType);
413386
}
414387

415388
private static Expression createRangeExpression(

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceSourceAttributes.java

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,16 @@ protected PhysicalPlan rule(EsSourceExec plan) {
6262
attributes.add(score);
6363
}
6464

65-
return new EsQueryExec(plan.source(), plan.indexPattern(), plan.indexMode(), plan.indexNameWithModes(), attributes, plan.query());
65+
return new EsQueryExec(
66+
plan.source(),
67+
plan.indexPattern(),
68+
plan.indexMode(),
69+
plan.indexNameWithModes(),
70+
attributes,
71+
null,
72+
null,
73+
null,
74+
List.of(new EsQueryExec.QueryBuilderAndTags(plan.query(), List.of()))
75+
);
6676
}
6777
}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EsQueryExec.java

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -105,27 +105,6 @@ public String toString() {
105105
}
106106
};
107107

108-
public EsQueryExec(
109-
Source source,
110-
String indexPattern,
111-
IndexMode indexMode,
112-
Map<String, IndexMode> indexNameWithModes,
113-
List<Attribute> attributes,
114-
QueryBuilder query
115-
) {
116-
this(
117-
source,
118-
indexPattern,
119-
indexMode,
120-
indexNameWithModes,
121-
attributes,
122-
null,
123-
null,
124-
null,
125-
List.of(new QueryBuilderAndTags(query, List.of()))
126-
);
127-
}
128-
129108
public EsQueryExec(
130109
Source source,
131110
String indexPattern,
@@ -201,7 +180,7 @@ public Map<String, IndexMode> indexNameWithModes() {
201180
}
202181

203182
/**
204-
* query is merged into queryBuilderAndTags, keep this method as it is called by too many places.
183+
* query is merged into queryBuilderAndTags, keep this method as it is called by many many places in both tests and product code.
205184
* If this method is called, the caller looks for the original queryBuilder, before {@code ReplaceRoundToWithQueryAndTags} converts it
206185
* to multiple queries with tags.
207186
*/

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -76,13 +76,13 @@ public ReplaceRoundToWithQueryAndTagsTests(String name, Configuration config) {
7676

7777
private static final Map<String, List<Object>> roundToAllTypes = new HashMap<>(
7878
Map.ofEntries(
79-
Map.entry("byte", List.of(1, 2, 3, 4)),
80-
Map.entry("short", List.of(1, 2, 3, 4)),
79+
Map.entry("byte", List.of(2, 1, 3, 4)),
80+
Map.entry("short", List.of(1, 3, 2, 4)),
8181
Map.entry("integer", List.of(1, 2, 3, 4)),
8282
Map.entry("long", List.of(1697760000000L, 1697846400000L, 1697932800000L, 1698019200000L)),
83-
Map.entry("float", List.of(1.0, 2.0, 3.0, 4.0)),
84-
Map.entry("half_float", List.of(1.0, 2.0, 3.0, 4.0)),
85-
Map.entry("scaled_float", List.of(1.0, 2.0, 3.0, 4.0)),
83+
Map.entry("float", List.of(3.0, 2.0, 1.0, 4.0)),
84+
Map.entry("half_float", List.of(4.0, 2.0, 3.0, 1.0)),
85+
Map.entry("scaled_float", List.of(4.0, 3.0, 2.0, 1.0)),
8686
Map.entry("double", List.of(1.0, 2.0, 3.0, 4.0)),
8787
Map.entry("date", List.of("\"2023-10-20\"::date", "\"2023-10-21\"::date", "\"2023-10-22\"::date", "\"2023-10-23\"::date")),
8888
Map.entry(
@@ -474,6 +474,8 @@ private static List<List<Object>> dateBuckets(boolean isDate) {
474474
}
475475

476476
private static List<List<Object>> numericBuckets(List<Object> roundingPoints) {
477+
// sort the rounding points in ascending order
478+
roundingPoints = roundingPoints.stream().sorted().collect(Collectors.toList());
477479
Object p1 = roundingPoints.get(0);
478480
Object p2 = roundingPoints.get(1);
479481
Object p3 = roundingPoints.get(2);

0 commit comments

Comments
 (0)