Increase parallelism for TS command

Author: dnhatn

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeBuffer.java

@@ -17,7 +17,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 
-final class ExchangeBuffer {
+public final class ExchangeBuffer {
 
     private final Queue<Page> queue = new ConcurrentLinkedQueue<>();
     // uses a separate counter for size for CAS; and ConcurrentLinkedQueue#size is not a constant time operation.
@@ -34,7 +34,7 @@ final class ExchangeBuffer {
 
     private volatile boolean noMoreInputs = false;
 
-    ExchangeBuffer(int maxSize) {
+    public ExchangeBuffer(int maxSize) {
         if (maxSize < 1) {
             throw new IllegalArgumentException("max_buffer_size must be at least one; got=" + maxSize);
         }
@@ -161,4 +161,73 @@ int size() {
     void addCompletionListener(ActionListener<Void> listener) {
         completionFuture.addListener(listener);
     }
+
+    /**
+     * Creates a new {@link ExchangeSource} associated with this buffer.
+     */
+    public ExchangeSource newExchangeSource() {
+        final ExchangeBuffer buffer = this;
+        return new ExchangeSource() {
+            @Override
+            public Page pollPage() {
+                return buffer.pollPage();
+            }
+
+            @Override
+            public void finish() {
+                buffer.finish(true);
+            }
+
+            @Override
+            public boolean isFinished() {
+                return buffer.isFinished();
+            }
+
+            @Override
+            public int bufferSize() {
+                return buffer.size();
+            }
+
+            @Override
+            public IsBlockedResult waitForReading() {
+                return buffer.waitForReading();
+            }
+        };
+    }
+
+    /**
+     * Creates a new {@link ExchangeSink} associated with this buffer.
+     */
+    public ExchangeSink newExchangeSink() {
+        final ExchangeBuffer buffer = this;
+        return new ExchangeSink() {
+            private boolean finished = false;
+
+            @Override
+            public void addPage(Page page) {
+                buffer.addPage(page);
+            }
+
+            @Override
+            public void finish() {
+                finished = true;
+                buffer.finish(false);
+            }
+
+            @Override
+            public boolean isFinished() {
+                return finished; // no need to wait for the buffer
+            }
+
+            @Override
+            public void addCompletionListener(ActionListener<Void> listener) {
+                buffer.addCompletionListener(listener);
+            }
+
+            @Override
+            public IsBlockedResult waitForWriting() {
+                return buffer.waitForWriting();
+            }
+        };
+    }
 }

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/TimeSeriesIT.java

@@ -11,6 +11,8 @@
 import org.elasticsearch.common.Randomness;
 import org.elasticsearch.common.Rounding;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.compute.lucene.TimeSeriesSourceOperator;
+import org.elasticsearch.compute.operator.DriverProfile;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
 import org.elasticsearch.xpack.esql.core.type.DataType;
@@ -729,4 +731,29 @@ public void testIndexMode() {
         });
         assertThat(failure.getMessage(), containsString("Unknown index [hosts-old]"));
     }
+
+    public void testProfile() {
+        EsqlQueryRequest request = new EsqlQueryRequest();
+        request.profile(true);
+        request.query("TS hosts | STATS sum(rate(request_count)) BY cluster, bucket(@timestamp, 1minute) | SORT cluster");
+        try (var resp = run(request)) {
+            EsqlQueryResponse.Profile profile = resp.profile();
+            List<DriverProfile> dataProfiles = profile.drivers().stream().filter(d -> d.description().equals("data")).toList();
+            int totalTimeSeries = 0;
+            for (DriverProfile p : dataProfiles) {
+                if (p.operators().stream().anyMatch(s -> s.status() instanceof TimeSeriesSourceOperator.Status)) {
+                    totalTimeSeries++;
+                    assertThat(p.operators(), hasSize(2));
+                    assertThat(p.operators().get(1).operator(), equalTo("ExchangeSinkOperator"));
+                } else {
+                    assertThat(p.operators(), hasSize(4));
+                    assertThat(p.operators().get(0).operator(), equalTo("ExchangeSourceOperator"));
+                    assertThat(p.operators().get(1).operator(), containsString("EvalOperator[evaluator=DateTruncDatetimeEvaluator"));
+                    assertThat(p.operators().get(2).operator(), containsString("TimeSeriesAggregationOperator"));
+                    assertThat(p.operators().get(3).operator(), equalTo("ExchangeSinkOperator"));
+                }
+            }
+            assertThat(totalTimeSeries, equalTo(dataProfiles.size() / 2));
+        }
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -46,6 +46,7 @@
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator.SourceOperatorFactory;
 import org.elasticsearch.compute.operator.StringExtractOperator;
+import org.elasticsearch.compute.operator.exchange.ExchangeBuffer;
 import org.elasticsearch.compute.operator.exchange.ExchangeSink;
 import org.elasticsearch.compute.operator.exchange.ExchangeSinkOperator.ExchangeSinkOperatorFactory;
 import org.elasticsearch.compute.operator.exchange.ExchangeSource;
@@ -195,6 +196,7 @@ public LocalExecutionPlanner(
      */
     public LocalExecutionPlan plan(String description, FoldContext foldCtx, PhysicalPlan localPhysicalPlan) {
         var context = new LocalExecutionPlannerContext(
+            description,
             new ArrayList<>(),
             new Holder<>(DriverParallelism.SINGLE),
             configuration.pragmas(),
@@ -272,9 +274,9 @@ else if (node instanceof EsQueryExec esQuery) {
         } else if (node instanceof ShowExec show) {
             return planShow(show);
         } else if (node instanceof ExchangeSourceExec exchangeSource) {
-            return planExchangeSource(exchangeSource, context);
+            return planExchangeSource(exchangeSource, exchangeSourceSupplier, context);
         } else if (node instanceof TimeSeriesSourceExec ts) {
-            return planTimeSeriesNode(ts, context);
+            return planTimeSeriesSource(ts, context);
         }
         // lookups and joins
         else if (node instanceof EnrichExec enrich) {
@@ -332,8 +334,36 @@ private PhysicalOperation planEsQueryNode(EsQueryExec esQueryExec, LocalExecutio
         return physicalOperationProviders.sourcePhysicalOperation(esQueryExec, context);
     }
 
-    private PhysicalOperation planTimeSeriesNode(TimeSeriesSourceExec esQueryExec, LocalExecutionPlannerContext context) {
-        return physicalOperationProviders.timeSeriesSourceOperation(esQueryExec, context);
+    private PhysicalOperation planTimeSeriesSource(TimeSeriesSourceExec ts, LocalExecutionPlannerContext context) {
+        var exchangeBuffer = new ExchangeBuffer(configuration.pragmas().exchangeBufferSize());
+        // Use a separate driver for the time-series source to split the pipeline to increase parallelism,
+        // since the time-series source must be executed with a single driver at the shard level.
+        {
+            var sourceOperator = physicalOperationProviders.timeSeriesSourceOperation(ts, context);
+            var sinkOperator = sourceOperator.withSink(
+                new ExchangeSinkOperatorFactory(exchangeBuffer::newExchangeSink),
+                sourceOperator.layout
+            );
+            final TimeValue statusInterval = configuration.pragmas().statusInterval();
+            context.addDriverFactory(
+                new DriverFactory(
+                    new DriverSupplier(
+                        context.description,
+                        ClusterName.CLUSTER_NAME_SETTING.get(settings).value(),
+                        Node.NODE_NAME_SETTING.get(settings),
+                        context.bigArrays,
+                        context.blockFactory,
+                        sinkOperator,
+                        statusInterval,
+                        settings
+                    ),
+                    DriverParallelism.SINGLE
+                )
+            );
+            context.driverParallelism.set(DriverParallelism.SINGLE);
+        }
+        var exchangeSource = new ExchangeSourceExec(ts.source(), ts.output(), false);
+        return planExchangeSource(exchangeSource, exchangeBuffer::newExchangeSource, context);
     }
 
     private PhysicalOperation planEsStats(EsStatsQueryExec statsQuery, LocalExecutionPlannerContext context) {
@@ -410,7 +440,11 @@ private PhysicalOperation planExchangeSink(ExchangeSinkExec exchangeSink, LocalE
         return source.withSink(new ExchangeSinkOperatorFactory(exchangeSinkSupplier), source.layout);
     }
 
-    private PhysicalOperation planExchangeSource(ExchangeSourceExec exchangeSource, LocalExecutionPlannerContext context) {
+    private PhysicalOperation planExchangeSource(
+        ExchangeSourceExec exchangeSource,
+        Supplier<ExchangeSource> exchangeSourceSupplier,
+        LocalExecutionPlannerContext context
+    ) {
         Objects.requireNonNull(exchangeSourceSupplier, "ExchangeSourceHandler wasn't provided");
 
         var builder = new Layout.Builder();
@@ -923,6 +957,7 @@ enum Type {
      * maintains information how many driver instances should be created for a given driver.
      */
     public record LocalExecutionPlannerContext(
+        String description,
         List<DriverFactory> driverFactories,
         Holder<DriverParallelism> driverParallelism,
         QueryPragmas queryPragmas,

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlannerTests.java

@@ -30,7 +30,9 @@
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy;
+import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperServiceTestCase;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.node.Node;
 import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.Plugin;
@@ -39,13 +41,16 @@
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
 import org.elasticsearch.xpack.esql.expression.Order;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
+import org.elasticsearch.xpack.esql.plan.physical.LimitExec;
+import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesSourceExec;
 import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.esql.session.Configuration;
@@ -59,8 +64,10 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class LocalExecutionPlannerTests extends MapperServiceTestCase {
@@ -204,6 +211,24 @@ public void testDriverClusterAndNodeName() throws IOException {
         assertThat(supplier.nodeName(), equalTo("node-1"));
     }
 
+    public void testTimeSeriesSource() throws Exception {
+        MetadataAttribute tsid = new MetadataAttribute(Source.EMPTY, "_tsid", DataType.KEYWORD, false);
+        var timeSeriesSource = new TimeSeriesSourceExec(
+            Source.EMPTY,
+            List.of(tsid),
+            new MatchAllQueryBuilder(),
+            new Literal(Source.EMPTY, between(1, 100), DataType.INTEGER),
+            MappedFieldType.FieldExtractPreference.NONE,
+            Set.of(),
+            Set.of(),
+            List.of(),
+            randomEstimatedRowSize(estimatedRowSizeIsHuge)
+        );
+        var limitExec = new LimitExec(Source.EMPTY, timeSeriesSource, new Literal(Source.EMPTY, between(1, 100), DataType.INTEGER));
+        LocalExecutionPlanner.LocalExecutionPlan plan = planner().plan("test", FoldContext.small(), limitExec);
+        assertThat(plan.driverFactories, hasSize(2));
+    }
+
     private int randomEstimatedRowSize(boolean huge) {
         int hugeBoundary = SourceOperator.MIN_TARGET_PAGE_SIZE * 10;
         return huge ? between(hugeBoundary, Integer.MAX_VALUE) : between(1, hugeBoundary);