ESQL: Heuristics to pick efficient partitioning (#125739)

Adds heuristics to pick an efficient partitioning strategy based on the
index and rewritten query. This speeds up some queries by throwing more
cores at the problem:
```
FROM test | STATS SUM(b)

Before: took: 31 CPU: 222.3%
 After: took: 15 CPU: 806.9%
```

It also lowers the overhead of simpler queries by throwing less cores at
the problem when it won't really speed anything up:
```
FROM test

Before: took: 1 CPU: 48.5%
 After: took: 1 CPU: 70.4%
```

We have had a `pragma` to control our data partitioning for a long time,
this just looks at the query to pick a partitioning scheme. The
partitioning options:
* `shard`: use one core per shard
* `segment`: use one core per large segment
* `doc`: break each shard into as many segments as there are cores

`doc` is the fastest, but has a lot of overhead, especially for complex
Lucene queries. `segment` is fast, but doesn't make the most out of CPUs
when there are few segments. `shard` has the lowest overhead.

Previously we always used `segment` partitioning because it doesn't have
the terrible overhead but is fast. With this change we use `doc` when
the top level query matches all documents - those have very very low
overhead even in the `doc` partitioning. That's the twice as fast
example above.

This also uses the `shard` partitioning for queries that don't have to
do much work like `FROM foo` or `FROM foo | LIMIT 1` or
`FROM foo | SORT a`. That's the lower CPU example above.

This forking choice is taken very late on the data node. So queries like
this:
```
FROM test | WHERE @timestamp > "2025-01-01T00:00:00Z" | STATS SUM(b)
```
can also use the `doc` partitioning when all documents are after the
timestamp and all documents have `b`.

Author: nik9000

docs/changelog/125739.yaml

@@ -0,0 +1,5 @@
+pr: 125739
+summary: Heuristics to pick efficient partitioning
+area: ES|QL
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -218,6 +218,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SEMANTIC_TEXT_CHUNKING_CONFIG = def(9_047_00_0);
     public static final TransportVersion REPO_ANALYSIS_COPY_BLOB = def(9_048_00_0);
     public static final TransportVersion AMAZON_BEDROCK_TASK_SETTINGS = def(9_049_00_0);
+    public static final TransportVersion ESQL_REPORT_SHARD_PARTITIONING = def(9_050_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/DataPartitioning.java

@@ -7,11 +7,39 @@
 
 package org.elasticsearch.compute.lucene;
 
-public enum DataPartitioning {
+import org.elasticsearch.compute.operator.Driver;
 
+/**
+ * How we partition the data across {@link Driver}s. Each request forks into
+ * {@code min(1.5 * cpus, partition_count)} threads on the data node. More partitions
+ * allow us to bring more threads to bear on CPU intensive data node side tasks.
+ */
+public enum DataPartitioning {
+    /**
+     * Automatically select the data partitioning based on the query and index.
+     * Usually that's {@link #SEGMENT}, but for small indices it's {@link #SHARD}.
+     * When the additional overhead from {@link #DOC} is fairly low then it'll
+     * pick {@link #DOC}.
+     */
+    AUTO,
+    /**
+     * Make one partition per shard. This is generally the slowest option, but it
+     * has the lowest CPU overhead.
+     */
     SHARD,
-
+    /**
+     * Partition on segment boundaries, this doesn't allow forking to as many CPUs
+     * as {@link #DOC} but it has much lower overhead.
+     * <p>
+     * It packs segments smaller than {@link LuceneSliceQueue#MAX_DOCS_PER_SLICE}
+     * docs together into a partition. Larger segments get their own partition.
+     * Each slice contains no more than {@link LuceneSliceQueue#MAX_SEGMENTS_PER_SLICE}.
+     */
     SEGMENT,
-
+    /**
+     * Partition each shard into {@code task_concurrency} partitions, splitting
+     * larger segments into slices. This allows bringing the most CPUs to bear on
+     * the problem but adds extra overhead, especially in query preparation.
+     */
     DOC,
 }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneCountOperator.java

@@ -49,7 +49,16 @@ public Factory(
             int taskConcurrency,
             int limit
         ) {
-            super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false);
+            super(
+                contexts,
+                queryFunction,
+                dataPartitioning,
+                query -> LuceneSliceQueue.PartitioningStrategy.SHARD,
+                taskConcurrency,
+                limit,
+                false,
+                ScoreMode.COMPLETE_NO_SCORES
+            );
         }
 
         @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMaxFactory.java

@@ -23,8 +23,6 @@
 import java.util.List;
 import java.util.function.Function;
 
-import static org.elasticsearch.compute.lucene.LuceneOperator.weightFunction;
-
 /**
  * Factory that generates an operator that finds the max value of a field using the {@link LuceneMinMaxOperator}.
  */
@@ -123,7 +121,16 @@ public LuceneMaxFactory(
         NumberType numberType,
         int limit
     ) {
-        super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false);
+        super(
+            contexts,
+            queryFunction,
+            dataPartitioning,
+            query -> LuceneSliceQueue.PartitioningStrategy.SHARD,
+            taskConcurrency,
+            limit,
+            false,
+            ScoreMode.COMPLETE_NO_SCORES
+        );
         this.fieldName = fieldName;
         this.numberType = numberType;
     }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneMinFactory.java

@@ -23,8 +23,6 @@
 import java.util.List;
 import java.util.function.Function;
 
-import static org.elasticsearch.compute.lucene.LuceneOperator.weightFunction;
-
 /**
  * Factory that generates an operator that finds the min value of a field using the {@link LuceneMinMaxOperator}.
  */
@@ -123,7 +121,16 @@ public LuceneMinFactory(
         NumberType numberType,
         int limit
     ) {
-        super(contexts, weightFunction(queryFunction, ScoreMode.COMPLETE_NO_SCORES), dataPartitioning, taskConcurrency, limit, false);
+        super(
+            contexts,
+            queryFunction,
+            dataPartitioning,
+            query -> LuceneSliceQueue.PartitioningStrategy.SHARD,
+            taskConcurrency,
+            limit,
+            false,
+            ScoreMode.COMPLETE_NO_SCORES
+        );
         this.fieldName = fieldName;
         this.numberType = numberType;
     }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java

@@ -9,7 +9,6 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.BulkScorer;
-import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Query;
@@ -37,12 +36,16 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.TransportVersions.ESQL_REPORT_SHARD_PARTITIONING;
+
 public abstract class LuceneOperator extends SourceOperator {
     private static final Logger logger = LogManager.getLogger(LuceneOperator.class);
 
@@ -93,15 +96,17 @@ public abstract static class Factory implements SourceOperator.SourceOperatorFac
          */
         protected Factory(
             List<? extends ShardContext> contexts,
-            Function<ShardContext, Weight> weightFunction,
+            Function<ShardContext, Query> queryFunction,
             DataPartitioning dataPartitioning,
+            Function<Query, LuceneSliceQueue.PartitioningStrategy> autoStrategy,
             int taskConcurrency,
             int limit,
-            boolean needsScore
+            boolean needsScore,
+            ScoreMode scoreMode
         ) {
             this.limit = limit;
             this.dataPartitioning = dataPartitioning;
-            this.sliceQueue = LuceneSliceQueue.create(contexts, weightFunction, dataPartitioning, taskConcurrency);
+            this.sliceQueue = LuceneSliceQueue.create(contexts, queryFunction, dataPartitioning, autoStrategy, taskConcurrency, scoreMode);
             this.taskConcurrency = Math.min(sliceQueue.totalSlices(), taskConcurrency);
             this.needsScore = needsScore;
         }
@@ -269,6 +274,7 @@ public static class Status implements Operator.Status {
         private final int sliceMax;
         private final int current;
         private final long rowsEmitted;
+        private final Map<String, LuceneSliceQueue.PartitioningStrategy> partitioningStrategies;
 
         private Status(LuceneOperator operator) {
             processedSlices = operator.processedSlices;
@@ -294,6 +300,7 @@ private Status(LuceneOperator operator) {
             }
             pagesEmitted = operator.pagesEmitted;
             rowsEmitted = operator.rowsEmitted;
+            partitioningStrategies = operator.sliceQueue.partitioningStrategies();
         }
 
         Status(
@@ -307,7 +314,8 @@ private Status(LuceneOperator operator) {
             int sliceMin,
             int sliceMax,
             int current,
-            long rowsEmitted
+            long rowsEmitted,
+            Map<String, LuceneSliceQueue.PartitioningStrategy> partitioningStrategies
         ) {
             this.processedSlices = processedSlices;
             this.processedQueries = processedQueries;
@@ -320,6 +328,7 @@ private Status(LuceneOperator operator) {
             this.sliceMax = sliceMax;
             this.current = current;
             this.rowsEmitted = rowsEmitted;
+            this.partitioningStrategies = partitioningStrategies;
         }
 
         Status(StreamInput in) throws IOException {
@@ -343,6 +352,9 @@ private Status(LuceneOperator operator) {
             } else {
                 rowsEmitted = 0;
             }
+            partitioningStrategies = in.getTransportVersion().onOrAfter(ESQL_REPORT_SHARD_PARTITIONING)
+                ? in.readMap(LuceneSliceQueue.PartitioningStrategy::readFrom)
+                : Map.of();
         }
 
         @Override
@@ -364,6 +376,9 @@ public void writeTo(StreamOutput out) throws IOException {
             if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_PROFILE_ROWS_PROCESSED)) {
                 out.writeVLong(rowsEmitted);
             }
+            if (out.getTransportVersion().onOrAfter(ESQL_REPORT_SHARD_PARTITIONING)) {
+                out.writeMap(partitioningStrategies, StreamOutput::writeString, StreamOutput::writeWriteable);
+            }
         }
 
         @Override
@@ -415,6 +430,10 @@ public long rowsEmitted() {
             return rowsEmitted;
         }
 
+        public Map<String, LuceneSliceQueue.PartitioningStrategy> partitioningStrategies() {
+            return partitioningStrategies;
+        }
+
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
@@ -432,6 +451,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field("slice_max", sliceMax);
             builder.field("current", current);
             builder.field("rows_emitted", rowsEmitted);
+            builder.field("partitioning_strategies", new TreeMap<>(this.partitioningStrategies));
             return builder.endObject();
         }
 
@@ -450,12 +470,23 @@ public boolean equals(Object o) {
                 && sliceMin == status.sliceMin
                 && sliceMax == status.sliceMax
                 && current == status.current
-                && rowsEmitted == status.rowsEmitted;
+                && rowsEmitted == status.rowsEmitted
+                && partitioningStrategies.equals(status.partitioningStrategies);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(processedSlices, sliceIndex, totalSlices, pagesEmitted, sliceMin, sliceMax, current, rowsEmitted);
+            return Objects.hash(
+                processedSlices,
+                sliceIndex,
+                totalSlices,
+                pagesEmitted,
+                sliceMin,
+                sliceMax,
+                current,
+                rowsEmitted,
+                partitioningStrategies
+            );
         }
 
         @Override
@@ -468,17 +499,4 @@ public TransportVersion getMinimalSupportedVersion() {
             return TransportVersions.V_8_11_X;
         }
     }
-
-    static Function<ShardContext, Weight> weightFunction(Function<ShardContext, Query> queryFunction, ScoreMode scoreMode) {
-        return ctx -> {
-            final var query = queryFunction.apply(ctx);
-            final var searcher = ctx.searcher();
-            try {
-                Query actualQuery = scoreMode.needsScores() ? query : new ConstantScoreQuery(query);
-                return searcher.createWeight(searcher.rewrite(actualQuery), scoreMode, 1);
-            } catch (IOException e) {
-                throw new UncheckedIOException(e);
-            }
-        };
-    }
 }