ESQL: Compute engine support for tagged queries (#128521)

Begins adding support for running "tagged queries" to the compute
engine. Here, it's just the `LuceneSourceOperator` because that's
useful and contained.

Example time! Say you are running:
```
FROM foo
| STATS MAX(v) BY ROUND_TO(g, 0, 100, 1000, 100000)
```

It's *often* faster to run this as four queries:
* The docs that round to `0`
* The docs that round to `100`
* The docs that round to `1000`
* The docs that round to `100000`

This creates an ESQL operator that can run these queries, one after the
other and attach those tags.

Aggs uses this trick and it's *way* faster when it can push down count
queries, but it's still faster when it pushes doc loading things. This
implementation in `LuceneSourceOperator` is quite similar to the doc
loading version in _search.

I don't have performance measurements yet because I haven't plugged this
into the language. In _search we call this `filter-by-filter` and enable
it when each group averages to more than 5000 documents and when there
isn't an `_doc_count` field. It's faster in those cases not to push. I
expect we'll be pretty similar.

Author: nik9000

docs/reference/query-languages/esql/_snippets/functions/appendix/values.md

@@ -9,7 +9,6 @@
 
 import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.LeafCollector;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Weight;
@@ -44,7 +43,7 @@ public static class Factory extends LuceneOperator.Factory {
 
         public Factory(
             List<? extends ShardContext> contexts,
-            Function<ShardContext, Query> queryFunction,
+            Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
             DataPartitioning dataPartitioning,
             int taskConcurrency,
             int limit
@@ -121,6 +120,9 @@ protected Page getCheckedOutput() throws IOException {
             if (scorer == null) {
                 remainingDocs = 0;
             } else {
+                if (scorer.tags().isEmpty() == false) {
+                    throw new UnsupportedOperationException("tags not supported by " + getClass());
+                }
                 Weight weight = scorer.weight();
                 var leafReaderContext = scorer.leafReaderContext();
                 // see org.apache.lucene.search.TotalHitCountCollector

docs/reference/query-languages/esql/kibana/definition/functions/values.json

@@ -10,7 +10,6 @@
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.PointValues;
 import org.apache.lucene.index.SortedNumericDocValues;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.NumericUtils;
 import org.elasticsearch.compute.data.Block;
@@ -114,7 +113,7 @@ public final long evaluate(long value1, long value2) {
 
     public LuceneMaxFactory(
         List<? extends ShardContext> contexts,
-        Function<ShardContext, Query> queryFunction,
+        Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
         DataPartitioning dataPartitioning,
         int taskConcurrency,
         String fieldName,

docs/reference/query-languages/esql/kibana/docs/functions/values.md

@@ -10,7 +10,6 @@
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.PointValues;
 import org.apache.lucene.index.SortedNumericDocValues;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.NumericUtils;
 import org.elasticsearch.compute.data.Block;
@@ -114,7 +113,7 @@ public final long evaluate(long value1, long value2) {
 
     public LuceneMinFactory(
         List<? extends ShardContext> contexts,
-        Function<ShardContext, Query> queryFunction,
+        Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
         DataPartitioning dataPartitioning,
         int taskConcurrency,
         String fieldName,

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

@@ -102,6 +102,9 @@ public Page getCheckedOutput() throws IOException {
             if (scorer == null) {
                 remainingDocs = 0;
             } else {
+                if (scorer.tags().isEmpty() == false) {
+                    throw new UnsupportedOperationException("tags not supported by " + getClass());
+                }
                 final LeafReader reader = scorer.leafReaderContext().reader();
                 final Query query = scorer.weight().getQuery();
                 if (query == null || query instanceof MatchAllDocsQuery) {

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

@@ -97,7 +97,7 @@ public abstract static class Factory implements SourceOperator.SourceOperatorFac
          */
         protected Factory(
             List<? extends ShardContext> contexts,
-            Function<ShardContext, Query> queryFunction,
+            Function<ShardContext, List<LuceneSliceQueue.QueryAndTags>> queryFunction,
             DataPartitioning dataPartitioning,
             Function<Query, LuceneSliceQueue.PartitioningStrategy> autoStrategy,
             int taskConcurrency,
@@ -155,10 +155,13 @@ LuceneScorer getCurrentOrLoadNextScorer() {
             final PartialLeafReaderContext partialLeaf = currentSlice.getLeaf(sliceIndex++);
             logger.trace("Starting {}", partialLeaf);
             final LeafReaderContext leaf = partialLeaf.leafReaderContext();
-            if (currentScorer == null || currentScorer.leafReaderContext() != leaf) {
+            if (currentScorer == null // First time
+                || currentScorer.leafReaderContext() != leaf // Moved to a new leaf
+                || currentScorer.weight != currentSlice.weight() // Moved to a new query
+            ) {
                 final Weight weight = currentSlice.weight();
                 processedQueries.add(weight.getQuery());
-                currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, leaf);
+                currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.tags(), leaf);
             }
             assert currentScorer.maxPosition <= partialLeaf.maxDoc() : currentScorer.maxPosition + ">" + partialLeaf.maxDoc();
             currentScorer.maxPosition = partialLeaf.maxDoc();
@@ -177,15 +180,17 @@ static final class LuceneScorer {
         private final ShardContext shardContext;
         private final Weight weight;
         private final LeafReaderContext leafReaderContext;
+        private final List<Object> tags;
 
         private BulkScorer bulkScorer;
         private int position;
         private int maxPosition;
         private Thread executingThread;
 
-        LuceneScorer(ShardContext shardContext, Weight weight, LeafReaderContext leafReaderContext) {
+        LuceneScorer(ShardContext shardContext, Weight weight, List<Object> tags, LeafReaderContext leafReaderContext) {
             this.shardContext = shardContext;
             this.weight = weight;
+            this.tags = tags;
             this.leafReaderContext = leafReaderContext;
             reinitialize();
         }
@@ -230,6 +235,13 @@ Weight weight() {
         int position() {
             return position;
         }
+
+        /**
+         * Tags to add to the data returned by this query.
+         */
+        List<Object> tags() {
+            return tags;
+        }
     }
 
     @Override

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

@@ -14,7 +14,7 @@
 /**
  * Holds a list of multiple partial Lucene segments
  */
-public record LuceneSlice(ShardContext shardContext, List<PartialLeafReaderContext> leaves, Weight weight) {
+public record LuceneSlice(ShardContext shardContext, List<PartialLeafReaderContext> leaves, Weight weight, List<Object> tags) {
     int numLeaves() {
         return leaves.size();
     }

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

@@ -32,8 +32,47 @@
 
 /**
  * Shared Lucene slices between Lucene operators.
+ * <p>
+ *     Each shard is {@link #create built} with a list of queries to run and
+ *     tags to add to the queries ({@code List<QueryAndTags>}). Some examples:
+ * </p>
+ * <ul>
+ *     <li>
+ *         For queries like {@code FROM foo} we'll use a one element list
+ *         containing {@code match_all, []}. It loads all documents in the
+ *         index and append no extra fields to the loaded documents.
+ *     </li>
+ *     <li>
+ *         For queries like {@code FROM foo | WHERE a > 10} we'll use a one
+ *         element list containing {@code +single_value(a) +(a > 10), []}.
+ *         It loads all documents where {@code a} is single valued and
+ *         greater than 10.
+ *     </li>
+ *     <li>
+ *         For queries like {@code FROM foo | STATS MAX(b) BY ROUND_TO(a, 0, 100)}
+ *         we'll use a two element list containing
+ *         <ul>
+ *             <li>{@code +single_value(a) +(a < 100), [0]}</li>
+ *             <li>{@code +single_value(a) +(a >= 100), [100]}</li>
+ *         </ul>
+ *         It loads all documents in the index where {@code a} is single
+ *         valued and adds a constant {@code 0} to the documents where
+ *         {@code a < 100} and the constant {@code 100} to the documents
+ *         where {@code a >= 100}.
+ *     </li>
+ * </ul>
+ * <p>
+ *     IMPORTANT: Runners make no effort to deduplicate the results from multiple
+ *     queries. If you need to only see each document one time then make sure the
+ *     queries are mutually exclusive.
+ * </p>
  */
 public final class LuceneSliceQueue {
+    /**
+     * Query to run and tags to add to the results.
+     */
+    public record QueryAndTags(Query query, List<Object> tags) {}
+
     public static final int MAX_DOCS_PER_SLICE = 250_000; // copied from IndexSearcher
     public static final int MAX_SEGMENTS_PER_SLICE = 5; // copied from IndexSearcher
 
@@ -69,7 +108,7 @@ public Collection<String> remainingShardsIdentifiers() {
 
     public static LuceneSliceQueue create(
         List<? extends ShardContext> contexts,
-        Function<ShardContext, Query> queryFunction,
+        Function<ShardContext, List<QueryAndTags>> queryFunction,
         DataPartitioning dataPartitioning,
         Function<Query, PartitioningStrategy> autoStrategy,
         int taskConcurrency,
@@ -78,27 +117,29 @@ public static LuceneSliceQueue create(
         List<LuceneSlice> slices = new ArrayList<>();
         Map<String, PartitioningStrategy> partitioningStrategies = new HashMap<>(contexts.size());
         for (ShardContext ctx : contexts) {
-            Query query = queryFunction.apply(ctx);
-            query = scoreMode.needsScores() ? query : new ConstantScoreQuery(query);
-            /*
-             * Rewrite the query on the local index so things like fully
-             * overlapping range queries become match all. It's important
-             * to do this before picking the partitioning strategy so we
-             * can pick more aggressive strategies when the query rewrites
-             * into MatchAll.
-             */
-            try {
-                query = ctx.searcher().rewrite(query);
-            } catch (IOException e) {
-                throw new UncheckedIOException(e);
-            }
-            PartitioningStrategy partitioning = PartitioningStrategy.pick(dataPartitioning, autoStrategy, ctx, query);
-            partitioningStrategies.put(ctx.shardIdentifier(), partitioning);
-            List<List<PartialLeafReaderContext>> groups = partitioning.groups(ctx.searcher(), taskConcurrency);
-            Weight weight = weight(ctx, query, scoreMode);
-            for (List<PartialLeafReaderContext> group : groups) {
-                if (group.isEmpty() == false) {
-                    slices.add(new LuceneSlice(ctx, group, weight));
+            for (QueryAndTags queryAndExtra : queryFunction.apply(ctx)) {
+                Query query = queryAndExtra.query;
+                query = scoreMode.needsScores() ? query : new ConstantScoreQuery(query);
+                /*
+                 * Rewrite the query on the local index so things like fully
+                 * overlapping range queries become match all. It's important
+                 * to do this before picking the partitioning strategy so we
+                 * can pick more aggressive strategies when the query rewrites
+                 * into MatchAll.
+                 */
+                try {
+                    query = ctx.searcher().rewrite(query);
+                } catch (IOException e) {
+                    throw new UncheckedIOException(e);
+                }
+                PartitioningStrategy partitioning = PartitioningStrategy.pick(dataPartitioning, autoStrategy, ctx, query);
+                partitioningStrategies.put(ctx.shardIdentifier(), partitioning);
+                List<List<PartialLeafReaderContext>> groups = partitioning.groups(ctx.searcher(), taskConcurrency);
+                Weight weight = weight(ctx, query, scoreMode);
+                for (List<PartialLeafReaderContext> group : groups) {
+                    if (group.isEmpty() == false) {
+                        slices.add(new LuceneSlice(ctx, group, weight, queryAndExtra.tags));
+                    }
                 }
             }
         }