Two more tests fixes

Author: shmuelhanoch

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java

@@ -31,7 +31,7 @@ public static List<Object[]> params() {
                 continue;
             }
             for (boolean nullAllowed : new boolean[] { false, true }) {
-                if (e == ElementType.AGGREGATE_METRIC_DOUBLE) {
+                if (e == ElementType.AGGREGATE_METRIC_DOUBLE || e == ElementType.DATE_RANGE) {
                     // doesn't support multi-values
                     params.add(new Object[] { e, nullAllowed, 0, 1 });
                     continue;

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockValueAsserter.java

@@ -42,6 +42,7 @@ static void assertBlockValues(Block block, List<List<Object>> expectedBlockValue
                         valueCount,
                         expectedRowValues
                     );
+                    case DATE_RANGE -> assertDteRangeValues((DateRangeBlock) block, firstValueIndex, valueCount, expectedRowValues);
                     default -> throw new IllegalArgumentException("Unsupported element type [" + block.elementType() + "]");
                 }
             }
@@ -117,4 +118,12 @@ private static void assertAggregateMetricRowValues(
             assertThat(block.countBlock().getInt(firstValueIndex + valueIndex), is(equalTo(expectedValue.count())));
         }
     }
+
+    private static void assertDteRangeValues(DateRangeBlock block, int firstValueIdx, int valueCount, List<Object> expectedRowValues) {
+        for (int idx = 0; idx < valueCount; idx++) {
+            var expectedValue = (DateRangeBlockBuilder.DateRangeLiteral) expectedRowValues.get(idx);
+            assertThat(block.getFromBlock().getLong(firstValueIdx + idx), equalTo(expectedValue.from()));
+            assertThat(block.getToBlock().getLong(firstValueIdx + idx), equalTo(expectedValue.to()));
+        }
+    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/ExtractorTests.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.DateRangeBlockBuilder;
 import org.elasticsearch.compute.data.DocVector;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.lucene.AlwaysReferencedIndexedByShardId;
@@ -30,6 +31,7 @@
 import java.util.List;
 import java.util.function.Supplier;
 
+import static org.elasticsearch.common.time.DateUtils.MAX_MILLIS_BEFORE_9999;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 
@@ -63,6 +65,9 @@ public static Iterable<Object[]> parameters() {
                         )
                     );
                 }
+                case DATE_RANGE -> cases.add(
+                    valueTestCase("date_range with nulls", e, TopNEncoder.DEFAULT_UNSORTABLE, ExtractorTests::randomDateRange)
+                );
                 case FLOAT -> {
                 }
                 case BYTES_REF -> {
@@ -252,4 +257,10 @@ public static AggregateMetricDoubleLiteral randomAggregateMetricDouble(boolean a
             randomBoolean() ? randomInt() : null
         );
     }
+
+    private static DateRangeBlockBuilder.DateRangeLiteral randomDateRange() {
+        var from = randomMillisUpToYear9999();
+        var to = randomLongBetween(from + 1, MAX_MILLIS_BEFORE_9999);
+        return new DateRangeBlockBuilder.DateRangeLiteral(from, to);
+    }
 }