Enable doc partitioning by default for time-series queries (elastic#133038)

dnhatn · javanna · commit c26332952e39 · 2025-08-18T09:37:31.000+02:00
With elastic#132774, the overhead of running queries with DOC partitioning is small. While we might switch the default data partitioning to DOC for all queries in the future, this PR defaults data partitioning to DOC for time-series queries only to minimize any unexpected impact. Relates elastic#132774
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/DataPartitioning.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/DataPartitioning.java
@@ -7,9 +7,11 @@
 
 package org.elasticsearch.compute.lucene;
 
+import org.apache.lucene.search.Query;
 import org.elasticsearch.compute.operator.Driver;
 
 import java.util.List;
+import java.util.function.Function;
 
 /**
  * How we partition the data across {@link Driver}s. Each request forks into
@@ -54,5 +56,19 @@ public enum DataPartitioning {
      *       their own tasks. See {@link LuceneSliceQueue#nextSlice(LuceneSlice)}.</li>
      * </ol>
      */
-    DOC
+    DOC;
+
+    @FunctionalInterface
+    public interface AutoStrategy {
+        Function<Query, LuceneSliceQueue.PartitioningStrategy> pickStrategy(int limit);
+
+        AutoStrategy DEFAULT = LuceneSourceOperator.Factory::autoStrategy;
+        AutoStrategy DEFAULT_TIME_SERIES = limit -> {
+            if (limit == LuceneOperator.NO_LIMIT) {
+                return q -> LuceneSliceQueue.PartitioningStrategy.DOC;
+            } else {
+                return DEFAULT.pickStrategy(limit);
+            }
+        };
+    }
 }
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
@@ -68,6 +68,7 @@ public Factory(
             List<? extends ShardContext> contexts,
             Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
             DataPartitioning dataPartitioning,
+            DataPartitioning.AutoStrategy autoStrategy,
             int taskConcurrency,
             int maxPageSize,
             int limit,
@@ -77,7 +78,7 @@ public Factory(
                 contexts,
                 queryFunction,
                 dataPartitioning,
-                autoStrategy(limit),
+                autoStrategy.pickStrategy(limit),
                 taskConcurrency,
                 limit,
                 needsScore,
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java
@@ -381,6 +381,7 @@ static LuceneOperator.Factory luceneOperatorFactory(IndexReader reader, List<Luc
             List.of(searchContext),
             ctx -> queryAndTags,
             randomFrom(DataPartitioning.values()),
+            DataPartitioning.AutoStrategy.DEFAULT,
             randomIntBetween(1, 10),
             randomPageSize(),
             limit,
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluatorTests.java
@@ -279,6 +279,7 @@ private static LuceneOperator.Factory luceneOperatorFactory(IndexReader reader,
             List.of(searchContext),
             ctx -> List.of(new LuceneSliceQueue.QueryAndTags(query, List.of())),
             randomFrom(DataPartitioning.values()),
+            DataPartitioning.AutoStrategy.DEFAULT,
             randomIntBetween(1, 10),
             randomPageSize(),
             LuceneOperator.NO_LIMIT,
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java
@@ -224,6 +224,7 @@ private LuceneSourceOperator.Factory simple(DataPartitioning dataPartitioning, i
             List.of(ctx),
             queryFunction,
             dataPartitioning,
+            DataPartitioning.AutoStrategy.DEFAULT,
             taskConcurrency,
             maxPageSize,
             limit,
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/read/ValueSourceReaderTypeConversionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/read/ValueSourceReaderTypeConversionTests.java
@@ -277,6 +277,7 @@ private SourceOperator simpleInput(DriverContext context, int size, int commitEv
             shardContexts,
             ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
             DataPartitioning.SHARD,
+            DataPartitioning.AutoStrategy.DEFAULT,
             1,// randomIntBetween(1, 10),
             pageSize,
             LuceneOperator.NO_LIMIT,
@@ -1312,6 +1313,7 @@ public void testWithNulls() throws IOException {
             List.of(shardContext),
             ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
             randomFrom(DataPartitioning.values()),
+            DataPartitioning.AutoStrategy.DEFAULT,
             randomIntBetween(1, 10),
             randomPageSize(),
             LuceneOperator.NO_LIMIT,
@@ -1473,6 +1475,7 @@ public void testManyShards() throws IOException {
                 contexts,
                 ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
                 DataPartitioning.SHARD,
+                DataPartitioning.AutoStrategy.DEFAULT,
                 randomIntBetween(1, 10),
                 1000,
                 LuceneOperator.NO_LIMIT,
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/read/ValuesSourceReaderOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/read/ValuesSourceReaderOperatorTests.java
@@ -195,6 +195,7 @@ private SourceOperator sourceOperator(DriverContext context, int pageSize) {
             List.of(new LuceneSourceOperatorTests.MockShardContext(reader, 0)),
             ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
             DataPartitioning.SHARD,
+            DataPartitioning.AutoStrategy.DEFAULT,
             randomIntBetween(1, 10),
             pageSize,
             LuceneOperator.NO_LIMIT,
@@ -1506,6 +1507,7 @@ public void testWithNulls() throws IOException {
             List.of(new LuceneSourceOperatorTests.MockShardContext(reader, 0)),
             ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
             randomFrom(DataPartitioning.values()),
+            DataPartitioning.AutoStrategy.DEFAULT,
             randomIntBetween(1, 10),
             randomPageSize(),
             LuceneOperator.NO_LIMIT,
@@ -1755,6 +1757,7 @@ public void testManyShards() throws IOException {
                 contexts,
                 ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
                 DataPartitioning.SHARD,
+                DataPartitioning.AutoStrategy.DEFAULT,
                 randomIntBetween(1, 10),
                 1000,
                 LuceneOperator.NO_LIMIT,
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupFromIndexIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupFromIndexIT.java
@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.action;
 
-import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.action.ActionListener;
@@ -25,6 +24,7 @@
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.data.LongVector;
 import org.elasticsearch.compute.lucene.DataPartitioning;
+import org.elasticsearch.compute.lucene.LuceneOperator;
 import org.elasticsearch.compute.lucene.LuceneSliceQueue;
 import org.elasticsearch.compute.lucene.LuceneSourceOperator;
 import org.elasticsearch.compute.lucene.ShardContext;
@@ -280,9 +280,10 @@ private void runLookup(List<DataType> keyTypes, PopulateIndices populateIndices)
                 List.of(esqlContext),
                 ctx -> List.of(new LuceneSliceQueue.QueryAndTags(new MatchAllDocsQuery(), List.of())),
                 DataPartitioning.SEGMENT,
+                DataPartitioning.AutoStrategy.DEFAULT,
                 1,
                 10000,
-                DocIdSetIterator.NO_MORE_DOCS,
+                LuceneOperator.NO_LIMIT,
                 false // no scoring
             );
             List<ValuesSourceReaderOperator.FieldInfo> fieldInfos = new ArrayList<>();
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java
@@ -292,6 +292,7 @@ public final PhysicalOperation sourcePhysicalOperation(EsQueryExec esQueryExec,
                 shardContexts,
                 querySupplier(esQueryExec.query()),
                 context.queryPragmas().dataPartitioning(physicalSettings.defaultDataPartitioning()),
+                context.autoPartitioningStrategy().get(),
                 context.queryPragmas().taskConcurrency(),
                 context.pageSize(rowEstimatedSize),
                 limit,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.LocalCircuitBreaker;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.lucene.DataPartitioning;
 import org.elasticsearch.compute.lucene.LuceneOperator;
 import org.elasticsearch.compute.operator.ChangePointOperator;
 import org.elasticsearch.compute.operator.ColumnExtractOperator;
@@ -117,6 +118,7 @@
 import org.elasticsearch.xpack.esql.plan.physical.RrfScoreEvalExec;
 import org.elasticsearch.xpack.esql.plan.physical.SampleExec;
 import org.elasticsearch.xpack.esql.plan.physical.ShowExec;
+import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesAggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
 import org.elasticsearch.xpack.esql.plan.physical.inference.CompletionExec;
@@ -201,6 +203,7 @@ public LocalExecutionPlanner(
      * turn the given plan into a list of drivers to execute
      */
     public LocalExecutionPlan plan(String description, FoldContext foldCtx, PhysicalPlan localPhysicalPlan) {
+
         var context = new LocalExecutionPlannerContext(
             description,
             new ArrayList<>(),
@@ -210,6 +213,11 @@ public LocalExecutionPlan plan(String description, FoldContext foldCtx, Physical
             blockFactory,
             foldCtx,
             settings,
+            new Holder<>(
+                localPhysicalPlan.anyMatch(p -> p instanceof TimeSeriesAggregateExec)
+                    ? DataPartitioning.AutoStrategy.DEFAULT_TIME_SERIES
+                    : DataPartitioning.AutoStrategy.DEFAULT
+            ),
             shardContexts
         );
 
@@ -1012,6 +1020,7 @@ public record LocalExecutionPlannerContext(
         BlockFactory blockFactory,
         FoldContext foldCtx,
         Settings settings,
+        Holder<DataPartitioning.AutoStrategy> autoPartitioningStrategy,
         List<EsPhysicalOperationProviders.ShardContext> shardContexts
     ) {
         void addDriverFactory(DriverFactory driverFactory) {
