Fix QuantileStates for non-existent group (#99558)

OrdinalsGroupingOperator can retrieve the state of a group that doesn't exist in the grouping state. The grouping state needs to return null in this case. However, QuantileStates doesn't follow this.
elastic · Sep 18, 2023 · 9ce92a9 · 9ce92a9
1 parent 91e4662
commit 9ce92a9
Show file tree

Hide file tree

Showing 8 changed files with 42 additions and 9 deletions.
diff --git a/...n/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationDoubleAggregator.java b/...n/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationDoubleAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selected) {

diff --git a/...main/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationIntAggregator.java b/...main/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationIntAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selected) {

diff --git a/...ain/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationLongAggregator.java b/...ain/java/org/elasticsearch/compute/aggregation/MedianAbsoluteDeviationLongAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selected) {

diff --git a/...mpute/src/main/java/org/elasticsearch/compute/aggregation/PercentileDoubleAggregator.java b/...mpute/src/main/java/org/elasticsearch/compute/aggregation/PercentileDoubleAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selected) {

diff --git a/.../compute/src/main/java/org/elasticsearch/compute/aggregation/PercentileIntAggregator.java b/.../compute/src/main/java/org/elasticsearch/compute/aggregation/PercentileIntAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selected) {

diff --git a/...compute/src/main/java/org/elasticsearch/compute/aggregation/PercentileLongAggregator.java b/...compute/src/main/java/org/elasticsearch/compute/aggregation/PercentileLongAggregator.java
@@ -57,7 +57,7 @@ public static void combineStates(
         QuantileStates.GroupingState state,
         int statePosition
     ) {
-        current.add(currentGroupId, state.get(statePosition));
+        current.add(currentGroupId, state.getOrNull(statePosition));
     }
 
     public static Block evaluateFinal(QuantileStates.GroupingState state, IntVector selectedGroups) {

diff --git a/...ugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/QuantileStates.java b/...ugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/QuantileStates.java
@@ -148,8 +148,12 @@ void add(int groupId, BytesRef other) {
             getOrAddGroup(groupId).add(deserializeDigest(other));
         }
 
-        TDigestState get(int position) {
-            return digests.get(position);
+        TDigestState getOrNull(int position) {
+            if (position < digests.size()) {
+                return digests.get(position);
+            } else {
+                return null;
+            }
         }
 
         /** Extracts an intermediate view of the contents of this state.  */
@@ -161,7 +165,7 @@ public void toIntermediate(Block[] blocks, int offset, IntVector selected) {
                 int group = selected.getInt(i);
                 TDigestState state;
                 if (group < digests.size()) {
-                    state = get(group);
+                    state = getOrNull(group);
                     if (state == null) {
                         state = TDigestState.create(DEFAULT_COMPRESSION);
                     }

diff --git a/...n/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java b/...n/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java
@@ -37,6 +37,7 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.OptionalDouble;
 import java.util.concurrent.CountDownLatch;
@@ -1053,6 +1054,34 @@ public void testReturnNoNestedDocuments() throws IOException, ExecutionException
         assertNoNestedDocuments("from " + alias + " | where data >= 50", Arrays.stream(countValuesGreaterThanFifty).sum(), 50L, 100L);
     }
 
+    public void testGroupingMultiValueByOrdinals() {
+        String indexName = "test-ordinals";
+        assertAcked(client().admin().indices().prepareCreate(indexName).setMapping("kw", "type=keyword", "v", "type=long").get());
+        int numDocs = randomIntBetween(10, 200);
+        for (int i = 0; i < numDocs; i++) {
+            Map<String, Object> source = new HashMap<>();
+            source.put("kw", "key-" + randomIntBetween(1, 20));
+            List<Integer> values = new ArrayList<>();
+            int numValues = between(0, 2);
+            for (int v = 0; v < numValues; v++) {
+                values.add(randomIntBetween(1, 1000));
+            }
+            if (values.isEmpty() == false) {
+                source.put("v", values);
+            }
+            client().prepareIndex(indexName).setSource(source).get();
+            if (randomInt(100) < 20) {
+                client().admin().indices().prepareRefresh(indexName).get();
+            }
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+        var functions = List.of("min(v)", "max(v)", "count_distinct(v)", "count(v)", "sum(v)", "avg(v)", "percentile(v, 90)");
+        for (String fn : functions) {
+            String query = String.format(Locale.ROOT, "from %s | stats s = %s by kw", indexName, fn);
+            run(query);
+        }
+    }
+
     private void createNestedMappingIndex(String indexName) throws IOException {
         XContentBuilder builder = JsonXContent.contentBuilder();
         builder.startObject();