Skip to content

Commit 45bfaab

Browse files
nik9000nik9000
authored
ESQL: ROUND_TO function (#128278)
Creates a `ROUND_TO` function that rounds it's input to one of the provided values. Like so: ``` ROUND_TO(v, 0, 5000, 10000, 20000, 40000, 100000) v | ROUND_TO 0 | 0 100 | 0 6000 | 5000 45001 | 40000 999999 | 100000 ``` For some sequences of numbers you could do this with the `/` operator - but for arbitrary sequences of numbers you needed `CASE` which is quite slow. And hard to read! Rewriting the example above would look like: ``` CASE ( v < 5000, 0, v < 10000, 5000, v < 20000, 10000, v < 40000, 20000, v < 100000, 40000, 100000 ) ``` Even better, this is *fast*: ``` (operation) Mode Cnt Score Error Units round_to_4_via_case avgt 7 138.124 ± 0.738 ns/op round_to_4 avgt 7 0.805 ± 0.011 ns/op round_to_3 avgt 7 0.739 ± 0.011 ns/op round_to_2 avgt 7 0.651 ± 0.009 ns/op date_trunc avgt 7 2.425 ± 0.018 ns/op ``` I've included a comparison to `DATE_TRUNC` above because we should be able to rewrite `DATE_TRUNC` into `ROUND_TO` when we know the date range of the index. This doesn't do it now, but it should be possible.
1 parent 0c3d47d commit 45bfaab

File tree

42 files changed

+4082
-4
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+4082
-4
lines changed

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

Lines changed: 150 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.elasticsearch.common.breaker.NoopCircuitBreaker;
1414
import org.elasticsearch.common.logging.LogConfigurator;
1515
import org.elasticsearch.common.settings.Settings;
16+
import org.elasticsearch.common.unit.ByteSizeUnit;
1617
import org.elasticsearch.common.util.BigArrays;
1718
import org.elasticsearch.compute.data.Block;
1819
import org.elasticsearch.compute.data.BlockFactory;
@@ -44,13 +45,15 @@
4445
import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
4546
import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
4647
import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs;
48+
import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
4749
import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin;
4850
import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
4951
import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
5052
import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToLower;
5153
import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToUpper;
5254
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
5355
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
56+
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
5457
import org.elasticsearch.xpack.esql.planner.Layout;
5558
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
5659
import org.elasticsearch.xpack.esql.session.Configuration;
@@ -128,6 +131,10 @@ static void selfTest() {
128131
"long_equal_to_int",
129132
"mv_min",
130133
"mv_min_ascending",
134+
"round_to_4_via_case",
135+
"round_to_2",
136+
"round_to_3",
137+
"round_to_4",
131138
"rlike",
132139
"to_lower",
133140
"to_lower_ords",
@@ -240,6 +247,65 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
240247
RLike rlike = new RLike(Source.EMPTY, keywordField, new RLikePattern(".ar"));
241248
yield EvalMapper.toEvaluator(FOLD_CONTEXT, rlike, layout(keywordField)).get(driverContext);
242249
}
250+
case "round_to_4_via_case" -> {
251+
FieldAttribute f = longField();
252+
253+
Expression ltkb = new LessThan(Source.EMPTY, f, kb());
254+
Expression ltmb = new LessThan(Source.EMPTY, f, mb());
255+
Expression ltgb = new LessThan(Source.EMPTY, f, gb());
256+
EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
257+
FOLD_CONTEXT,
258+
new Case(Source.EMPTY, ltkb, List.of(b(), ltmb, kb(), ltgb, mb(), gb())),
259+
layout(f)
260+
).get(driverContext);
261+
String desc = "CaseLazyEvaluator";
262+
if (evaluator.toString().contains(desc) == false) {
263+
throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
264+
}
265+
yield evaluator;
266+
}
267+
case "round_to_2" -> {
268+
FieldAttribute f = longField();
269+
270+
EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
271+
FOLD_CONTEXT,
272+
new RoundTo(Source.EMPTY, f, List.of(b(), kb())),
273+
layout(f)
274+
).get(driverContext);
275+
String desc = "RoundToLong2";
276+
if (evaluator.toString().contains(desc) == false) {
277+
throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
278+
}
279+
yield evaluator;
280+
}
281+
case "round_to_3" -> {
282+
FieldAttribute f = longField();
283+
284+
EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
285+
FOLD_CONTEXT,
286+
new RoundTo(Source.EMPTY, f, List.of(b(), kb(), mb())),
287+
layout(f)
288+
).get(driverContext);
289+
String desc = "RoundToLong3";
290+
if (evaluator.toString().contains(desc) == false) {
291+
throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
292+
}
293+
yield evaluator;
294+
}
295+
case "round_to_4" -> {
296+
FieldAttribute f = longField();
297+
298+
EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
299+
FOLD_CONTEXT,
300+
new RoundTo(Source.EMPTY, f, List.of(b(), kb(), mb(), gb())),
301+
layout(f)
302+
).get(driverContext);
303+
String desc = "RoundToLong4";
304+
if (evaluator.toString().contains(desc) == false) {
305+
throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
306+
}
307+
yield evaluator;
308+
}
243309
case "to_lower", "to_lower_ords" -> {
244310
FieldAttribute keywordField = keywordField();
245311
ToLower toLower = new ToLower(Source.EMPTY, keywordField, configuration());
@@ -419,6 +485,69 @@ private static void checkExpected(String operation, Page actual) {
419485
}
420486
}
421487
}
488+
case "round_to_4_via_case", "round_to_4" -> {
489+
long b = 1;
490+
long kb = ByteSizeUnit.KB.toBytes(1);
491+
long mb = ByteSizeUnit.MB.toBytes(1);
492+
long gb = ByteSizeUnit.GB.toBytes(1);
493+
494+
LongVector f = actual.<LongBlock>getBlock(0).asVector();
495+
LongVector result = actual.<LongBlock>getBlock(1).asVector();
496+
for (int i = 0; i < BLOCK_LENGTH; i++) {
497+
long expected = f.getLong(i);
498+
if (expected < kb) {
499+
expected = b;
500+
} else if (expected < mb) {
501+
expected = kb;
502+
} else if (expected < gb) {
503+
expected = mb;
504+
} else {
505+
expected = gb;
506+
}
507+
if (result.getLong(i) != expected) {
508+
throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
509+
}
510+
}
511+
}
512+
case "round_to_3" -> {
513+
long b = 1;
514+
long kb = ByteSizeUnit.KB.toBytes(1);
515+
long mb = ByteSizeUnit.MB.toBytes(1);
516+
517+
LongVector f = actual.<LongBlock>getBlock(0).asVector();
518+
LongVector result = actual.<LongBlock>getBlock(1).asVector();
519+
for (int i = 0; i < BLOCK_LENGTH; i++) {
520+
long expected = f.getLong(i);
521+
if (expected < kb) {
522+
expected = b;
523+
} else if (expected < mb) {
524+
expected = kb;
525+
} else {
526+
expected = mb;
527+
}
528+
if (result.getLong(i) != expected) {
529+
throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
530+
}
531+
}
532+
}
533+
case "round_to_2" -> {
534+
long b = 1;
535+
long kb = ByteSizeUnit.KB.toBytes(1);
536+
537+
LongVector f = actual.<LongBlock>getBlock(0).asVector();
538+
LongVector result = actual.<LongBlock>getBlock(1).asVector();
539+
for (int i = 0; i < BLOCK_LENGTH; i++) {
540+
long expected = f.getLong(i);
541+
if (expected < kb) {
542+
expected = b;
543+
} else {
544+
expected = kb;
545+
}
546+
if (result.getLong(i) != expected) {
547+
throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
548+
}
549+
}
550+
}
422551
case "to_lower" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
423552
case "to_lower_ords" -> checkBytes(operation, actual, true, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
424553
case "to_upper" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("FOO"), new BytesRef("BAR") });
@@ -450,7 +579,7 @@ private static void checkBytes(String operation, Page actual, boolean expectOrds
450579

451580
private static Page page(String operation) {
452581
return switch (operation) {
453-
case "abs", "add", "date_trunc", "equal_to_const" -> {
582+
case "abs", "add", "date_trunc", "equal_to_const", "round_to_4_via_case", "round_to_2", "round_to_3", "round_to_4" -> {
454583
var builder = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
455584
for (int i = 0; i < BLOCK_LENGTH; i++) {
456585
builder.appendLong(i * 100_000);
@@ -540,6 +669,26 @@ private static Page page(String operation) {
540669
};
541670
}
542671

672+
private static Literal b() {
673+
return lit(1L);
674+
}
675+
676+
private static Literal kb() {
677+
return lit(ByteSizeUnit.KB.toBytes(1));
678+
}
679+
680+
private static Literal mb() {
681+
return lit(ByteSizeUnit.MB.toBytes(1));
682+
}
683+
684+
private static Literal gb() {
685+
return lit(ByteSizeUnit.GB.toBytes(1));
686+
}
687+
688+
private static Literal lit(long v) {
689+
return new Literal(Source.EMPTY, v, DataType.LONG);
690+
}
691+
543692
@Benchmark
544693
@OperationsPerInvocation(1024 * BLOCK_LENGTH)
545694
public void run() {

‎docs/changelog/128278.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 128278
2+
summary: ROUND_TO function
3+
area: ES|QL
4+
type: enhancement
5+
issues: []

‎docs/reference/query-languages/esql/_snippets/functions/description/round_to.md‎

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

‎docs/reference/query-languages/esql/_snippets/functions/examples/round_to.md‎

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

‎docs/reference/query-languages/esql/_snippets/functions/layout/round_to.md‎

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

‎docs/reference/query-languages/esql/_snippets/functions/parameters/round_to.md‎

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

‎docs/reference/query-languages/esql/_snippets/functions/types/round_to.md‎

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

‎docs/reference/query-languages/esql/images/functions/round_to.svg‎

Lines changed: 1 addition & 0 deletions
Loading

0 commit comments

Comments
 (0)