[Transform] Optimize composite agg execution using ordered groupings (#75424)

Automatically reorder group_by for composite aggs, ensuring date histogram
group by comes first. The order is only changed for execution, the provided
config remains unchanged.

In case of 2 group_by's of the same order type, the configuration order is
respected. Script and runtime field based group_by's are penalized.

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/pivot/PivotConfig.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.core.transform.transforms.pivot;
 
 import org.elasticsearch.action.ActionRequestValidationException;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
@@ -18,8 +17,8 @@ import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.search.aggregations.MultiBucketConsumerService;
-import org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
 import org.elasticsearch.xpack.core.transform.TransformField;
 import org.elasticsearch.xpack.core.transform.utils.ExceptionsHelper;
 
@@ -27,7 +26,6 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map.Entry;
 import java.util.Objects;
 
 import static org.elasticsearch.action.ValidateActions.addValidationError;
@@ -111,23 +109,6 @@ public class PivotConfig implements Writeable, ToXContentObject {
         return builder;
     }
 
-    public void toCompositeAggXContent(XContentBuilder builder) throws IOException {
-        builder.startObject();
-        builder.field(CompositeAggregationBuilder.SOURCES_FIELD_NAME.getPreferredName());
-        builder.startArray();
-
-        for (Entry<String, SingleGroupSource> groupBy : groups.getGroups().entrySet()) {
-            builder.startObject();
-            builder.startObject(groupBy.getKey());
-            builder.field(groupBy.getValue().getType().value(), groupBy.getValue());
-            builder.endObject();
-            builder.endObject();
-        }
-
-        builder.endArray();
-        builder.endObject(); // sources
-    }
-
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         groups.writeTo(out);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/TransformConfigTests.java

@@ -11,7 +11,6 @@ import org.elasticsearch.Version;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.Writeable.Reader;
-import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.common.xcontent.DeprecationHandler;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -19,8 +18,7 @@ import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
-import org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
-import org.elasticsearch.search.aggregations.bucket.composite.CompositeValuesSourceBuilder;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.xpack.core.common.validation.SourceDestValidator.RemoteClusterMinimumVersionValidation;
 import org.elasticsearch.xpack.core.common.validation.SourceDestValidator.SourceDestValidation;
 import org.elasticsearch.xpack.core.transform.AbstractSerializingTransformTestCase;
@@ -36,19 +34,17 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
-import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.test.TestMatchers.matchesPattern;
 import static org.elasticsearch.xpack.core.transform.transforms.DestConfigTests.randomDestConfig;
 import static org.elasticsearch.xpack.core.transform.transforms.SourceConfigTests.randomInvalidSourceConfig;
 import static org.elasticsearch.xpack.core.transform.transforms.SourceConfigTests.randomSourceConfig;
+import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.contains;
 
 public class TransformConfigTests extends AbstractSerializingTransformTestCase<TransformConfig> {
 
@@ -598,7 +594,9 @@ public class TransformConfigTests extends AbstractSerializingTransformTestCase<T
 
     public void testGroupByStayInOrder() throws IOException {
         String json = "{"
-            + " \"id\" : \"" + transformId +"\","
+            + " \"id\" : \""
+            + transformId
+            + "\","
             + " \"source\" : {"
             + "   \"index\":\"src\""
             + "},"
@@ -625,20 +623,12 @@ public class TransformConfigTests extends AbstractSerializingTransformTestCase<T
             + "} } } } }";
         TransformConfig transformConfig = createTransformConfigFromString(json, transformId, true);
         List<String> originalGroups = new ArrayList<>(transformConfig.getPivotConfig().getGroupConfig().getGroups().keySet());
-        assertThat(
-            originalGroups,
-            contains("time", "alert", "id")
-        );
+        assertThat(originalGroups, contains("time", "alert", "id"));
         for (int runs = 0; runs < NUMBER_OF_TEST_RUNS; runs++) {
             // Wire serialization order guarantees
             TransformConfig serialized = this.copyInstance(transformConfig);
             List<String> serializedGroups = new ArrayList<>(serialized.getPivotConfig().getGroupConfig().getGroups().keySet());
             assertThat(serializedGroups, equalTo(originalGroups));
-            CompositeAggregationBuilder compositeAggregationBuilder = createCompositeAggregationSources(serialized.getPivotConfig());
-            assertThat(
-                compositeAggregationBuilder.sources().stream().map(CompositeValuesSourceBuilder::name).collect(Collectors.toList()),
-                equalTo(originalGroups)
-            );
 
             // Now test xcontent serialization and parsing on wire serialized object
             XContentType xContentType = randomFrom(XContentType.values()).canonical();
@@ -647,11 +637,6 @@ public class TransformConfigTests extends AbstractSerializingTransformTestCase<T
             TransformConfig parsed = doParseInstance(parser);
             List<String> parsedGroups = new ArrayList<>(parsed.getPivotConfig().getGroupConfig().getGroups().keySet());
             assertThat(parsedGroups, equalTo(originalGroups));
-            compositeAggregationBuilder = createCompositeAggregationSources(parsed.getPivotConfig());
-            assertThat(
-                compositeAggregationBuilder.sources().stream().map(CompositeValuesSourceBuilder::name).collect(Collectors.toList()),
-                equalTo(originalGroups)
-            );
         }
     }
 
@@ -665,21 +650,4 @@ public class TransformConfigTests extends AbstractSerializingTransformTestCase<T
         return TransformConfig.fromXContent(parser, id, lenient);
     }
 
-    private CompositeAggregationBuilder createCompositeAggregationSources(PivotConfig config) throws IOException {
-        CompositeAggregationBuilder compositeAggregation;
-
-        try (XContentBuilder builder = jsonBuilder()) {
-            config.toCompositeAggXContent(builder);
-            XContentParser parser = builder.generator()
-                .contentType()
-                .xContent()
-                .createParser(
-                    xContentRegistry(),
-                    DeprecationHandler.THROW_UNSUPPORTED_OPERATION,
-                    BytesReference.bytes(builder).streamInput()
-                );
-            compositeAggregation = CompositeAggregationBuilder.PARSER.parse(parser, "composite_agg");
-        }
-        return compositeAggregation;
-    }
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/pivot/DateHistogramGroupSourceTests.java

@@ -30,12 +30,20 @@ public class DateHistogramGroupSourceTests extends AbstractSerializingTestCase<D
         return randomDateHistogramGroupSource(Version.CURRENT);
     }
 
+    public static DateHistogramGroupSource randomDateHistogramGroupSourceNoScript() {
+        return randomDateHistogramGroupSource(Version.CURRENT, false);
+    }
+
     public static DateHistogramGroupSource randomDateHistogramGroupSource(Version version) {
+        return randomDateHistogramGroupSource(version, randomBoolean());
+    }
+
+    public static DateHistogramGroupSource randomDateHistogramGroupSource(Version version, boolean withScript) {
         ScriptConfig scriptConfig = null;
         String field;
 
         // either a field or a script must be specified, it's possible to have both, but disallowed to have none
-        if (version.onOrAfter(Version.V_7_7_0) && randomBoolean()) {
+        if (version.onOrAfter(Version.V_7_7_0) && withScript) {
             scriptConfig = ScriptConfigTests.randomScriptConfig();
             field = randomBoolean() ? null : randomAlphaOfLengthBetween(1, 20);
         } else {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/pivot/HistogramGroupSourceTests.java

@@ -20,12 +20,20 @@ public class HistogramGroupSourceTests extends AbstractSerializingTestCase<Histo
         return randomHistogramGroupSource(Version.CURRENT);
     }
 
+    public static HistogramGroupSource randomHistogramGroupSourceNoScript() {
+        return randomHistogramGroupSource(Version.CURRENT, false);
+    }
+
     public static HistogramGroupSource randomHistogramGroupSource(Version version) {
+        return randomHistogramGroupSource(version, randomBoolean());
+    }
+
+    public static HistogramGroupSource randomHistogramGroupSource(Version version, boolean withScript) {
         ScriptConfig scriptConfig = null;
         String field;
 
         // either a field or a script must be specified, it's possible to have both, but disallowed to have none
-        if (version.onOrAfter(Version.V_7_7_0) && randomBoolean()) {
+        if (version.onOrAfter(Version.V_7_7_0) && withScript) {
             scriptConfig = ScriptConfigTests.randomScriptConfig();
             field = randomBoolean() ? null : randomAlphaOfLengthBetween(1, 20);
         } else {

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/pivot/TermsGroupSourceTests.java

@@ -27,12 +27,20 @@ public class TermsGroupSourceTests extends AbstractSerializingTestCase<TermsGrou
         return randomTermsGroupSource(Version.CURRENT);
     }
 
+    public static TermsGroupSource randomTermsGroupSourceNoScript() {
+        return randomTermsGroupSource(Version.CURRENT, false);
+    }
+
     public static TermsGroupSource randomTermsGroupSource(Version version) {
+        return randomTermsGroupSource(Version.CURRENT, randomBoolean());
+    }
+
+    public static TermsGroupSource randomTermsGroupSource(Version version, boolean withScript) {
         ScriptConfig scriptConfig = null;
         String field;
 
         // either a field or a script must be specified, it's possible to have both, but disallowed to have none
-        if (version.onOrAfter(Version.V_7_7_0) && randomBoolean()) {
+        if (version.onOrAfter(Version.V_7_7_0) && withScript) {
             scriptConfig = ScriptConfigTests.randomScriptConfig();
             field = randomBoolean() ? null : randomAlphaOfLengthBetween(1, 20);
         } else {
@@ -43,11 +51,6 @@ public class TermsGroupSourceTests extends AbstractSerializingTestCase<TermsGrou
         return new TermsGroupSource(field, scriptConfig, missingBucket);
     }
 
-    public static TermsGroupSource randomTermsGroupSourceNoScript() {
-        String field = randomAlphaOfLengthBetween(1, 20);
-        return new TermsGroupSource(field, null, randomBoolean());
-    }
-
     @Override
     protected TermsGroupSource doParseInstance(XContentParser parser) throws IOException {
         return TermsGroupSource.fromXContent(parser, false);

x-pack/plugin/transform/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/transform/integration/TransformProgressIT.java

@@ -168,7 +168,7 @@ public class TransformProgressIT extends ESRestTestCase {
             null
         );
 
-        Pivot pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT);
+        Pivot pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT, Collections.emptySet());
 
         TransformProgress progress = getProgress(pivot, getProgressQuery(pivot, config.getSource().getIndex(), null));
 
@@ -196,7 +196,7 @@ public class TransformProgressIT extends ESRestTestCase {
             Collections.singletonMap("every_50", new HistogramGroupSource("missing_field", null, missingBucket, 50.0))
         );
         pivotConfig = new PivotConfig(histgramGroupConfig, aggregationConfig, null);
-        pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT);
+        pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT, Collections.emptySet());
 
         progress = getProgress(
             pivot,

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/FunctionFactory.java

@@ -26,7 +26,12 @@ public final class FunctionFactory {
      */
     public static Function create(TransformConfig config) {
         if (config.getPivotConfig() != null) {
-            return new Pivot(config.getPivotConfig(), config.getSettings(), config.getVersion());
+            return new Pivot(
+                config.getPivotConfig(),
+                config.getSettings(),
+                config.getVersion(),
+                config.getSource().getScriptBasedRuntimeMappings().keySet()
+            );
         } else if (config.getLatestConfig() != null) {
             return new Latest(config.getLatestConfig());
         } else {

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/GroupByOptimizer.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.transform.transforms.pivot;
+
+import org.elasticsearch.core.Tuple;
+import org.elasticsearch.xpack.core.transform.transforms.pivot.SingleGroupSource;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public final class GroupByOptimizer {
+
+    private GroupByOptimizer() {}
+
+    /**
+     * Returns an ordered collection of group by fields in order to get better performance.
+     *
+     * The decision is based on the type and whether the input field is a indexed/runtime/script field
+     *
+     * TODO: take index sorting into account
+     *
+     * @param groups group by as defined by the user
+     * @param runtimeFields set of runtime fields
+     * @return collection in order of priority
+     */
+    static Collection<Entry<String, SingleGroupSource>> reorderGroups(Map<String, SingleGroupSource> groups, Set<String> runtimeFields) {
+        if (groups.size() == 1) {
+            return groups.entrySet();
+        }
+
+        List<Tuple<Entry<String, SingleGroupSource>, Integer>> prioritizedGroups = new ArrayList<>(groups.size());
+
+        // respect the order in the configuration by giving every entry a base priority
+        int basePriority = groups.size();
+
+        for (Entry<String, SingleGroupSource> groupBy : groups.entrySet()) {
+            // prefer indexed fields over runtime fields over scripts
+            int priority = basePriority-- + (groupBy.getValue().getScriptConfig() == null
+                ? runtimeFields.contains(groupBy.getValue().getField()) ? 250 : 500
+                : 0);
+
+            switch (groupBy.getValue().getType()) {
+                case DATE_HISTOGRAM:
+                    priority += 4000;
+                    break;
+                case HISTOGRAM:
+                    priority += 3000;
+                    break;
+                case TERMS:
+                    priority += 2000;
+                    break;
+                case GEOTILE_GRID:
+                    priority += 1000;
+                    break;
+                default:
+                    assert false : "new group source type misses priority definition";
+            }
+
+            prioritizedGroups.add(new Tuple<>(groupBy, priority));
+        }
+
+        prioritizedGroups.sort(Comparator.comparing(Tuple<Entry<String, SingleGroupSource>, Integer>::v2).reversed());
+
+        return prioritizedGroups.stream().map(x -> x.v1()).collect(Collectors.toList());
+    }
+}

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/Pivot.java

@@ -36,8 +36,11 @@ import org.elasticsearch.xpack.transform.transforms.common.AbstractCompositeAggF
 import org.elasticsearch.xpack.transform.transforms.common.DocumentConversionUtils;
 
 import java.io.IOException;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toList;
@@ -59,8 +62,8 @@ public class Pivot extends AbstractCompositeAggFunction {
      * @param settings Any miscellaneous settings for the function
      * @param version The version of the transform
      */
-    public Pivot(PivotConfig config, SettingsConfig settings, Version version) {
-        super(createCompositeAggregation(config));
+    public Pivot(PivotConfig config, SettingsConfig settings, Version version, Set<String> runtimeFields) {
+        super(createCompositeAggregation(config, runtimeFields));
         this.config = config;
         this.settings = settings;
         this.version = version == null ? Version.CURRENT : version;
@@ -71,10 +74,9 @@ public class Pivot extends AbstractCompositeAggFunction {
         for (AggregationBuilder agg : config.getAggregationConfig().getAggregatorFactories()) {
             if (TransformAggregations.isSupportedByTransform(agg.getType()) == false) {
                 listener.onFailure(
-                    new ValidationException()
-                        .addValidationError(
-                            new ParameterizedMessage("Unsupported aggregation type [{}]", agg.getType()).getFormattedMessage()
-                        )
+                    new ValidationException().addValidationError(
+                        new ParameterizedMessage("Unsupported aggregation type [{}]", agg.getType()).getFormattedMessage()
+                    )
                 );
                 return;
             }
@@ -159,8 +161,8 @@ public class Pivot extends AbstractCompositeAggFunction {
         return searchSourceBuilder.query(existsClauses).size(0).trackTotalHits(true);
     }
 
-    private static CompositeAggregationBuilder createCompositeAggregation(PivotConfig config) {
-        final CompositeAggregationBuilder compositeAggregation = createCompositeAggregationSources(config);
+    private static CompositeAggregationBuilder createCompositeAggregation(PivotConfig config, Set<String> runtimeFields) {
+        final CompositeAggregationBuilder compositeAggregation = createCompositeAggregationSources(config, runtimeFields);
 
         config.getAggregationConfig().getAggregatorFactories().forEach(compositeAggregation::subAggregation);
         config.getAggregationConfig().getPipelineAggregatorFactories().forEach(compositeAggregation::subAggregation);
@@ -168,11 +170,29 @@ public class Pivot extends AbstractCompositeAggFunction {
         return compositeAggregation;
     }
 
-    private static CompositeAggregationBuilder createCompositeAggregationSources(PivotConfig config) {
+    private static CompositeAggregationBuilder createCompositeAggregationSources(PivotConfig config, Set<String> runtimeFields) {
         CompositeAggregationBuilder compositeAggregation;
 
+        Collection<Entry<String, SingleGroupSource>> groups = GroupByOptimizer.reorderGroups(
+            config.getGroupConfig().getGroups(),
+            runtimeFields
+        );
+
         try (XContentBuilder builder = jsonBuilder()) {
-            config.toCompositeAggXContent(builder);
+            builder.startObject();
+            builder.field(CompositeAggregationBuilder.SOURCES_FIELD_NAME.getPreferredName());
+            builder.startArray();
+
+            for (Entry<String, SingleGroupSource> groupBy : groups) {
+                builder.startObject();
+                builder.startObject(groupBy.getKey());
+                builder.field(groupBy.getValue().getType().value(), groupBy.getValue());
+                builder.endObject();
+                builder.endObject();
+            }
+
+            builder.endArray();
+            builder.endObject(); // sources
             XContentParser parser = builder.generator()
                 .contentType()
                 .xContent()
@@ -180,7 +200,9 @@ public class Pivot extends AbstractCompositeAggFunction {
             compositeAggregation = CompositeAggregationBuilder.PARSER.parse(parser, COMPOSITE_AGGREGATION_NAME);
         } catch (IOException e) {
             throw new RuntimeException(
-                TransformMessages.getMessage(TransformMessages.TRANSFORM_FAILED_TO_CREATE_COMPOSITE_AGGREGATION, "pivot"), e);
+                TransformMessages.getMessage(TransformMessages.TRANSFORM_FAILED_TO_CREATE_COMPOSITE_AGGREGATION, "pivot"),
+                e
+            );
         }
         return compositeAggregation;
     }

x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/action/TransformConfigLinterTests.java

@@ -26,6 +26,7 @@ import org.elasticsearch.xpack.transform.transforms.Function;
 import org.elasticsearch.xpack.transform.transforms.latest.Latest;
 import org.elasticsearch.xpack.transform.transforms.pivot.Pivot;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -38,12 +39,12 @@ import static org.hamcrest.Matchers.is;
 public class TransformConfigLinterTests extends ESTestCase {
 
     public void testGetWarnings_Pivot_WithScriptBasedRuntimeFields() {
-        PivotConfig pivotConfig =
-            new PivotConfig(
-                GroupConfigTests.randomGroupConfig(() -> new TermsGroupSource(randomAlphaOfLengthBetween(1, 20), null, false)),
-                AggregationConfigTests.randomAggregationConfig(),
-                null);
-        Function function = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT);
+        PivotConfig pivotConfig = new PivotConfig(
+            GroupConfigTests.randomGroupConfig(() -> new TermsGroupSource(randomAlphaOfLengthBetween(1, 20), null, false)),
+            AggregationConfigTests.randomAggregationConfig(),
+            null
+        );
+        Function function = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT, Collections.emptySet());
         SourceConfig sourceConfig = SourceConfigTests.randomSourceConfig();
         assertThat(TransformConfigLinter.getWarnings(function, sourceConfig, null), is(empty()));
 
@@ -51,19 +52,25 @@ public class TransformConfigLinterTests extends ESTestCase {
 
         assertThat(TransformConfigLinter.getWarnings(function, sourceConfig, syncConfig), is(empty()));
 
-        Map<String, Object> runtimeMappings = new HashMap<>() {{
-            put("rt-field-A", singletonMap("type", "keyword"));
-            put("rt-field-B", singletonMap("script", "some script"));
-            put("rt-field-C", singletonMap("script", "some other script"));
-        }};
-        sourceConfig =
-            new SourceConfig(generateRandomStringArray(10, 10, false, false), QueryConfigTests.randomQueryConfig(), runtimeMappings);
+        Map<String, Object> runtimeMappings = new HashMap<>() {
+            {
+                put("rt-field-A", singletonMap("type", "keyword"));
+                put("rt-field-B", singletonMap("script", "some script"));
+                put("rt-field-C", singletonMap("script", "some other script"));
+            }
+        };
+        sourceConfig = new SourceConfig(
+            generateRandomStringArray(10, 10, false, false),
+            QueryConfigTests.randomQueryConfig(),
+            runtimeMappings
+        );
         assertThat(TransformConfigLinter.getWarnings(function, sourceConfig, syncConfig), is(empty()));
 
         syncConfig = new TimeSyncConfig("rt-field-B", null);
         assertThat(
             TransformConfigLinter.getWarnings(function, sourceConfig, syncConfig),
-            contains("sync time field is a script-based runtime field, this transform might run slowly, please check your configuration."));
+            contains("sync time field is a script-based runtime field, this transform might run slowly, please check your configuration.")
+        );
     }
 
     public void testGetWarnings_Latest_WithScriptBasedRuntimeFields() {
@@ -74,36 +81,51 @@ public class TransformConfigLinterTests extends ESTestCase {
 
         SyncConfig syncConfig = new TimeSyncConfig("rt-field-C", null);
 
-        Map<String, Object> runtimeMappings = new HashMap<>() {{
-            put("rt-field-A", singletonMap("type", "keyword"));
-            put("rt-field-B", singletonMap("script", "some script"));
-            put("rt-field-C", singletonMap("script", "some other script"));
-        }};
-        sourceConfig =
-            new SourceConfig(generateRandomStringArray(10, 10, false, false), QueryConfigTests.randomQueryConfig(), runtimeMappings);
+        Map<String, Object> runtimeMappings = new HashMap<>() {
+            {
+                put("rt-field-A", singletonMap("type", "keyword"));
+                put("rt-field-B", singletonMap("script", "some script"));
+                put("rt-field-C", singletonMap("script", "some other script"));
+            }
+        };
+        sourceConfig = new SourceConfig(
+            generateRandomStringArray(10, 10, false, false),
+            QueryConfigTests.randomQueryConfig(),
+            runtimeMappings
+        );
 
         assertThat(
             TransformConfigLinter.getWarnings(function, sourceConfig, syncConfig),
             contains(
                 "all the group-by fields are script-based runtime fields, "
                     + "this transform might run slowly, please check your configuration.",
-                "sync time field is a script-based runtime field, this transform might run slowly, please check your configuration."));
+                "sync time field is a script-based runtime field, this transform might run slowly, please check your configuration."
+            )
+        );
     }
 
     public void testGetWarnings_Pivot_CouldNotFindAnyOptimization() {
-        PivotConfig pivotConfig =
-            new PivotConfig(
-                GroupConfigTests.randomGroupConfig(
-                    () -> new HistogramGroupSource(
-                        randomAlphaOfLengthBetween(1, 20), null, true, randomDoubleBetween(Math.nextUp(0), Double.MAX_VALUE, false))),
-                AggregationConfigTests.randomAggregationConfig(),
-                null);
-        Function function = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT);
+        PivotConfig pivotConfig = new PivotConfig(
+            GroupConfigTests.randomGroupConfig(
+                () -> new HistogramGroupSource(
+                    randomAlphaOfLengthBetween(1, 20),
+                    null,
+                    true,
+                    randomDoubleBetween(Math.nextUp(0), Double.MAX_VALUE, false)
+                )
+            ),
+            AggregationConfigTests.randomAggregationConfig(),
+            null
+        );
+        Function function = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT, Collections.emptySet());
         SourceConfig sourceConfig = SourceConfigTests.randomSourceConfig();
         SyncConfig syncConfig = TimeSyncConfigTests.randomTimeSyncConfig();
         assertThat(
             TransformConfigLinter.getWarnings(function, sourceConfig, syncConfig),
-            contains("could not find any optimizations for continuous execution, "
-                + "this transform might run slowly, please check your configuration."));
+            contains(
+                "could not find any optimizations for continuous execution, "
+                    + "this transform might run slowly, please check your configuration."
+            )
+        );
     }
 }

x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/GroupByOptimizerTests.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.transform.transforms.pivot;
+
+import org.elasticsearch.Version;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.core.transform.transforms.pivot.SingleGroupSource;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.DateHistogramGroupSourceTests.randomDateHistogramGroupSource;
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.DateHistogramGroupSourceTests.randomDateHistogramGroupSourceNoScript;
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.GeoTileGroupSourceTests.randomGeoTileGroupSource;
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.HistogramGroupSourceTests.randomHistogramGroupSourceNoScript;
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.TermsGroupSourceTests.randomTermsGroupSource;
+import static org.elasticsearch.xpack.core.transform.transforms.pivot.TermsGroupSourceTests.randomTermsGroupSourceNoScript;
+
+public class GroupByOptimizerTests extends ESTestCase {
+
+    public void testOneGroupBy() {
+        Map<String, SingleGroupSource> groups = Collections.singletonMap("date1", randomDateHistogramGroupSourceNoScript());
+
+        Collection<Entry<String, SingleGroupSource>> reorderedGroups = GroupByOptimizer.reorderGroups(groups, Collections.emptySet());
+        assertEquals(1, reorderedGroups.size());
+        Entry<String, SingleGroupSource> entry = reorderedGroups.iterator().next();
+
+        assertEquals("date1", entry.getKey());
+        assertEquals(groups.get("date1"), entry.getValue());
+    }
+
+    public void testOrderByType() {
+        Map<String, SingleGroupSource> groups = new LinkedHashMap<>();
+
+        groups.put("terms1", randomTermsGroupSourceNoScript());
+        groups.put("date1", randomDateHistogramGroupSourceNoScript());
+        groups.put("terms2", randomTermsGroupSourceNoScript());
+        groups.put("date2", randomDateHistogramGroupSourceNoScript());
+        groups.put("hist1", randomHistogramGroupSourceNoScript());
+        groups.put("geo1", randomGeoTileGroupSource());
+        groups.put("hist2", randomHistogramGroupSourceNoScript());
+
+        List<String> groupNames = GroupByOptimizer.reorderGroups(Collections.unmodifiableMap(groups), Collections.emptySet())
+            .stream()
+            .map(e -> e.getKey())
+            .collect(Collectors.toList());
+        assertEquals(List.of("date1", "date2", "hist1", "hist2", "terms1", "terms2", "geo1"), groupNames);
+
+        // index field preferred over runtime field
+        groupNames = GroupByOptimizer.reorderGroups(
+            Collections.unmodifiableMap(groups),
+            Collections.singleton(groups.get("terms1").getField())
+        ).stream().map(e -> e.getKey()).collect(Collectors.toList());
+        assertEquals(List.of("date1", "date2", "hist1", "hist2", "terms2", "terms1", "geo1"), groupNames);
+
+        // if both are runtime fields, order as defined in the config
+        groupNames = GroupByOptimizer.reorderGroups(
+            Collections.unmodifiableMap(groups),
+            Set.of(groups.get("terms1").getField(), groups.get("terms2").getField())
+        ).stream().map(e -> e.getKey()).collect(Collectors.toList());
+        assertEquals(List.of("date1", "date2", "hist1", "hist2", "terms1", "terms2", "geo1"), groupNames);
+    }
+
+    public void testOrderByScriptAndType() {
+        Map<String, SingleGroupSource> groups = new LinkedHashMap<>();
+
+        groups.put("terms1", randomTermsGroupSourceNoScript());
+        // create with scripts
+        groups.put("date1", randomDateHistogramGroupSource(Version.CURRENT, true));
+        groups.put("terms2", randomTermsGroupSource(Version.CURRENT, true));
+        groups.put("date2", randomDateHistogramGroupSourceNoScript());
+        groups.put("date3", randomDateHistogramGroupSourceNoScript());
+
+        List<String> groupNames = GroupByOptimizer.reorderGroups(Collections.unmodifiableMap(groups), Collections.emptySet())
+            .stream()
+            .map(e -> e.getKey())
+            .collect(Collectors.toList());
+        assertEquals(List.of("date2", "date3", "date1", "terms1", "terms2"), groupNames);
+
+        // prefer no script, runtime field, script
+        groupNames = GroupByOptimizer.reorderGroups(
+            Collections.unmodifiableMap(groups),
+            Collections.singleton(groups.get("date2").getField())
+        ).stream().map(e -> e.getKey()).collect(Collectors.toList());
+        assertEquals(List.of("date3", "date2", "date1", "terms1", "terms2"), groupNames);
+    }
+
+}

x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/PivotTests.java

@@ -105,14 +105,14 @@ public class PivotTests extends ESTestCase {
 
     public void testValidateExistingIndex() throws Exception {
         SourceConfig source = new SourceConfig("existing_source_index");
-        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT);
+        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT, Collections.emptySet());
 
         assertValidTransform(client, source, pivot);
     }
 
     public void testValidateNonExistingIndex() throws Exception {
         SourceConfig source = new SourceConfig("non_existing_source_index");
-        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT);
+        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT, Collections.emptySet());
 
         assertInvalidTransform(client, source, pivot);
     }
@@ -123,14 +123,16 @@ public class PivotTests extends ESTestCase {
         Function pivot = new Pivot(
             new PivotConfig(GroupConfigTests.randomGroupConfig(), getValidAggregationConfig(), expectedPageSize),
             new SettingsConfig(),
-            Version.CURRENT
+            Version.CURRENT,
+            Collections.emptySet()
         );
         assertThat(pivot.getInitialPageSize(), equalTo(expectedPageSize));
 
         pivot = new Pivot(
             new PivotConfig(GroupConfigTests.randomGroupConfig(), getValidAggregationConfig(), null),
             new SettingsConfig(),
-            Version.CURRENT
+            Version.CURRENT,
+            Collections.emptySet()
         );
         assertThat(pivot.getInitialPageSize(), equalTo(Transform.DEFAULT_INITIAL_MAX_PAGE_SEARCH_SIZE));
 
@@ -142,7 +144,7 @@ public class PivotTests extends ESTestCase {
         // search has failures although they might just be temporary
         SourceConfig source = new SourceConfig("existing_source_index_with_failing_shards");
 
-        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT);
+        Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), Version.CURRENT, Collections.emptySet());
 
         assertInvalidTransform(client, source, pivot);
     }
@@ -153,7 +155,12 @@ public class PivotTests extends ESTestCase {
         for (String agg : supportedAggregations) {
             AggregationConfig aggregationConfig = getAggregationConfig(agg);
 
-            Function pivot = new Pivot(getValidPivotConfig(aggregationConfig), new SettingsConfig(), Version.CURRENT);
+            Function pivot = new Pivot(
+                getValidPivotConfig(aggregationConfig),
+                new SettingsConfig(),
+                Version.CURRENT,
+                Collections.emptySet()
+            );
             assertValidTransform(client, source, pivot);
         }
     }
@@ -162,13 +169,20 @@ public class PivotTests extends ESTestCase {
         for (String agg : unsupportedAggregations) {
             AggregationConfig aggregationConfig = getAggregationConfig(agg);
 
-            Function pivot = new Pivot(getValidPivotConfig(aggregationConfig), new SettingsConfig(), Version.CURRENT);
+            Function pivot = new Pivot(
+                getValidPivotConfig(aggregationConfig),
+                new SettingsConfig(),
+                Version.CURRENT,
+                Collections.emptySet()
+            );
 
             pivot.validateConfig(ActionListener.wrap(r -> { fail("expected an exception but got a response"); }, e -> {
                 assertThat(e, is(instanceOf(ValidationException.class)));
                 assertThat(
                     "expected aggregations to be unsupported, but they were",
-                    e.getMessage(), containsString("Unsupported aggregation type [" + agg + "]"));
+                    e.getMessage(),
+                    containsString("Unsupported aggregation type [" + agg + "]")
+                );
             }));
         }
     }
@@ -186,7 +200,7 @@ public class PivotTests extends ESTestCase {
         assertThat(groupConfig.validate(null), is(nullValue()));
 
         PivotConfig pivotConfig = new PivotConfig(groupConfig, AggregationConfigTests.randomAggregationConfig(), null);
-        Function pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT);
+        Function pivot = new Pivot(pivotConfig, new SettingsConfig(), Version.CURRENT, Collections.emptySet());
         assertThat(pivot.getPerformanceCriticalFields(), contains("field-A", "field-B", "field-C"));
     }
 
@@ -311,9 +325,7 @@ public class PivotTests extends ESTestCase {
             );
         }
         if (agg.equals("global")) {
-            return parseAggregations(
-                "{\"pivot_global\": {\"global\": {}}}"
-            );
+            return parseAggregations("{\"pivot_global\": {\"global\": {}}}");
         }
 
         return parseAggregations(