Merge remote-tracking branch 'upstream/main' into 04212025/PauseIndexingES11516

refresh

Author: ankikuma

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java

@@ -21,10 +21,13 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.data.LongVector;
+import org.elasticsearch.compute.data.OrdinalBytesRefVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AggregationOperator;
 import org.elasticsearch.compute.operator.DriverContext;
@@ -282,11 +285,18 @@ private static Block dataBlock(int groups, String dataType) {
         int blockLength = blockLength(groups);
         return switch (dataType) {
             case BYTES_REF -> {
-                try (BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(blockLength)) {
+                try (
+                    BytesRefVector.Builder dict = blockFactory.newBytesRefVectorBuilder(blockLength);
+                    IntVector.Builder ords = blockFactory.newIntVectorBuilder(blockLength)
+                ) {
+                    final int dictLength = Math.min(blockLength, KEYWORDS.length);
+                    for (int i = 0; i < dictLength; i++) {
+                        dict.appendBytesRef(KEYWORDS[i]);
+                    }
                     for (int i = 0; i < blockLength; i++) {
-                        builder.appendBytesRef(KEYWORDS[i % KEYWORDS.length]);
+                        ords.appendInt(i % dictLength);
                     }
-                    yield builder.build();
+                    yield new OrdinalBytesRefVector(ords.build(), dict.build()).asBlock();
                 }
             }
             case INT -> {

docs/changelog/127849.yaml

@@ -0,0 +1,5 @@
+pr: 127849
+summary: Optimize ordinal inputs in Values aggregation
+area: "ES|QL"
+type: enhancement
+issues: []

muted-tests.yml

@@ -459,6 +459,12 @@ tests:
 - class: org.elasticsearch.xpack.inference.qa.mixed.HuggingFaceServiceMixedIT
   method: testElser
   issue: https://github.com/elastic/elasticsearch/issues/127876
+- class: org.elasticsearch.xpack.inference.qa.mixed.OpenAIServiceMixedIT
+  method: testOpenAiCompletions
+  issue: https://github.com/elastic/elasticsearch/issues/127878
+- class: org.elasticsearch.xpack.esql.qa.mixed.EsqlClientYamlIT
+  method: test {p0=esql/120_profile/avg 8.14 or after}
+  issue: https://github.com/elastic/elasticsearch/issues/127879
 
 # Examples:
 #

server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java

@@ -1998,34 +1998,31 @@ public Builder numberOfShards(int numberOfShards) {
 
         /**
          * Builder to create IndexMetadata that has an increased shard count (used for re-shard).
-         * The new shard count must be a multiple of the original shardcount.
+         * The new shard count must be a multiple of the original shardcount as well as a factor
+         * of routingNumShards.
          * We do not support shrinking the shard count.
-         * @param shardCount   updated shardCount
-         *
-         * TODO: Check if this.version needs to be incremented
+         * @param targetShardCount   target shard count after resharding
          */
-        public Builder reshardAddShards(int shardCount) {
-            // Assert routingNumShards is null ?
-            // Assert numberOfShards > 0
-            if (shardCount % numberOfShards() != 0) {
+        public Builder reshardAddShards(int targetShardCount) {
+            final int sourceNumShards = numberOfShards();
+            if (targetShardCount % sourceNumShards != 0) {
                 throw new IllegalArgumentException(
                     "New shard count ["
-                        + shardCount
+                        + targetShardCount
                         + "] should be a multiple"
                         + " of current shard count ["
-                        + numberOfShards()
+                        + sourceNumShards
                         + "] for ["
                         + index
                         + "]"
                 );
             }
-            IndexVersion indexVersionCreated = indexCreatedVersion(settings);
-            settings = Settings.builder().put(settings).put(SETTING_NUMBER_OF_SHARDS, shardCount).build();
-            var newPrimaryTerms = new long[shardCount];
+            settings = Settings.builder().put(settings).put(SETTING_NUMBER_OF_SHARDS, targetShardCount).build();
+            var newPrimaryTerms = new long[targetShardCount];
             Arrays.fill(newPrimaryTerms, this.primaryTerms.length, newPrimaryTerms.length, SequenceNumbers.UNASSIGNED_PRIMARY_TERM);
             System.arraycopy(primaryTerms, 0, newPrimaryTerms, 0, this.primaryTerms.length);
             primaryTerms = newPrimaryTerms;
-            routingNumShards = MetadataCreateIndexService.calculateNumRoutingShards(shardCount, indexVersionCreated);
+            routingNumShards = MetadataCreateIndexService.getIndexNumberOfRoutingShards(settings, sourceNumShards, this.routingNumShards);
             return this;
         }
 
@@ -3034,7 +3031,7 @@ public static ShardId selectCloneShard(int shardId, IndexMetadata sourceIndexMet
         return new ShardId(sourceIndexMetadata.getIndex(), shardId);
     }
 
-    private static void assertSplitMetadata(int numSourceShards, int numTargetShards, IndexMetadata sourceIndexMetadata) {
+    public static void assertSplitMetadata(int numSourceShards, int numTargetShards, IndexMetadata sourceIndexMetadata) {
         if (numSourceShards > numTargetShards) {
             throw new IllegalArgumentException(
                 "the number of source shards ["

server/src/main/java/org/elasticsearch/cluster/metadata/MetadataCreateIndexService.java

@@ -1282,10 +1282,23 @@ private static void validateSoftDeleteSettings(Settings indexSettings) {
      * it will return the value configured for that index.
      */
     static int getIndexNumberOfRoutingShards(Settings indexSettings, @Nullable IndexMetadata sourceMetadata) {
+        final int routingNumShards = getIndexNumberOfRoutingShards(
+            indexSettings,
+            sourceMetadata == null ? 1 : sourceMetadata.getNumberOfShards(),
+            sourceMetadata == null ? 0 : sourceMetadata.getRoutingNumShards()
+        );
+        return routingNumShards;
+    }
+
+    /**
+     * Calculates the number of routing shards based on the configured value in indexSettings or if recovering from another index
+     * it will return the value configured for that index.
+     */
+    static int getIndexNumberOfRoutingShards(Settings indexSettings, final int sourceNumShards, final int sourceRoutingNumShards) {
         final int numTargetShards = INDEX_NUMBER_OF_SHARDS_SETTING.get(indexSettings);
         final IndexVersion indexVersionCreated = IndexMetadata.SETTING_INDEX_VERSION_CREATED.get(indexSettings);
         final int routingNumShards;
-        if (sourceMetadata == null || sourceMetadata.getNumberOfShards() == 1) {
+        if (sourceNumShards == 1) {
             // in this case we either have no index to recover from or
             // we have a source index with 1 shard and without an explicit split factor
             // or one that is valid in that case we can split into whatever and auto-generate a new factor.
@@ -1299,7 +1312,7 @@ static int getIndexNumberOfRoutingShards(Settings indexSettings, @Nullable Index
         } else {
             assert IndexMetadata.INDEX_NUMBER_OF_ROUTING_SHARDS_SETTING.exists(indexSettings) == false
                 : "index.number_of_routing_shards should not be present on the target index on resize";
-            routingNumShards = sourceMetadata.getRoutingNumShards();
+            routingNumShards = sourceRoutingNumShards;
         }
         return routingNumShards;
     }

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/GroupingAggregatorImplementer.java

@@ -35,6 +35,7 @@
 
 import static java.util.stream.Collectors.joining;
 import static org.elasticsearch.compute.gen.AggregatorImplementer.capitalize;
+import static org.elasticsearch.compute.gen.Methods.optionalStaticMethod;
 import static org.elasticsearch.compute.gen.Methods.requireAnyArgs;
 import static org.elasticsearch.compute.gen.Methods.requireAnyType;
 import static org.elasticsearch.compute.gen.Methods.requireArgs;
@@ -336,10 +337,32 @@ private MethodSpec prepareProcessPage() {
             builder.beginControlFlow("if (valuesBlock.mayHaveNulls())");
             builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
             builder.endControlFlow();
-            builder.addStatement("return $L", addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra)));
+            if (shouldWrapAddInput(blockType(aggParam.type()))) {
+                builder.addStatement(
+                    "var addInput = $L",
+                    addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra))
+                );
+                builder.addStatement("return $T.wrapAddInput(addInput, state, valuesBlock)", declarationType);
+            } else {
+                builder.addStatement(
+                    "return $L",
+                    addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra))
+                );
+            }
         }
         builder.endControlFlow();
-        builder.addStatement("return $L", addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra)));
+        if (shouldWrapAddInput(vectorType(aggParam.type()))) {
+            builder.addStatement(
+                "var addInput = $L",
+                addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra))
+            );
+            builder.addStatement("return $T.wrapAddInput(addInput, state, valuesVector)", declarationType);
+        } else {
+            builder.addStatement(
+                "return $L",
+                addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra))
+            );
+        }
         return builder.build();
     }
 
@@ -526,6 +549,15 @@ private void combineRawInputForArray(MethodSpec.Builder builder, String arrayVar
         warningsBlock(builder, () -> builder.addStatement("$T.combine(state, groupId, $L)", declarationType, arrayVariable));
     }
 
+    private boolean shouldWrapAddInput(TypeName valuesType) {
+        return optionalStaticMethod(
+            declarationType,
+            requireType(GROUPING_AGGREGATOR_FUNCTION_ADD_INPUT),
+            requireName("wrapAddInput"),
+            requireArgs(requireType(GROUPING_AGGREGATOR_FUNCTION_ADD_INPUT), requireType(aggState.declaredType()), requireType(valuesType))
+        ) != null;
+    }
+
     private void warningsBlock(MethodSpec.Builder builder, Runnable block) {
         if (warnExceptions.isEmpty() == false) {
             builder.beginControlFlow("try");

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/Methods.java

@@ -59,20 +59,31 @@ static ExecutableElement requireStaticMethod(
         TypeMatcher returnTypeMatcher,
         NameMatcher nameMatcher,
         ArgumentMatcher argumentMatcher
+    ) {
+        ExecutableElement method = optionalStaticMethod(declarationType, returnTypeMatcher, nameMatcher, argumentMatcher);
+        if (method == null) {
+            var message = nameMatcher.names.size() == 1 ? "Requires method: " : "Requires one of methods: ";
+            var signatures = nameMatcher.names.stream()
+                .map(name -> "public static " + returnTypeMatcher + " " + declarationType + "#" + name + "(" + argumentMatcher + ")")
+                .collect(joining(" or "));
+            throw new IllegalArgumentException(message + signatures);
+        }
+        return method;
+    }
+
+    static ExecutableElement optionalStaticMethod(
+        TypeElement declarationType,
+        TypeMatcher returnTypeMatcher,
+        NameMatcher nameMatcher,
+        ArgumentMatcher argumentMatcher
     ) {
         return typeAndSuperType(declarationType).flatMap(type -> ElementFilter.methodsIn(type.getEnclosedElements()).stream())
             .filter(method -> method.getModifiers().contains(Modifier.STATIC))
             .filter(method -> nameMatcher.test(method.getSimpleName().toString()))
             .filter(method -> returnTypeMatcher.test(TypeName.get(method.getReturnType())))
             .filter(method -> argumentMatcher.test(method.getParameters().stream().map(it -> TypeName.get(it.asType())).toList()))
             .findFirst()
-            .orElseThrow(() -> {
-                var message = nameMatcher.names.size() == 1 ? "Requires method: " : "Requires one of methods: ";
-                var signatures = nameMatcher.names.stream()
-                    .map(name -> "public static " + returnTypeMatcher + " " + declarationType + "#" + name + "(" + argumentMatcher + ")")
-                    .collect(joining(" or "));
-                return new IllegalArgumentException(message + signatures);
-            });
+            .orElse(null);
     }
 
     static NameMatcher requireName(String... names) {