From b11cd0af99fad5ce0ec48e4228ba1951ad478c9d Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Tue, 10 Oct 2023 23:55:38 -0700
Subject: [PATCH] Relax ValueSources check in OrdinalsGroupingOperator
 (#100566) (#100654)

ValuesSource can be Null instead of Bytes when a shard has no data for a
specific field. This PR relaxes the check for ValueSources in the
OrdinalsGroupingOperator.

We will need to add more tests for OrdinalsGroupingOperator.

Closes #100438
---
 .../operator/OrdinalsGroupingOperator.java    |  6 ----
 .../xpack/esql/action/EsqlActionIT.java       | 34 +++++++++++++++++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java
index 4dab7faa2a074..7c930118903cf 100644
--- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java
+++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java
@@ -105,12 +105,6 @@ public OrdinalsGroupingOperator(
         DriverContext driverContext
     ) {
         Objects.requireNonNull(aggregatorFactories);
-        boolean bytesValues = sources.get(0).source() instanceof ValuesSource.Bytes;
-        for (int i = 1; i < sources.size(); i++) {
-            if (sources.get(i).source() instanceof ValuesSource.Bytes != bytesValues) {
-                throw new IllegalStateException("ValuesSources are mismatched");
-            }
-        }
         this.sources = sources;
         this.docChannel = docChannel;
         this.groupingField = groupingField;
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java
index f10ca17d741d8..0017a8600a013 100644
--- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java
+++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java
@@ -35,6 +35,7 @@
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -1187,6 +1188,39 @@ public void testGroupingMultiValueByOrdinals() {
         }
     }
 
+    public void testUnsupportedTypesOrdinalGrouping() {
+        assertAcked(
+            client().admin().indices().prepareCreate("index-1").setMapping("f1", "type=keyword", "f2", "type=keyword", "v", "type=long")
+        );
+        assertAcked(
+            client().admin().indices().prepareCreate("index-2").setMapping("f1", "type=object", "f2", "type=keyword", "v", "type=long")
+        );
+        Map<String, Long> groups = new HashMap<>();
+        int numDocs = randomIntBetween(10, 20);
+        for (int i = 0; i < numDocs; i++) {
+            String k = randomFrom("a", "b", "c");
+            long v = randomIntBetween(1, 10);
+            groups.merge(k, v, Long::sum);
+            groups.merge(null, v, Long::sum); // null group
+            client().prepareIndex("index-1").setSource("f1", k, "v", v).get();
+            client().prepareIndex("index-2").setSource("f2", k, "v", v).get();
+        }
+        client().admin().indices().prepareRefresh("index-1", "index-2").get();
+        for (String field : List.of("f1", "f2")) {
+            try (var resp = run("from index-1,index-2 | stats sum(v) by " + field)) {
+                Iterator<Iterator<Object>> values = resp.values();
+                Map<String, Long> actual = new HashMap<>();
+                while (values.hasNext()) {
+                    Iterator<Object> row = values.next();
+                    Long v = (Long) row.next();
+                    String k = (String) row.next();
+                    actual.put(k, v);
+                }
+                assertThat(actual, equalTo(groups));
+            }
+        }
+    }
+
     private void createNestedMappingIndex(String indexName) throws IOException {
         XContentBuilder builder = JsonXContent.contentBuilder();
         builder.startObject();