[ESQL] Add support for date trunc on date nanos type

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

@@ -362,3 +362,21 @@ 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
+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

@@ -318,6 +318,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/expression/function/scalar/date/DateTruncRoundingTests.java

@@ -15,7 +15,7 @@
 import java.time.Period;
 
 import static org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc.createRounding;
-import static org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc.process;
+import static org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc.processDatetime;
 import static org.hamcrest.Matchers.containsString;
 
 /**
@@ -97,10 +97,10 @@ public void testCreateRoundingNullInterval() {
     public void testDateTruncFunction() {
         long ts = toMillis("2023-02-17T10:25:33.38Z");
 
-        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> process(ts, createRounding(Period.ofDays(-1))));
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> processDatetime(ts, createRounding(Period.ofDays(-1))));
         assertThat(e.getMessage(), containsString("Zero or negative time interval is not supported"));
 
-        e = expectThrows(IllegalArgumentException.class, () -> process(ts, createRounding(Duration.ofHours(-1))));
+        e = expectThrows(IllegalArgumentException.class, () -> processDatetime(ts, createRounding(Duration.ofHours(-1))));
         assertThat(e.getMessage(), containsString("Zero or negative time interval is not supported"));
     }