elastic · astefan · Apr 1, 2024 · Mar 8, 2024 · Mar 19, 2024 · Mar 19, 2024
diff --git a/docs/changelog/106516.yaml b/docs/changelog/106516.yaml
@@ -0,0 +1,5 @@
+pr: 106516
+summary: "ESQL: perform a reduction on the data node"
+area: ES|QL
+type: enhancement
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -157,6 +157,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SEARCH_NODE_LOAD_AUTOSCALING = def(8_617_00_0);
     public static final TransportVersion ESQL_ES_SOURCE_OPTIONS = def(8_618_00_0);
     public static final TransportVersion ADD_PERSISTENT_TASK_EXCEPTIONS = def(8_619_00_0);
+    public static final TransportVersion ESQL_REDUCER_NODE_FRAGMENT = def(8_620_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

diff --git a/...n/esql/qa/security/src/javaRestTest/java/org/elasticsearch/xpack/esql/EsqlSecurityIT.java b/...n/esql/qa/security/src/javaRestTest/java/org/elasticsearch/xpack/esql/EsqlSecurityIT.java
@@ -389,6 +389,9 @@ static Settings randomPragmas() {
         if (randomBoolean()) {
             settings.put("enrich_max_workers", between(1, 5));
         }
+        if (randomBoolean()) {
+            settings.put("node_level_reduction", randomBoolean());
+        }
         return settings.build();
     }
 }
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec
@@ -62,9 +62,9 @@ eth2           |epsilon        |[fe81::cae2:65ff:fece:feb9, fe82::cae2:65ff:fece
 lessThan
 required_feature: esql.mv_warn
 
-from hosts | sort host, card | where ip0 < ip1 | keep card, host, ip0, ip1;
-warning:Line 1:38: evaluation of [ip0 < ip1] failed, treating result as null. Only first 20 failures recorded.
-warning:Line 1:38: java.lang.IllegalArgumentException: single-value function encountered multi-value
+from hosts | sort host, card, ip1 | where ip0 < ip1 | keep card, host, ip0, ip1;
+warning:Line 1:43: evaluation of [ip0 < ip1] failed, treating result as null. Only first 20 failures recorded.
+warning:Line 1:43: java.lang.IllegalArgumentException: single-value function encountered multi-value
 
 card:keyword   |host:keyword   |ip0:ip                   |ip1:ip
 eth1           |beta           |127.0.0.1                |127.0.0.2

diff --git a/...ternalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java b/...ternalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java
@@ -188,6 +188,9 @@ protected static QueryPragmas randomPragmas() {
             if (randomBoolean()) {
                 settings.put("max_concurrent_shards_per_node", randomIntBetween(1, 10));
             }
+            if (randomBoolean()) {
+                settings.put("node_level_reduction", randomBoolean());
+            }
         }
         return new QueryPragmas(settings.build());
     }

diff --git a/...ql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java b/...ql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionTaskIT.java
@@ -76,6 +76,7 @@ public class EsqlActionTaskIT extends AbstractPausableIntegTestCase {
     private String READ_DESCRIPTION;
     private String MERGE_DESCRIPTION;
     private String REDUCE_DESCRIPTION;
+    private boolean nodeLevelReduction;
 
     @Before
     public void setup() {
@@ -94,6 +95,7 @@ public void setup() {
         REDUCE_DESCRIPTION = """
             \\_ExchangeSourceOperator[]
             \\_ExchangeSinkOperator""";
+        nodeLevelReduction = randomBoolean();
     }
 
     public void testTaskContents() throws Exception {
@@ -209,22 +211,31 @@ public void testCancelEsqlTask() throws Exception {
     }
 
     private ActionFuture<EsqlQueryResponse> startEsql() {
+        return startEsql("from test | stats sum(pause_me)");
+    }
+
+    private ActionFuture<EsqlQueryResponse> startEsql(String query) {
         scriptPermits.drainPermits();
         scriptPermits.release(between(1, 5));
-        var pragmas = new QueryPragmas(
-            Settings.builder()
-                // Force shard partitioning because that's all the tests know how to match. It is easier to reason about too.
-                .put("data_partitioning", "shard")
-                // Limit the page size to something small so we do more than one page worth of work, so we get more status updates.
-                .put("page_size", pageSize())
-                // Report the status after every action
-                .put("status_interval", "0ms")
-                .build()
-        );
-        return EsqlQueryRequestBuilder.newSyncEsqlQueryRequestBuilder(client())
-            .query("from test | stats sum(pause_me)")
-            .pragmas(pragmas)
-            .execute();
+        var settingsBuilder = Settings.builder()
+            // Force shard partitioning because that's all the tests know how to match. It is easier to reason about too.
+            .put("data_partitioning", "shard")
+            // Limit the page size to something small so we do more than one page worth of work, so we get more status updates.
+            .put("page_size", pageSize())
+            // Report the status after every action
+            .put("status_interval", "0ms");
+
+        if (nodeLevelReduction == false) {
+            // explicitly set the default (false) or don't
+            if (randomBoolean()) {
+                settingsBuilder.put("node_level_reduction", nodeLevelReduction);
+            }
+        } else {
+            settingsBuilder.put("node_level_reduction", nodeLevelReduction);
+        }
+
+        var pragmas = new QueryPragmas(settingsBuilder.build());
+        return EsqlQueryRequestBuilder.newSyncEsqlQueryRequestBuilder(client()).query(query).pragmas(pragmas).execute();
     }
 
     private void cancelTask(TaskId taskId) {
@@ -407,6 +418,67 @@ protected void doRun() throws Exception {
         }
     }
 
+    public void testTaskContentsForTopNQuery() throws Exception {
+        READ_DESCRIPTION = ("\\_LuceneTopNSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = 1000, "
+            + "sorts = [{\"pause_me\":{\"order\":\"asc\",\"missing\":\"_last\",\"unmapped_type\":\"long\"}}]]\n"
+            + "\\_ValuesSourceReaderOperator[fields = [pause_me]]\n"
+            + "\\_ProjectOperator[projection = [1]]\n"
+            + "\\_ExchangeSinkOperator").replace("pageSize()", Integer.toString(pageSize()));
+        MERGE_DESCRIPTION = "\\_ExchangeSourceOperator[]\n"
+            + "\\_TopNOperator[count=1000, elementTypes=[LONG], encoders=[DefaultSortable], "
+            + "sortOrders=[SortOrder[channel=0, asc=true, nullsFirst=false]]]\n"
+            + "\\_ProjectOperator[projection = [0]]\n"
+            + "\\_OutputOperator[columns = [pause_me]]";
+        REDUCE_DESCRIPTION = "\\_ExchangeSourceOperator[]\n"
+            + (nodeLevelReduction
+                ? "\\_TopNOperator[count=1000, elementTypes=[LONG], encoders=[DefaultSortable], "
+                    + "sortOrders=[SortOrder[channel=0, asc=true, nullsFirst=false]]]\n"
+                : "")
+            + "\\_ExchangeSinkOperator";
+
+        ActionFuture<EsqlQueryResponse> response = startEsql("from test | sort pause_me | keep pause_me");
+        try {
+            getTasksStarting();
+            scriptPermits.release(pageSize());
+            getTasksRunning();
+        } finally {
+            // each scripted field "emit" is called by LuceneTopNSourceOperator and by ValuesSourceReaderOperator
+            scriptPermits.release(2 * numberOfDocs());
+            try (EsqlQueryResponse esqlResponse = response.get()) {
+                assertThat(Iterators.flatMap(esqlResponse.values(), i -> i).next(), equalTo(1L));
+            }
+        }
+    }
+
+    public void testTaskContentsForLimitQuery() throws Exception {
+        String limit = Integer.toString(randomIntBetween(pageSize() + 1, 2 * numberOfDocs()));
+        READ_DESCRIPTION = """
+            \\_LuceneSourceOperator[dataPartitioning = SHARD, maxPageSize = pageSize(), limit = limit()]
+            \\_ValuesSourceReaderOperator[fields = [pause_me]]
+            \\_ProjectOperator[projection = [1]]
+            \\_ExchangeSinkOperator""".replace("pageSize()", Integer.toString(pageSize())).replace("limit()", limit);
+        MERGE_DESCRIPTION = """
+            \\_ExchangeSourceOperator[]
+            \\_LimitOperator[limit = limit()]
+            \\_ProjectOperator[projection = [0]]
+            \\_OutputOperator[columns = [pause_me]]""".replace("limit()", limit);
+        REDUCE_DESCRIPTION = ("\\_ExchangeSourceOperator[]\n"
+            + (nodeLevelReduction ? "\\_LimitOperator[limit = limit()]\n" : "")
+            + "\\_ExchangeSinkOperator").replace("limit()", limit);
+
+        ActionFuture<EsqlQueryResponse> response = startEsql("from test | keep pause_me | limit " + limit);
+        try {
+            getTasksStarting();
+            scriptPermits.release(pageSize());
+            getTasksRunning();
+        } finally {
+            scriptPermits.release(numberOfDocs());
+            try (EsqlQueryResponse esqlResponse = response.get()) {
+                assertThat(Iterators.flatMap(esqlResponse.values(), i -> i).next(), equalTo(1L));
+            }
+        }
+    }
+
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return CollectionUtils.appendToCopy(super.nodePlugins(), MockTransportService.TestPlugin.class);

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java
@@ -629,7 +629,8 @@ static FragmentExec readFragmentExec(PlanStreamInput in) throws IOException {
             in.readSource(),
             in.readLogicalPlanNode(),
             in.readOptionalNamedWriteable(QueryBuilder.class),
-            in.readOptionalVInt()
+            in.readOptionalVInt(),
+            in.getTransportVersion().onOrAfter(TransportVersions.ESQL_REDUCER_NODE_FRAGMENT) ? in.readOptionalPhysicalPlanNode() : null
         );
     }
 
@@ -638,6 +639,9 @@ static void writeFragmentExec(PlanStreamOutput out, FragmentExec fragmentExec) t
         out.writeLogicalPlanNode(fragmentExec.fragment());
         out.writeOptionalNamedWriteable(fragmentExec.esFilter());
         out.writeOptionalVInt(fragmentExec.estimatedRowSize());
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_REDUCER_NODE_FRAGMENT)) {
+            out.writeOptionalPhysicalPlanNode(fragmentExec.reducer());
+        }
     }
 
     static GrokExec readGrokExec(PlanStreamInput in) throws IOException {

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamInput.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamInput.java
@@ -104,6 +104,10 @@ public PhysicalPlan readPhysicalPlanNode() throws IOException {
         return readNamed(PhysicalPlan.class);
     }
 
+    public PhysicalPlan readOptionalPhysicalPlanNode() throws IOException {
+        return readOptionalNamed(PhysicalPlan.class);
+    }
+
     public Source readSource() throws IOException {
         boolean hasSource = readBoolean();
         return hasSource ? readSourceWithText(this, configuration.query()) : Source.EMPTY;

diff --git a/...ck/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamOutput.java b/...ck/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamOutput.java
@@ -53,6 +53,15 @@ public void writePhysicalPlanNode(PhysicalPlan physicalPlan) throws IOException
         writeNamed(PhysicalPlan.class, physicalPlan);
     }
 
+    public void writeOptionalPhysicalPlanNode(PhysicalPlan physicalPlan) throws IOException {
+        if (physicalPlan == null) {
+            writeBoolean(false);
+        } else {
+            writeBoolean(true);
+            writePhysicalPlanNode(physicalPlan);
+        }
+    }
+
     public void writeSource(Source source) throws IOException {
         writeBoolean(true);
         writeSourceNoText(this, source);

diff --git a/...ugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizer.java b/...ugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizer.java
@@ -151,7 +151,8 @@ public PhysicalPlan apply(PhysicalPlan plan) {
                                     Source.EMPTY,
                                     new Project(logicalFragment.source(), logicalFragment, output),
                                     fragmentExec.esFilter(),
-                                    fragmentExec.estimatedRowSize()
+                                    fragmentExec.estimatedRowSize(),
+                                    fragmentExec.reducer()
                                 )
                             );
                         }

diff --git a/...ck/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/FragmentExec.java b/...ck/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/FragmentExec.java
@@ -20,6 +20,7 @@ public class FragmentExec extends LeafExec implements EstimatesRowSize {
 
     private final LogicalPlan fragment;
     private final QueryBuilder esFilter;
+    private final PhysicalPlan reducer; // datanode-level physical plan node that performs an intermediate (not partial) reduce
 
     /**
      * Estimate of the number of bytes that'll be loaded per position before
@@ -28,14 +29,15 @@ public class FragmentExec extends LeafExec implements EstimatesRowSize {
     private final int estimatedRowSize;
 
     public FragmentExec(LogicalPlan fragment) {
-        this(fragment.source(), fragment, null, 0);
+        this(fragment.source(), fragment, null, 0, null);
     }
 
-    public FragmentExec(Source source, LogicalPlan fragment, QueryBuilder esFilter, int estimatedRowSize) {
+    public FragmentExec(Source source, LogicalPlan fragment, QueryBuilder esFilter, int estimatedRowSize, PhysicalPlan reducer) {
         super(source);
         this.fragment = fragment;
         this.esFilter = esFilter;
         this.estimatedRowSize = estimatedRowSize;
+        this.reducer = reducer;
     }
 
     public LogicalPlan fragment() {
@@ -50,9 +52,13 @@ public Integer estimatedRowSize() {
         return estimatedRowSize;
     }
 
+    public PhysicalPlan reducer() {
+        return reducer;
+    }
+
     @Override
     protected NodeInfo<FragmentExec> info() {
-        return NodeInfo.create(this, FragmentExec::new, fragment, esFilter, estimatedRowSize);
+        return NodeInfo.create(this, FragmentExec::new, fragment, esFilter, estimatedRowSize, reducer);
     }
 
     @Override
@@ -65,12 +71,20 @@ public PhysicalPlan estimateRowSize(State state) {
         int estimatedRowSize = state.consumeAllFields(false);
         return Objects.equals(estimatedRowSize, this.estimatedRowSize)
             ? this
-            : new FragmentExec(source(), fragment, esFilter, estimatedRowSize);
+            : new FragmentExec(source(), fragment, esFilter, estimatedRowSize, reducer);
+    }
+
+    public FragmentExec withFilter(QueryBuilder filter) {
+        return Objects.equals(filter, this.esFilter) ? this : new FragmentExec(source(), fragment, filter, estimatedRowSize, reducer);
+    }
+
+    public FragmentExec withReducer(PhysicalPlan reducer) {
+        return Objects.equals(reducer, this.reducer) ? this : new FragmentExec(source(), fragment, esFilter, estimatedRowSize, reducer);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(fragment, esFilter, estimatedRowSize);
+        return Objects.hash(fragment, esFilter, estimatedRowSize, reducer);
     }
 
     @Override
@@ -86,7 +100,8 @@ public boolean equals(Object obj) {
         FragmentExec other = (FragmentExec) obj;
         return Objects.equals(fragment, other.fragment)
             && Objects.equals(esFilter, other.esFilter)
-            && Objects.equals(estimatedRowSize, other.estimatedRowSize);
+            && Objects.equals(estimatedRowSize, other.estimatedRowSize)
+            && Objects.equals(reducer, other.reducer);
     }
 
     @Override
@@ -97,7 +112,9 @@ public String nodeString() {
         sb.append(esFilter);
         sb.append(", estimatedRowSize=");
         sb.append(estimatedRowSize);
-        sb.append(", fragment=[<>\n");
+        sb.append(", reducer=[");
+        sb.append(reducer == null ? "" : reducer.toString());
+        sb.append("], fragment=[<>\n");
         sb.append(fragment.toString());
         sb.append("<>]]");
         return sb.toString();

diff --git a/...plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java b/...plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
@@ -149,7 +149,7 @@ public LocalExecutionPlanner(
     /**
      * turn the given plan into a list of drivers to execute
      */
-    public LocalExecutionPlan plan(PhysicalPlan node) {
+    public LocalExecutionPlan plan(PhysicalPlan localPhysicalPlan) {
         var context = new LocalExecutionPlannerContext(
             new ArrayList<>(),
             new Holder<>(DriverParallelism.SINGLE),
@@ -160,11 +160,11 @@ public LocalExecutionPlan plan(PhysicalPlan node) {
         );
 
         // workaround for https://github.com/elastic/elasticsearch/issues/99782
-        node = node.transformUp(
+        localPhysicalPlan = localPhysicalPlan.transformUp(
             AggregateExec.class,
             a -> a.getMode() == AggregateExec.Mode.FINAL ? new ProjectExec(a.source(), a, Expressions.asAttributes(a.aggregates())) : a
         );
-        PhysicalOperation physicalOperation = plan(node, context);
+        PhysicalOperation physicalOperation = plan(localPhysicalPlan, context);
 
         final TimeValue statusInterval = configuration.pragmas().statusInterval();
         context.addDriverFactory(
@@ -181,7 +181,7 @@ private PhysicalOperation plan(PhysicalPlan node, LocalExecutionPlannerContext c
         if (node instanceof AggregateExec aggregate) {
             return planAggregation(aggregate, context);
         } else if (node instanceof FieldExtractExec fieldExtractExec) {
-            return planFieldExtractNode(context, fieldExtractExec);
+            return planFieldExtractNode(fieldExtractExec, context);
         } else if (node instanceof ExchangeExec exchangeExec) {
             return planExchange(exchangeExec, context);
         } else if (node instanceof TopNExec topNExec) {
@@ -259,7 +259,7 @@ private PhysicalOperation planEsStats(EsStatsQueryExec statsQuery, LocalExecutio
         return PhysicalOperation.fromSource(luceneFactory, layout.build());
     }
 
-    private PhysicalOperation planFieldExtractNode(LocalExecutionPlannerContext context, FieldExtractExec fieldExtractExec) {
+    private PhysicalOperation planFieldExtractNode(FieldExtractExec fieldExtractExec, LocalExecutionPlannerContext context) {
         return physicalOperationProviders.fieldExtractPhysicalOperation(fieldExtractExec, plan(fieldExtractExec.child(), context));
     }
 

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/Mapper.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/Mapper.java
@@ -116,7 +116,7 @@ public PhysicalPlan map(LogicalPlan p) {
         throw new EsqlIllegalArgumentException("unsupported logical plan node [" + p.nodeName() + "]");
     }
 
-    private static boolean isPipelineBreaker(LogicalPlan p) {
+    static boolean isPipelineBreaker(LogicalPlan p) {
         return p instanceof Aggregate || p instanceof TopN || p instanceof Limit || p instanceof OrderBy;
     }