ESQL: ROUND_TO function (#128278) (#128397)

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.

Author: nik9000

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

@@ -12,6 +12,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
@@ -41,13 +42,15 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs;
+import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToLower;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToUpper;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
 import org.elasticsearch.xpack.esql.planner.Layout;
 import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.elasticsearch.xpack.esql.session.Configuration;
@@ -116,6 +119,10 @@ public class EvalBenchmark {
             "long_equal_to_int",
             "mv_min",
             "mv_min_ascending",
+            "round_to_4_via_case",
+            "round_to_2",
+            "round_to_3",
+            "round_to_4",
             "rlike",
             "to_lower",
             "to_lower_ords",
@@ -228,6 +235,65 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                 RLike rlike = new RLike(Source.EMPTY, keywordField, new RLikePattern(".ar"));
                 yield EvalMapper.toEvaluator(FOLD_CONTEXT, rlike, layout(keywordField)).get(driverContext);
             }
+            case "round_to_4_via_case" -> {
+                FieldAttribute f = longField();
+
+                Expression ltkb = new LessThan(Source.EMPTY, f, kb());
+                Expression ltmb = new LessThan(Source.EMPTY, f, mb());
+                Expression ltgb = new LessThan(Source.EMPTY, f, gb());
+                EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
+                    new Case(Source.EMPTY, ltkb, List.of(b(), ltmb, kb(), ltgb, mb(), gb())),
+                    layout(f)
+                ).get(driverContext);
+                String desc = "CaseLazyEvaluator";
+                if (evaluator.toString().contains(desc) == false) {
+                    throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
+                }
+                yield evaluator;
+            }
+            case "round_to_2" -> {
+                FieldAttribute f = longField();
+
+                EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
+                    new RoundTo(Source.EMPTY, f, List.of(b(), kb())),
+                    layout(f)
+                ).get(driverContext);
+                String desc = "RoundToLong2";
+                if (evaluator.toString().contains(desc) == false) {
+                    throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
+                }
+                yield evaluator;
+            }
+            case "round_to_3" -> {
+                FieldAttribute f = longField();
+
+                EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
+                    new RoundTo(Source.EMPTY, f, List.of(b(), kb(), mb())),
+                    layout(f)
+                ).get(driverContext);
+                String desc = "RoundToLong3";
+                if (evaluator.toString().contains(desc) == false) {
+                    throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
+                }
+                yield evaluator;
+            }
+            case "round_to_4" -> {
+                FieldAttribute f = longField();
+
+                EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator(
+                    FOLD_CONTEXT,
+                    new RoundTo(Source.EMPTY, f, List.of(b(), kb(), mb(), gb())),
+                    layout(f)
+                ).get(driverContext);
+                String desc = "RoundToLong4";
+                if (evaluator.toString().contains(desc) == false) {
+                    throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]");
+                }
+                yield evaluator;
+            }
             case "to_lower", "to_lower_ords" -> {
                 FieldAttribute keywordField = keywordField();
                 ToLower toLower = new ToLower(Source.EMPTY, keywordField, configuration());
@@ -390,6 +456,69 @@ private static void checkExpected(String operation, Page actual) {
                     }
                 }
             }
+            case "round_to_4_via_case", "round_to_4" -> {
+                long b = 1;
+                long kb = ByteSizeUnit.KB.toBytes(1);
+                long mb = ByteSizeUnit.MB.toBytes(1);
+                long gb = ByteSizeUnit.GB.toBytes(1);
+
+                LongVector f = actual.<LongBlock>getBlock(0).asVector();
+                LongVector result = actual.<LongBlock>getBlock(1).asVector();
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    long expected = f.getLong(i);
+                    if (expected < kb) {
+                        expected = b;
+                    } else if (expected < mb) {
+                        expected = kb;
+                    } else if (expected < gb) {
+                        expected = mb;
+                    } else {
+                        expected = gb;
+                    }
+                    if (result.getLong(i) != expected) {
+                        throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
+                    }
+                }
+            }
+            case "round_to_3" -> {
+                long b = 1;
+                long kb = ByteSizeUnit.KB.toBytes(1);
+                long mb = ByteSizeUnit.MB.toBytes(1);
+
+                LongVector f = actual.<LongBlock>getBlock(0).asVector();
+                LongVector result = actual.<LongBlock>getBlock(1).asVector();
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    long expected = f.getLong(i);
+                    if (expected < kb) {
+                        expected = b;
+                    } else if (expected < mb) {
+                        expected = kb;
+                    } else {
+                        expected = mb;
+                    }
+                    if (result.getLong(i) != expected) {
+                        throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
+                    }
+                }
+            }
+            case "round_to_2" -> {
+                long b = 1;
+                long kb = ByteSizeUnit.KB.toBytes(1);
+
+                LongVector f = actual.<LongBlock>getBlock(0).asVector();
+                LongVector result = actual.<LongBlock>getBlock(1).asVector();
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    long expected = f.getLong(i);
+                    if (expected < kb) {
+                        expected = b;
+                    } else {
+                        expected = kb;
+                    }
+                    if (result.getLong(i) != expected) {
+                        throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
+                    }
+                }
+            }
             case "to_lower" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
             case "to_lower_ords" -> checkBytes(operation, actual, true, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
             case "to_upper" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("FOO"), new BytesRef("BAR") });
@@ -421,7 +550,7 @@ private static void checkBytes(String operation, Page actual, boolean expectOrds
 
     private static Page page(String operation) {
         return switch (operation) {
-            case "abs", "add", "date_trunc", "equal_to_const" -> {
+            case "abs", "add", "date_trunc", "equal_to_const", "round_to_4_via_case", "round_to_2", "round_to_3", "round_to_4" -> {
                 var builder = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
                 for (int i = 0; i < BLOCK_LENGTH; i++) {
                     builder.appendLong(i * 100_000);
@@ -511,6 +640,26 @@ private static Page page(String operation) {
         };
     }
 
+    private static Literal b() {
+        return lit(1L);
+    }
+
+    private static Literal kb() {
+        return lit(ByteSizeUnit.KB.toBytes(1));
+    }
+
+    private static Literal mb() {
+        return lit(ByteSizeUnit.MB.toBytes(1));
+    }
+
+    private static Literal gb() {
+        return lit(ByteSizeUnit.GB.toBytes(1));
+    }
+
+    private static Literal lit(long v) {
+        return new Literal(Source.EMPTY, v, DataType.LONG);
+    }
+
     @Benchmark
     @OperationsPerInvocation(1024 * BLOCK_LENGTH)
     public void run() {

docs/changelog/128278.yaml

@@ -0,0 +1,5 @@
+pr: 128278
+summary: ROUND_TO function
+area: ES|QL
+type: enhancement
+issues: []

docs/changelog/128397.yaml

@@ -0,0 +1,5 @@
+pr: 128397
+summary: ROUND_TO function
+area: ES|QL
+type: enhancement
+issues: []