Update Count types docs and tests

Author: ivancea

docs/reference/query-languages/esql/_snippets/functions/types/count.md

@@ -79,6 +79,7 @@ public Count(
                 "cartesian_point",
                 "cartesian_shape",
                 "date",
+                "date_nanos",
                 "double",
                 "geo_point",
                 "geo_shape",

docs/reference/query-languages/esql/kibana/definition/functions/count.json

@@ -36,6 +36,14 @@ public final class MultiRowTestCaseSupplier {
 
     private MultiRowTestCaseSupplier() {}
 
+    public static List<TypedDataSupplier> nullCases(int minRows, int maxRows) {
+        List<TypedDataSupplier> cases = new ArrayList<>();
+
+        addSuppliers(cases, minRows, maxRows, "null", DataType.NULL, () -> null);
+
+        return cases;
+    }
+
     public static List<TypedDataSupplier> intCases(int minRows, int maxRows, int min, int max, boolean includeZero) {
         List<TypedDataSupplier> cases = new ArrayList<>();
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Count.java

@@ -21,6 +21,7 @@
 import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -37,11 +38,13 @@ public static Iterable<Object[]> parameters() {
         var suppliers = new ArrayList<TestCaseSupplier>();
 
         Stream.of(
+            MultiRowTestCaseSupplier.nullCases(1, 1000),
             MultiRowTestCaseSupplier.intCases(1, 1000, Integer.MIN_VALUE, Integer.MAX_VALUE, true),
             MultiRowTestCaseSupplier.longCases(1, 1000, Long.MIN_VALUE, Long.MAX_VALUE, true),
             MultiRowTestCaseSupplier.ulongCases(1, 1000, BigInteger.ZERO, UNSIGNED_LONG_MAX, true),
             MultiRowTestCaseSupplier.doubleCases(1, 1000, -Double.MAX_VALUE, Double.MAX_VALUE, true),
             MultiRowTestCaseSupplier.dateCases(1, 1000),
+            MultiRowTestCaseSupplier.dateNanosCases(1, 1000),
             MultiRowTestCaseSupplier.booleanCases(1, 1000),
             MultiRowTestCaseSupplier.ipCases(1, 1000),
             MultiRowTestCaseSupplier.versionCases(1, 1000),
@@ -54,10 +57,12 @@ public static Iterable<Object[]> parameters() {
 
         // No rows
         for (var dataType : List.of(
+            DataType.NULL,
             DataType.INTEGER,
             DataType.LONG,
             DataType.DOUBLE,
             DataType.DATETIME,
+            DataType.DATE_NANOS,
             DataType.BOOLEAN,
             DataType.IP,
             DataType.VERSION,
@@ -93,13 +98,13 @@ protected Expression build(Source source, List<Expression> args) {
     private static TestCaseSupplier makeSupplier(TestCaseSupplier.TypedDataSupplier fieldSupplier) {
         return new TestCaseSupplier(fieldSupplier.name(), List.of(fieldSupplier.type()), () -> {
             var fieldTypedData = fieldSupplier.get();
-            var rowCount = fieldTypedData.multiRowData().size();
+            var rowCount = fieldTypedData.multiRowData().stream().filter(Objects::nonNull).count();
 
             return new TestCaseSupplier.TestCase(
                 List.of(fieldTypedData),
                 "Count[field=Attribute[channel=0]]",
                 DataType.LONG,
-                equalTo((long) rowCount)
+                equalTo(rowCount)
             );
         });
     }