[ESQL] Add support for date trunc on date nanos type (#116354)

Resolves #110008

As discussed elsewhere, this does NOT allow for truncating to a value smaller than a millisecond. Our timespan literal syntax doesn't allow specifying less than a millisecond, and the rounding infrastructure also does not support it.

We also had a discussion regarding the return type, and decided that it made sense to keep the type as date_nanos, even though the truncation will always produce a millisecond-rounded (or higher) value.

---------

Co-authored-by: Elastic Machine <[email protected]>

x-pack/plugin/esql/qa/testFixtures/src/main/resources/date_nanos.csv-spec

@@ -439,3 +439,23 @@ FROM date_nanos | WHERE millis > "2020-01-01" | STATS v = MV_SORT(VALUES(nanos),
 v:date_nanos
 [2023-10-23T13:55:01.543123456Z, 2023-10-23T13:53:55.832987654Z, 2023-10-23T13:52:55.015787878Z, 2023-10-23T13:51:54.732102837Z, 2023-10-23T13:33:34.937193000Z, 2023-10-23T12:27:28.948000000Z, 2023-10-23T12:15:03.360103847Z]
 ;
+
+Date trunc on date nanos
+required_capability: date_trunc_date_nanos
+
+FROM date_nanos 
+| WHERE millis > "2020-01-01" 
+| EVAL yr = DATE_TRUNC(1 year, nanos), mo = DATE_TRUNC(1 month, nanos), mn = DATE_TRUNC(10 minutes, nanos), ms = DATE_TRUNC(1 millisecond, nanos)
+| SORT nanos DESC
+| KEEP yr, mo, mn, ms;
+
+yr:date_nanos                  | mo:date_nanos                  | mn:date_nanos                  | ms:date_nanos                 
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T13:50:00.000000000Z | 2023-10-23T13:55:01.543000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T13:50:00.000000000Z | 2023-10-23T13:53:55.832000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T13:50:00.000000000Z | 2023-10-23T13:52:55.015000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T13:50:00.000000000Z | 2023-10-23T13:51:54.732000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T13:30:00.000000000Z | 2023-10-23T13:33:34.937000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T12:20:00.000000000Z | 2023-10-23T12:27:28.948000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T12:10:00.000000000Z | 2023-10-23T12:15:03.360000000Z
+2023-01-01T00:00:00.000000000Z | 2023-10-01T00:00:00.000000000Z | 2023-10-23T12:10:00.000000000Z | 2023-10-23T12:15:03.360000000Z
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -329,6 +329,11 @@ public enum Cap {
          */
         LEAST_GREATEST_FOR_DATENANOS(EsqlCorePlugin.DATE_NANOS_FEATURE_FLAG),
 
+        /**
+         * Support for date_trunc function on date nanos type
+         */
+        DATE_TRUNC_DATE_NANOS(EsqlCorePlugin.DATE_NANOS_FEATURE_FLAG),
+
         /**
          * support aggregations on date nanos
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Bucket.java

@@ -252,7 +252,7 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
                 assert DataType.isTemporalAmount(buckets.dataType()) : "Unexpected span data type [" + buckets.dataType() + "]";
                 preparedRounding = DateTrunc.createRounding(buckets.fold(), DEFAULT_TZ);
             }
-            return DateTrunc.evaluator(source(), toEvaluator.apply(field), preparedRounding);
+            return DateTrunc.evaluator(field.dataType(), source(), toEvaluator.apply(field), preparedRounding);
         }
         if (field.dataType().isNumeric()) {
             double roundTo;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/date/DateTrunc.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.time.DateUtils;
 import org.elasticsearch.compute.ann.Evaluator;
 import org.elasticsearch.compute.ann.Fixed;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -31,12 +32,14 @@
 import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
-import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isDate;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
 
 public class DateTrunc extends EsqlScalarFunction {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
@@ -45,6 +48,15 @@ public class DateTrunc extends EsqlScalarFunction {
         DateTrunc::new
     );
 
+    @FunctionalInterface
+    public interface DateTruncFactoryProvider {
+        ExpressionEvaluator.Factory apply(Source source, ExpressionEvaluator.Factory lhs, Rounding.Prepared rounding);
+    }
+
+    private static final Map<DataType, DateTruncFactoryProvider> evaluatorMap = Map.ofEntries(
+        Map.entry(DATETIME, DateTruncDatetimeEvaluator.Factory::new),
+        Map.entry(DATE_NANOS, DateTruncDateNanosEvaluator.Factory::new)
+    );
     private final Expression interval;
     private final Expression timestampField;
     protected static final ZoneId DEFAULT_TZ = ZoneOffset.UTC;
@@ -108,20 +120,28 @@ protected TypeResolution resolveType() {
             return new TypeResolution("Unresolved children");
         }
 
+        String operationName = sourceText();
         return isType(interval, DataType::isTemporalAmount, sourceText(), FIRST, "dateperiod", "timeduration").and(
-            isDate(timestampField, sourceText(), SECOND)
+            isType(timestampField, evaluatorMap::containsKey, operationName, SECOND, "date_nanos or datetime")
         );
     }
 
     public DataType dataType() {
-        return DataType.DATETIME;
+        // Default to DATETIME in the case of nulls. This mimics the behavior before DATE_NANOS support
+        return timestampField.dataType() == DataType.NULL ? DATETIME : timestampField.dataType();
     }
 
-    @Evaluator
-    static long process(long fieldVal, @Fixed Rounding.Prepared rounding) {
+    @Evaluator(extraName = "Datetime")
+    static long processDatetime(long fieldVal, @Fixed Rounding.Prepared rounding) {
         return rounding.round(fieldVal);
     }
 
+    @Evaluator(extraName = "DateNanos")
+    static long processDateNanos(long fieldVal, @Fixed Rounding.Prepared rounding) {
+        // Currently, ES|QL doesn't support rounding to sub-millisecond values, so it's safe to cast before rounding.
+        return DateUtils.toNanoSeconds(rounding.round(DateUtils.toMilliSeconds(fieldVal)));
+    }
+
     @Override
     public Expression replaceChildren(List<Expression> newChildren) {
         return new DateTrunc(source(), newChildren.get(0), newChildren.get(1));
@@ -214,14 +234,15 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
                 "Function [" + sourceText() + "] has invalid interval [" + interval.sourceText() + "]. " + e.getMessage()
             );
         }
-        return evaluator(source(), fieldEvaluator, DateTrunc.createRounding(foldedInterval, DEFAULT_TZ));
+        return evaluator(dataType(), source(), fieldEvaluator, DateTrunc.createRounding(foldedInterval, DEFAULT_TZ));
     }
 
     public static ExpressionEvaluator.Factory evaluator(
+        DataType forType,
         Source source,
         ExpressionEvaluator.Factory fieldEvaluator,
         Rounding.Prepared rounding
     ) {
-        return new DateTruncEvaluator.Factory(source, fieldEvaluator, rounding);
+        return evaluatorMap.get(forType).apply(source, fieldEvaluator, rounding);
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

@@ -1118,21 +1118,24 @@ public void testDateTruncOnInt() {
         verifyUnsupported("""
             from test
             | eval date_trunc(1 month, int)
-            """, "second argument of [date_trunc(1 month, int)] must be [datetime], found value [int] type [integer]");
+            """, "second argument of [date_trunc(1 month, int)] must be [date_nanos or datetime], found value [int] type [integer]");
     }
 
     public void testDateTruncOnFloat() {
         verifyUnsupported("""
             from test
             | eval date_trunc(1 month, float)
-            """, "second argument of [date_trunc(1 month, float)] must be [datetime], found value [float] type [double]");
+            """, "second argument of [date_trunc(1 month, float)] must be [date_nanos or datetime], found value [float] type [double]");
     }
 
     public void testDateTruncOnText() {
-        verifyUnsupported("""
-            from test
-            | eval date_trunc(1 month, keyword)
-            """, "second argument of [date_trunc(1 month, keyword)] must be [datetime], found value [keyword] type [keyword]");
+        verifyUnsupported(
+            """
+                from test
+                | eval date_trunc(1 month, keyword)
+                """,
+            "second argument of [date_trunc(1 month, keyword)] must be [date_nanos or datetime], found value [keyword] type [keyword]"
+        );
     }
 
     public void testDateTruncWithNumericInterval() {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/BucketTests.java

@@ -87,7 +87,8 @@ private static void dateCases(List<TestCaseSupplier> suppliers, String name, Lon
                     args.add(dateBound("to", toType, "2023-03-01T09:00:00.00Z"));
                     return new TestCaseSupplier.TestCase(
                         args,
-                        "DateTruncEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding[DAY_OF_MONTH in Z][fixed to midnight]]",
+                        "DateTruncDatetimeEvaluator[fieldVal=Attribute[channel=0], "
+                            + "rounding=Rounding[DAY_OF_MONTH in Z][fixed to midnight]]",
                         DataType.DATETIME,
                         resultsMatcher(args)
                     );
@@ -101,7 +102,7 @@ private static void dateCases(List<TestCaseSupplier> suppliers, String name, Lon
                     args.add(dateBound("to", toType, "2023-02-17T12:00:00Z"));
                     return new TestCaseSupplier.TestCase(
                         args,
-                        "DateTruncEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding[3600000 in Z][fixed]]",
+                        "DateTruncDatetimeEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding[3600000 in Z][fixed]]",
                         DataType.DATETIME,
                         equalTo(Rounding.builder(Rounding.DateTimeUnit.HOUR_OF_DAY).build().prepareForUnknown().round(date.getAsLong()))
                     );
@@ -134,7 +135,7 @@ private static void dateCasesWithSpan(
             args.add(new TestCaseSupplier.TypedData(span, spanType, "buckets").forceLiteral());
             return new TestCaseSupplier.TestCase(
                 args,
-                "DateTruncEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding" + spanStr + "]",
+                "DateTruncDatetimeEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding" + spanStr + "]",
                 DataType.DATETIME,
                 resultsMatcher(args)
             );