[ESQL] Support date_nanos on functions that take "any" type (#114056)

not-napoleon · elasticmachine · web-flow · commit 9c923db8f615 · 2024-10-22T12:15:44.000-04:00
Resolves #109998 For the most part, this is just adding tests. Greater and Least have actual production code changes - notably toEvaluator is modified to map date nanos to the long evaluator. This parallels the work done in #113961. I've added CSV tests and unit tests for all the functions listed in the original ticket. --------- Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/date_nanos.csv b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/date_nanos.csv
@@ -7,3 +7,4 @@ millis:date,nanos:date_nanos,num:long
 2023-10-23T12:27:28.948Z,2023-10-23T12:27:28.948000000Z,1698064048948000000
 2023-10-23T12:15:03.360Z,2023-10-23T12:15:03.360103847Z,1698063303360103847
 1999-10-23T12:15:03.360Z,[2023-03-23T12:15:03.360103847Z, 2023-02-23T13:33:34.937193000Z, 2023-01-23T13:55:01.543123456Z], 0
+1999-10-22T12:15:03.360Z,[2023-03-23T12:15:03.360103847Z, 2023-03-23T12:15:03.360103847Z, 2023-03-23T12:15:03.360103847Z], 0
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/date_nanos.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/date_nanos.csv-spec
@@ -38,9 +38,10 @@ nanos:date_nanos
 mv_min on date nanos
 required_capability: date_nanos_type
 
-FROM date_nanos | SORT millis ASC | EVAL nanos = MV_MIN(nanos) | KEEP nanos | LIMIT 1;
+FROM date_nanos | SORT millis ASC | WHERE millis < "2000-01-01" | EVAL nanos = MV_MIN(nanos) |  KEEP nanos; 
 
 nanos:date_nanos
+2023-03-23T12:15:03.360103847Z
 2023-01-23T13:55:01.543123456Z
 ;
 
@@ -56,9 +57,10 @@ ct:integer
 mv_first on date nanos
 required_capability: date_nanos_type
 
-FROM date_nanos | SORT millis ASC | EVAL nanos = MV_FIRST(nanos) | KEEP nanos | LIMIT 1;
+FROM date_nanos | SORT millis ASC | WHERE millis < "2000-01-01" | EVAL nanos = MV_FIRST(nanos) |  KEEP nanos;
 
 nanos:date_nanos
+2023-03-23T12:15:03.360103847Z
 2023-01-23T13:55:01.543123456Z
 ;
 
@@ -263,3 +265,69 @@ ROW a = TO_DATE_NANOS(null), b = TO_DATE_NANOS(null + 1::long), c = TO_DATE_NANO
 a:date_nanos | b:date_nanos | c:date_nanos
 null         | null         | null
 ;
+
+Coalasce date nanos
+required_capability: to_date_nanos
+
+ROW a = COALESCE(null, TO_DATE_NANOS(1698069301543123456));
+
+a:date_nanos
+2023-10-23T13:55:01.543123456Z
+;
+
+Case date nanos result
+required_capability: to_date_nanos
+
+ROW a = CASE(false, TO_DATE_NANOS(0::long), TO_DATE_NANOS(1698069301543123456));
+
+a:date_nanos
+2023-10-23T13:55:01.543123456Z
+;
+
+Greatest date nanos
+required_capability: least_greatest_for_datenanos
+
+ROW a = GREATEST(TO_DATE_NANOS("2023-10-23T13:55:01.543123456"), TO_DATE_NANOS("2023-10-23T13:53:55.832987654"));
+
+a:date_nanos
+2023-10-23T13:55:01.543123456Z
+;
+
+Least date nanos
+required_capability: least_greatest_for_datenanos
+
+ROW a = LEAST(TO_DATE_NANOS("2023-10-23T13:55:01.543123456"), TO_DATE_NANOS("2023-10-23T13:53:55.832987654"));
+
+a:date_nanos
+2023-10-23T13:53:55.832987654Z
+;
+
+mv_dedup over date nanos
+required_capability: date_nanos_type
+
+FROM date_nanos | WHERE millis < "2000-01-01" | EVAL a = MV_DEDUPE(nanos) | SORT millis DESC | KEEP a;
+
+a:date_nanos
+[2023-01-23T13:55:01.543123456Z, 2023-02-23T13:33:34.937193000Z, 2023-03-23T12:15:03.360103847Z]
+2023-03-23T12:15:03.360103847Z
+;
+
+mv_sort over date nanos
+required_capability: date_nanos_type
+
+FROM date_nanos | WHERE millis < "2000-01-01" | EVAL a = MV_SORT(nanos, "asc") | SORT millis DESC | KEEP a;
+
+a:date_nanos
+[2023-01-23T13:55:01.543123456Z, 2023-02-23T13:33:34.937193000Z, 2023-03-23T12:15:03.360103847Z]
+[2023-03-23T12:15:03.360103847Z, 2023-03-23T12:15:03.360103847Z, 2023-03-23T12:15:03.360103847Z]
+;
+
+mv_slice over date nanos
+required_capability: date_nanos_type
+
+FROM date_nanos | WHERE millis < "2000-01-01" | EVAL a = MV_SLICE(MV_SORT(nanos, "asc"), 1, 2) | SORT millis DESC | KEEP a;
+
+a:date_nanos
+[2023-02-23T13:33:34.937193000Z, 2023-03-23T12:15:03.360103847Z]
+[2023-03-23T12:15:03.360103847Z, 2023-03-23T12:15:03.360103847Z]
+;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
@@ -308,6 +308,11 @@ public enum Cap {
          */
         TO_DATE_NANOS(EsqlCorePlugin.DATE_NANOS_FEATURE_FLAG),
 
+        /**
+         * Support Least and Greatest functions on Date Nanos type
+         */
+        LEAST_GREATEST_FOR_DATENANOS(EsqlCorePlugin.DATE_NANOS_FEATURE_FLAG),
+
         /**
          * Support for datetime in least and greatest functions
          */
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Greatest.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Greatest.java
@@ -43,7 +43,7 @@ public class Greatest extends EsqlScalarFunction implements OptionalArgument {
     private DataType dataType;
 
     @FunctionInfo(
-        returnType = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+        returnType = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
         description = "Returns the maximum value from multiple columns. This is similar to <<esql-mv_max>>\n"
             + "except it is intended to run on multiple columns at once.",
         note = "When run on `keyword` or `text` fields, this returns the last string in alphabetical order. "
@@ -54,12 +54,12 @@ public Greatest(
         Source source,
         @Param(
             name = "first",
-            type = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+            type = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
             description = "First of the columns to evaluate."
         ) Expression first,
         @Param(
             name = "rest",
-            type = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+            type = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
             description = "The rest of the columns to evaluate.",
             optional = true
         ) List<Expression> rest
@@ -152,7 +152,7 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
         if (dataType == DataType.INTEGER) {
             return new GreatestIntEvaluator.Factory(source(), factories);
         }
-        if (dataType == DataType.LONG || dataType == DataType.DATETIME) {
+        if (dataType == DataType.LONG || dataType == DataType.DATETIME || dataType == DataType.DATE_NANOS) {
             return new GreatestLongEvaluator.Factory(source(), factories);
         }
         if (DataType.isString(dataType) || dataType == DataType.IP || dataType == DataType.VERSION || dataType == DataType.UNSUPPORTED) {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Least.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Least.java
@@ -43,7 +43,7 @@ public class Least extends EsqlScalarFunction implements OptionalArgument {
     private DataType dataType;
 
     @FunctionInfo(
-        returnType = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+        returnType = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
         description = "Returns the minimum value from multiple columns. "
             + "This is similar to <<esql-mv_min>> except it is intended to run on multiple columns at once.",
         examples = @Example(file = "math", tag = "least")
@@ -52,12 +52,12 @@ public Least(
         Source source,
         @Param(
             name = "first",
-            type = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+            type = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
             description = "First of the columns to evaluate."
         ) Expression first,
         @Param(
             name = "rest",
-            type = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+            type = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
             description = "The rest of the columns to evaluate.",
             optional = true
         ) List<Expression> rest
@@ -151,7 +151,7 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
         if (dataType == DataType.INTEGER) {
             return new LeastIntEvaluator.Factory(source(), factories);
         }
-        if (dataType == DataType.LONG || dataType == DataType.DATETIME) {
+        if (dataType == DataType.LONG || dataType == DataType.DATETIME || dataType == DataType.DATE_NANOS) {
             return new LeastLongEvaluator.Factory(source(), factories);
         }
         if (DataType.isString(dataType) || dataType == DataType.IP || dataType == DataType.VERSION || dataType == DataType.UNSUPPORTED) {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvDedupe.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvDedupe.java
@@ -38,6 +38,7 @@ public class MvDedupe extends AbstractMultivalueFunction {
             "cartesian_point",
             "cartesian_shape",
             "date",
+            "date_nanos",
             "double",
             "geo_point",
             "geo_shape",
@@ -60,6 +61,7 @@ public MvDedupe(
                 "cartesian_point",
                 "cartesian_shape",
                 "date",
+                "date_nanos",
                 "double",
                 "geo_point",
                 "geo_shape",
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSlice.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSlice.java
@@ -59,6 +59,7 @@ public class MvSlice extends EsqlScalarFunction implements OptionalArgument, Eva
             "cartesian_point",
             "cartesian_shape",
             "date",
+            "date_nanos",
             "double",
             "geo_point",
             "geo_shape",
@@ -87,6 +88,7 @@ public MvSlice(
                 "cartesian_point",
                 "cartesian_shape",
                 "date",
+                "date_nanos",
                 "double",
                 "geo_point",
                 "geo_shape",
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSort.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSort.java
@@ -69,15 +69,15 @@ public class MvSort extends EsqlScalarFunction implements OptionalArgument, Vali
     private static final String INVALID_ORDER_ERROR = "Invalid order value in [{}], expected one of [{}, {}] but got [{}]";
 
     @FunctionInfo(
-        returnType = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+        returnType = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
         description = "Sorts a multivalued field in lexicographical order.",
         examples = @Example(file = "ints", tag = "mv_sort")
     )
     public MvSort(
         Source source,
         @Param(
             name = "field",
-            type = { "boolean", "date", "double", "integer", "ip", "keyword", "long", "text", "version" },
+            type = { "boolean", "date", "date_nanos", "double", "integer", "ip", "keyword", "long", "text", "version" },
             description = "Multivalue expression. If `null`, the function returns `null`."
         ) Expression field,
         @Param(
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java
@@ -53,6 +53,7 @@ public class Coalesce extends EsqlScalarFunction implements OptionalArgument {
             "boolean",
             "cartesian_point",
             "cartesian_shape",
+            "date_nanos",
             "date",
             "geo_point",
             "geo_shape",
@@ -73,6 +74,7 @@ public Coalesce(
                 "boolean",
                 "cartesian_point",
                 "cartesian_shape",
+                "date_nanos",
                 "date",
                 "geo_point",
                 "geo_shape",
@@ -90,6 +92,7 @@ public Coalesce(
                 "boolean",
                 "cartesian_point",
                 "cartesian_shape",
+                "date_nanos",
                 "date",
                 "geo_point",
                 "geo_shape",
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java
@@ -722,17 +722,19 @@ public static void testFunctionInfo() {
             for (int i = 0; i < args.size() && i < types.size(); i++) {
                 typesFromSignature.get(i).add(types.get(i).esNameIfPossible());
             }
-            returnFromSignature.add(entry.getValue().esNameIfPossible());
+            if (DataType.UNDER_CONSTRUCTION.containsKey(entry.getValue()) == false) {
+                returnFromSignature.add(entry.getValue().esNameIfPossible());
+            }
         }
 
         for (int i = 0; i < args.size(); i++) {
             EsqlFunctionRegistry.ArgSignature arg = args.get(i);
             Set<String> annotationTypes = Arrays.stream(arg.type())
-                .filter(DataType.UNDER_CONSTRUCTION::containsKey)
+                .filter(t -> DataType.UNDER_CONSTRUCTION.containsKey(DataType.fromNameOrAlias(t)) == false)
                 .collect(Collectors.toCollection(TreeSet::new));
             Set<String> signatureTypes = typesFromSignature.get(i)
                 .stream()
-                .filter(DataType.UNDER_CONSTRUCTION::containsKey)
+                .filter(t -> DataType.UNDER_CONSTRUCTION.containsKey(DataType.fromNameOrAlias(t)) == false)
                 .collect(Collectors.toCollection(TreeSet::new));
             if (signatureTypes.isEmpty()) {
                 log.info("{}: skipping", arg.name());
@@ -746,7 +748,9 @@ public static void testFunctionInfo() {
             );
         }
 
-        Set<String> returnTypes = Arrays.stream(description.returnType()).collect(Collectors.toCollection(TreeSet::new));
+        Set<String> returnTypes = Arrays.stream(description.returnType())
+            .filter(t -> DataType.UNDER_CONSTRUCTION.containsKey(DataType.fromNameOrAlias(t)) == false)
+            .collect(Collectors.toCollection(TreeSet::new));
         assertEquals(returnFromSignature, returnTypes);
     }
 
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java
@@ -46,6 +46,7 @@ public class CaseTests extends AbstractScalarFunctionTestCase {
             DataType.TEXT,
             DataType.BOOLEAN,
             DataType.DATETIME,
+            DataType.DATE_NANOS,
             DataType.DOUBLE,
             DataType.INTEGER,
             DataType.LONG,
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/GreatestTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/GreatestTests.java
@@ -115,6 +115,21 @@ public static Iterable<Object[]> parameters() {
                 )
             )
         );
+        suppliers.add(
+            new TestCaseSupplier(
+                "(a, b)",
+                List.of(DataType.DATE_NANOS, DataType.DATE_NANOS),
+                () -> new TestCaseSupplier.TestCase(
+                    List.of(
+                        new TestCaseSupplier.TypedData(1727877348000123456L, DataType.DATE_NANOS, "a"),
+                        new TestCaseSupplier.TypedData(1727790948000987654L, DataType.DATE_NANOS, "b")
+                    ),
+                    "GreatestLongEvaluator[values=[MvMax[field=Attribute[channel=0]], MvMax[field=Attribute[channel=1]]]]",
+                    DataType.DATE_NANOS,
+                    equalTo(1727877348000123456L)
+                )
+            )
+        );
         return parameterSuppliersFromTypedData(anyNullIsNull(false, suppliers));
     }
 
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/LeastTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/LeastTests.java
@@ -114,6 +114,21 @@ public static Iterable<Object[]> parameters() {
                 )
             )
         );
+        suppliers.add(
+            new TestCaseSupplier(
+                "(a, b)",
+                List.of(DataType.DATE_NANOS, DataType.DATE_NANOS),
+                () -> new TestCaseSupplier.TestCase(
+                    List.of(
+                        new TestCaseSupplier.TypedData(1727877348000123456L, DataType.DATE_NANOS, "a"),
+                        new TestCaseSupplier.TypedData(1727790948000987654L, DataType.DATE_NANOS, "b")
+                    ),
+                    "LeastLongEvaluator[values=[MvMin[field=Attribute[channel=0]], MvMin[field=Attribute[channel=1]]]]",
+                    DataType.DATE_NANOS,
+                    equalTo(1727790948000987654L)
+                )
+            )
+        );
         return parameterSuppliersFromTypedData(anyNullIsNull(false, suppliers));
     }
 
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvDedupeTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvDedupeTests.java
@@ -39,6 +39,7 @@ public static Iterable<Object[]> parameters() {
         booleans(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values));
         bytesRefs(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values));
         dateTimes(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values.mapToObj(Long::valueOf)));
+        dateNanos(cases, "mv_dedupe", "MvDedupe", DataType.DATE_NANOS, (size, values) -> getMatcher(values.mapToObj(Long::valueOf)));
         doubles(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values.mapToObj(Double::valueOf)));
         ints(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values.mapToObj(Integer::valueOf)));
         longs(cases, "mv_dedupe", "MvDedupe", (size, values) -> getMatcher(values.mapToObj(Long::valueOf)));
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSliceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSliceTests.java
@@ -182,6 +182,23 @@ private static void longs(List<TestCaseSupplier> suppliers) {
                 equalTo(start == end ? field.get(start) : field.subList(start, end + 1))
             );
         }));
+
+        suppliers.add(new TestCaseSupplier(List.of(DataType.DATE_NANOS, DataType.INTEGER, DataType.INTEGER), () -> {
+            List<Long> field = randomList(1, 10, () -> randomLong());
+            int length = field.size();
+            int start = randomIntBetween(0, length - 1);
+            int end = randomIntBetween(start, length - 1);
+            return new TestCaseSupplier.TestCase(
+                List.of(
+                    new TestCaseSupplier.TypedData(field, DataType.DATE_NANOS, "field"),
+                    new TestCaseSupplier.TypedData(start, DataType.INTEGER, "start"),
+                    new TestCaseSupplier.TypedData(end, DataType.INTEGER, "end")
+                ),
+                "MvSliceLongEvaluator[field=Attribute[channel=0], start=Attribute[channel=1], end=Attribute[channel=2]]",
+                DataType.DATE_NANOS,
+                equalTo(start == end ? field.get(start) : field.subList(start, end + 1))
+            );
+        }));
     }
 
     private static void doubles(List<TestCaseSupplier> suppliers) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSortTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSortTests.java
@@ -110,6 +110,20 @@ private static void longs(List<TestCaseSupplier> suppliers) {
                 equalTo(field.size() == 1 ? field.iterator().next() : field.stream().sorted(Collections.reverseOrder()).toList())
             );
         }));
+
+        suppliers.add(new TestCaseSupplier(List.of(DataType.DATE_NANOS, DataType.KEYWORD), () -> {
+            List<Long> field = randomList(1, 10, () -> randomLong());
+            BytesRef order = new BytesRef("DESC");
+            return new TestCaseSupplier.TestCase(
+                List.of(
+                    new TestCaseSupplier.TypedData(field, DataType.DATE_NANOS, "field"),
+                    new TestCaseSupplier.TypedData(order, DataType.KEYWORD, "order").forceLiteral()
+                ),
+                "MvSortLong[field=Attribute[channel=0], order=false]",
+                DataType.DATE_NANOS,
+                equalTo(field.size() == 1 ? field.iterator().next() : field.stream().sorted(Collections.reverseOrder()).toList())
+            );
+        }));
     }
 
     private static void doubles(List<TestCaseSupplier> suppliers) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java
