Some more infra for new tests

shmuelhanoch · shmuelhanoch · commit 67affcf270e2 · 2025-10-11T03:16:38.000+03:00
diff --git a/server/src/main/resources/transport/definitions/referable/esql_date_range_created_version.csv b/server/src/main/resources/transport/definitions/referable/esql_date_range_created_version.csv
@@ -1 +1 @@
-9191000
+9192000
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractor.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractor.java
@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.DateRangeBlock;
 import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.ElementType;
@@ -53,6 +54,7 @@ static ValueExtractor extractorFor(ElementType elementType, TopNEncoder encoder,
             case NULL -> new ValueExtractorForNull();
             case DOC -> new ValueExtractorForDoc(encoder, ((DocBlock) block).asVector());
             case AGGREGATE_METRIC_DOUBLE -> new ValueExtractorForAggregateMetricDouble(encoder, (AggregateMetricDoubleBlock) block);
+            case DATE_RANGE -> new ValueExtractorForDateRange(encoder, (DateRangeBlock) block);
             default -> {
                 assert false : "No value extractor for [" + block.elementType() + "]";
                 throw new UnsupportedOperationException("No value extractor for [" + block.elementType() + "]");
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractorForDateRange.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/ValueExtractorForDateRange.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator.topn;
+
+import org.elasticsearch.compute.data.DateRangeBlock;
+import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
+
+public class ValueExtractorForDateRange implements ValueExtractor {
+    private final DateRangeBlock block;
+
+    ValueExtractorForDateRange(TopNEncoder encoder, DateRangeBlock block) {
+        assert encoder == TopNEncoder.DEFAULT_UNSORTABLE;
+        this.block = block;
+    }
+
+    @Override
+    public void writeValue(BreakingBytesRefBuilder values, int position) {
+        TopNEncoder.DEFAULT_UNSORTABLE.encodeVInt(1, values);
+        if (block.isNull(position)) {
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeBoolean(false, values);
+        } else {
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeBoolean(true, values);
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(block.getFromBlock().getLong(position), values);
+            TopNEncoder.DEFAULT_UNSORTABLE.encodeDouble(block.getToBlock().getLong(position), values);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "ValueExtractorForDateRange";
+    }
+}
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/ExtractorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/ExtractorTests.java
@@ -66,9 +66,10 @@ public static Iterable<Object[]> parameters() {
                         )
                     );
                 }
-                case DATE_RANGE -> cases.add(
-                    valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, ExtractorTests::randomDateRange)
-                );
+                case DATE_RANGE -> {
+                    cases.add(valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, () -> randomDateRange(true)));
+                    cases.add(valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, () -> randomDateRange(false)));
+                }
                 case FLOAT -> {
                 }
                 case BYTES_REF -> {
@@ -262,9 +263,9 @@ public static AggregateMetricDoubleLiteral randomAggregateMetricDouble(boolean a
         );
     }
 
-    private static DateRangeBlockBuilder.DateRangeLiteral randomDateRange() {
+    private static DateRangeBlockBuilder.DateRangeLiteral randomDateRange(boolean isNull) {
         var from = randomMillisUpToYear9999();
         var to = randomLongBetween(from + 1, MAX_MILLIS_BEFORE_9999);
-        return new DateRangeBlockBuilder.DateRangeLiteral(from, to);
+        return isNull ? null : new DateRangeBlockBuilder.DateRangeLiteral(from, to);
     }
 }
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java
@@ -74,6 +74,7 @@
 import static org.elasticsearch.compute.data.ElementType.BOOLEAN;
 import static org.elasticsearch.compute.data.ElementType.BYTES_REF;
 import static org.elasticsearch.compute.data.ElementType.COMPOSITE;
+import static org.elasticsearch.compute.data.ElementType.DATE_RANGE;
 import static org.elasticsearch.compute.data.ElementType.DOUBLE;
 import static org.elasticsearch.compute.data.ElementType.FLOAT;
 import static org.elasticsearch.compute.data.ElementType.INT;
@@ -602,7 +603,7 @@ public void testCollectAllValues_RandomMultiValues() {
 
         for (int type = 0; type < blocksCount; type++) {
             ElementType e = randomFrom(ElementType.values());
-            if (e == ElementType.UNKNOWN || e == COMPOSITE || e == AGGREGATE_METRIC_DOUBLE) {
+            if (e == ElementType.UNKNOWN || e == COMPOSITE || e == AGGREGATE_METRIC_DOUBLE || e == DATE_RANGE) {
                 continue;
             }
             elementTypes.add(e);
@@ -1024,7 +1025,7 @@ public void testRandomMultiValuesTopN() {
 
         for (int type = 0; type < blocksCount; type++) {
             ElementType e = randomValueOtherThanMany(
-                t -> t == ElementType.UNKNOWN || t == ElementType.DOC || t == COMPOSITE || t == AGGREGATE_METRIC_DOUBLE,
+                t -> t == ElementType.UNKNOWN || t == ElementType.DOC || t == COMPOSITE || t == AGGREGATE_METRIC_DOUBLE || t == DATE_RANGE,
                 () -> randomFrom(ElementType.values())
             );
             elementTypes.add(e);
diff --git a/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java b/x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java
@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.common.time.DateUtils.MAX_MILLIS_BEFORE_9999;
 import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
 import static org.elasticsearch.test.ESTestCase.between;
 import static org.elasticsearch.test.ESTestCase.randomBoolean;
@@ -43,6 +44,8 @@
 import static org.elasticsearch.test.ESTestCase.randomInt;
 import static org.elasticsearch.test.ESTestCase.randomIntBetween;
 import static org.elasticsearch.test.ESTestCase.randomLong;
+import static org.elasticsearch.test.ESTestCase.randomLongBetween;
+import static org.elasticsearch.test.ESTestCase.randomMillisUpToYear9999;
 import static org.elasticsearch.test.ESTestCase.randomNonNegativeInt;
 import static org.elasticsearch.test.ESTestCase.randomRealisticUnicodeOfCodepointLengthBetween;
 import static org.hamcrest.Matchers.equalTo;
@@ -66,14 +69,18 @@ public static Object randomValue(ElementType e) {
                 randomDouble(),
                 randomNonNegativeInt()
             );
+            case DATE_RANGE -> {
+                var from = randomMillisUpToYear9999();
+                var to = randomLongBetween(from + 1, MAX_MILLIS_BEFORE_9999);
+                yield new DateRangeBlockBuilder.DateRangeLiteral(from, to);
+            }
             case DOC -> new BlockUtils.Doc(
                 randomIntBetween(0, 255), // Shard ID should be small and non-negative.
                 randomInt(),
                 between(0, Integer.MAX_VALUE)
             );
             case NULL -> null;
             case COMPOSITE -> throw new IllegalArgumentException("can't make random values for composite");
-            case DATE_RANGE -> throw new IllegalArgumentException("can't make random values for date range");
             case UNKNOWN -> throw new IllegalArgumentException("can't make random values for [" + e + "]");
         };
     }

Original file line number	Diff line number	Diff line change
`@@ -66,9 +66,10 @@ public static Iterable<Object[]> parameters() {`
`66`	`66`	`)`
`67`	`67`	`);`
`68`	`68`	`}`
`69`		`- case DATE_RANGE -> cases.add(`
`70`		`- valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, ExtractorTests::randomDateRange)`
`71`		`- );`
	`69`	`+ case DATE_RANGE -> {`
	`70`	`+ cases.add(valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, () -> randomDateRange(true)));`
	`71`	`+ cases.add(valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, () -> randomDateRange(false)));`
	`72`	`+ }`
`72`	`73`	`case FLOAT -> {`
`73`	`74`	`}`
`74`	`75`	`case BYTES_REF -> {`
`@@ -262,9 +263,9 @@ public static AggregateMetricDoubleLiteral randomAggregateMetricDouble(boolean a`
`262`	`263`	`);`
`263`	`264`	`}`
`264`	`265`
`265`		`- private static DateRangeBlockBuilder.DateRangeLiteral randomDateRange() {`
	`266`	`+ private static DateRangeBlockBuilder.DateRangeLiteral randomDateRange(boolean isNull) {`
`266`	`267`	`var from = randomMillisUpToYear9999();`
`267`	`268`	`var to = randomLongBetween(from + 1, MAX_MILLIS_BEFORE_9999);`
`268`		`- return new DateRangeBlockBuilder.DateRangeLiteral(from, to);`
	`269`	`+ return isNull ? null : new DateRangeBlockBuilder.DateRangeLiteral(from, to);`
`269`	`270`	`}`
`270`	`271`	`}`