Introduce grouping in TopN operator

Author: przemekwitek

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/TopNBenchmark.java

@@ -125,6 +125,7 @@ private static Operator operator(String data, int topCount) {
             topCount,
             elementTypes,
             encoders,
+            List.of(),
             IntStream.range(0, count).mapToObj(c -> new TopNOperator.SortOrder(c, false, false)).toList(),
             16 * 1024
         );

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/TopNOperator.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
@@ -23,11 +24,14 @@
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 
+import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
 
 /**
  * An operator that sorts "rows" of values by encoding the values to sort on, as bytes (using BytesRef). Each data type is encoded
@@ -194,6 +198,16 @@ private void writeValues(int position, BreakingBytesRefBuilder values) {
         }
     }
 
+    public record Partition(int channel) {
+
+        private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Partition.class);
+
+        @Override
+        public String toString() {
+            return "Partition[channel=" + this.channel + "]";
+        }
+    }
+
     public record SortOrder(int channel, boolean asc, boolean nullsFirst) {
 
         private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(SortOrder.class);
@@ -224,6 +238,7 @@ public record TopNOperatorFactory(
         int topCount,
         List<ElementType> elementTypes,
         List<TopNEncoder> encoders,
+        List<Partition> partitions,
         List<SortOrder> sortOrders,
         int maxPageSize
     ) implements OperatorFactory {
@@ -243,6 +258,7 @@ public TopNOperator get(DriverContext driverContext) {
                 topCount,
                 elementTypes,
                 encoders,
+                partitions,
                 sortOrders,
                 maxPageSize
             );
@@ -256,6 +272,8 @@ public String describe() {
                 + elementTypes
                 + ", encoders="
                 + encoders
+                + ", partitions="
+                + partitions
                 + ", sortOrders="
                 + sortOrders
                 + "]";
@@ -264,12 +282,14 @@ public String describe() {
 
     private final BlockFactory blockFactory;
     private final CircuitBreaker breaker;
-    private final Queue inputQueue;
+    private final Map<String, Queue> inputQueues;
 
+    private final int topCount;
     private final int maxPageSize;
 
     private final List<ElementType> elementTypes;
     private final List<TopNEncoder> encoders;
+    private final List<Partition> partitions;
     private final List<SortOrder> sortOrders;
 
     private Row spare;
@@ -304,16 +324,19 @@ public TopNOperator(
         int topCount,
         List<ElementType> elementTypes,
         List<TopNEncoder> encoders,
+        List<Partition> partitions,
         List<SortOrder> sortOrders,
         int maxPageSize
     ) {
         this.blockFactory = blockFactory;
         this.breaker = breaker;
+        this.topCount = topCount;
         this.maxPageSize = maxPageSize;
         this.elementTypes = elementTypes;
         this.encoders = encoders;
+        this.partitions = partitions;
         this.sortOrders = sortOrders;
-        this.inputQueue = new Queue(topCount);
+        this.inputQueues = new TreeMap<>();
     }
 
     static int compareRows(Row r1, Row r2) {
@@ -385,6 +408,8 @@ public void addInput(Page page) {
                 spareKeysPreAllocSize = Math.max(spare.keys.length(), spareKeysPreAllocSize / 2);
                 spareValuesPreAllocSize = Math.max(spare.values.length(), spareValuesPreAllocSize / 2);
 
+                String partitionKey = getPartitionKey(page, i);
+                Queue inputQueue = inputQueues.computeIfAbsent(partitionKey, key -> new Queue(topCount));
                 spare = inputQueue.insertWithOverflow(spare);
             }
         } finally {
@@ -394,6 +419,28 @@ public void addInput(Page page) {
         }
     }
 
+    /**
+     * Calculates the partition key of the i-th row of the given page.
+     *
+     * @param page page for which the partition key should be calculated
+     * @param i row index
+     * @return partition key of the i-th row of the given page
+     */
+    private String getPartitionKey(Page page, int i) {
+        if (partitions.isEmpty()) {
+            return "";
+        }
+        assert page.getPositionCount() > 0;
+        StringBuilder builder = new StringBuilder();
+        for (Partition partition : partitions) {
+            try (var block = page.getBlock(partition.channel).filter(i)) {
+                BytesRef partitionFieldValue = ((BytesRefBlock) block).getBytesRef(i, new BytesRef());
+                builder.append(partitionFieldValue.utf8ToString());
+            }
+        }
+        return builder.toString();
+    }
+
     @Override
     public void finish() {
         if (output == null) {
@@ -407,14 +454,17 @@ private Iterator<Page> toPages() {
             spare.close();
             spare = null;
         }
-        if (inputQueue.size() == 0) {
-            return Collections.emptyIterator();
-        }
-        List<Row> list = new ArrayList<>(inputQueue.size());
-        List<Page> result = new ArrayList<>();
-        ResultBuilder[] builders = null;
         boolean success = false;
+        List<Row> list = null;
+        ResultBuilder[] builders = null;
+        List<Page> result = new ArrayList<>();
+        // TODO: optimize case where all the queues are empty
         try {
+            for (var entry : inputQueues.entrySet()) {
+            Queue inputQueue = entry.getValue();
+
+            list = new ArrayList<>(inputQueue.size());
+            builders = null;
             while (inputQueue.size() > 0) {
                 list.add(inputQueue.pop());
             }
@@ -483,6 +533,7 @@ private Iterator<Page> toPages() {
                 }
             }
             assert builders == null;
+            }
             success = true;
             return result.iterator();
         } finally {
@@ -524,20 +575,20 @@ public Page getOutput() {
 
     @Override
     public void close() {
+        List<Releasable> releasables = new ArrayList<>();
+        releasables.addAll(inputQueues.values().stream().map(Releasables::wrap).toList());
+        releasables.add(output == null ? null : Releasables.wrap(() -> Iterators.map(output, p -> p::releaseBlocks)));
         /*
          * If we close before calling finish then spare and inputQueue will be live rows
          * that need closing. If we close after calling finish then the output iterator
          * will contain pages of results that have yet to be returned.
          */
-        Releasables.closeExpectNoException(
-            spare,
-            inputQueue == null ? null : Releasables.wrap(inputQueue),
-            output == null ? null : Releasables.wrap(() -> Iterators.map(output, p -> p::releaseBlocks))
-        );
+        Releasables.closeExpectNoException(spare, Releasables.wrap(releasables));
     }
 
-    private static long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TopNOperator.class) + RamUsageEstimator
-        .shallowSizeOfInstance(List.class) * 3;
+    private static long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TopNOperator.class)
+        + RamUsageEstimator.shallowSizeOfInstance(List.class) * 4
+        + RamUsageEstimator.shallowSizeOfInstance(Map.class);
 
     @Override
     public long ramBytesUsed() {
@@ -548,25 +599,34 @@ public long ramBytesUsed() {
         // These lists may slightly under-count, but it's not likely to be by much.
         size += RamUsageEstimator.alignObjectSize(arrHeader + ref * elementTypes.size());
         size += RamUsageEstimator.alignObjectSize(arrHeader + ref * encoders.size());
+        size += RamUsageEstimator.alignObjectSize(arrHeader + ref * partitions.size());
+        size += partitions.size() * Partition.SHALLOW_SIZE;
         size += RamUsageEstimator.alignObjectSize(arrHeader + ref * sortOrders.size());
         size += sortOrders.size() * SortOrder.SHALLOW_SIZE;
-        size += inputQueue.ramBytesUsed();
+        long ramBytesUsedSum = inputQueues.entrySet().stream()
+            .mapToLong(e -> e.getKey().getBytes(Charset.defaultCharset()).length + e.getValue().ramBytesUsed())
+            .sum();
+        size += ramBytesUsedSum;
         return size;
     }
 
     @Override
     public Status status() {
-        return new TopNOperatorStatus(inputQueue.size(), ramBytesUsed(), pagesReceived, pagesEmitted, rowsReceived, rowsEmitted);
+        int queueSizeSum = inputQueues.values().stream().mapToInt(Queue::size).sum();
+        return new TopNOperatorStatus(queueSizeSum, ramBytesUsed(), pagesReceived, pagesEmitted, rowsReceived, rowsEmitted);
     }
 
     @Override
     public String toString() {
+        int queueSizeSum = inputQueues.values().stream().mapToInt(Queue::size).sum();
         return "TopNOperator[count="
-            + inputQueue
+            + queueSizeSum + "/" + topCount
             + ", elementTypes="
             + elementTypes
             + ", encoders="
             + encoders
+            + ", partitions="
+            + partitions
             + ", sortOrders="
             + sortOrders
             + "]";

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java

@@ -135,6 +135,7 @@ protected TopNOperator.TopNOperatorFactory simple(SimpleOptions options) {
             4,
             List.of(LONG),
             List.of(DEFAULT_UNSORTABLE),
+            List.of(),
             List.of(new TopNOperator.SortOrder(0, true, false)),
             pageSize
         );
@@ -143,15 +144,15 @@ protected TopNOperator.TopNOperatorFactory simple(SimpleOptions options) {
     @Override
     protected Matcher<String> expectedDescriptionOfSimple() {
         return equalTo(
-            "TopNOperator[count=4, elementTypes=[LONG], encoders=[DefaultUnsortable], "
+            "TopNOperator[count=4, elementTypes=[LONG], encoders=[DefaultUnsortable], partitions=[], "
                 + "sortOrders=[SortOrder[channel=0, asc=true, nullsFirst=false]]]"
         );
     }
 
     @Override
     protected Matcher<String> expectedToStringOfSimple() {
         return equalTo(
-            "TopNOperator[count=0/4, elementTypes=[LONG], encoders=[DefaultUnsortable], "
+            "TopNOperator[count=0/4, elementTypes=[LONG], encoders=[DefaultUnsortable], partitions=[], "
                 + "sortOrders=[SortOrder[channel=0, asc=true, nullsFirst=false]]]"
         );
     }
@@ -216,6 +217,7 @@ public long accumulateObject(Object o, long shallowSize, Map<Field, Object> fiel
                 topCount,
                 List.of(LONG),
                 List.of(DEFAULT_UNSORTABLE),
+                List.of(),
                 List.of(new TopNOperator.SortOrder(0, true, false)),
                 pageSize
             ).get(context)
@@ -554,6 +556,7 @@ public void testCollectAllValues() {
                         topCount,
                         elementTypes,
                         encoders,
+                        List.of(),
                         List.of(new TopNOperator.SortOrder(0, false, false)),
                         randomPageSize()
                     )
@@ -643,6 +646,7 @@ public void testCollectAllValues_RandomMultiValues() {
                         topCount,
                         elementTypes,
                         encoders,
+                        List.of(),
                         List.of(new TopNOperator.SortOrder(0, false, false)),
                         randomPageSize()
                     )
@@ -677,6 +681,7 @@ private List<Tuple<Long, Long>> topNTwoColumns(
                         limit,
                         elementTypes,
                         encoder,
+                        List.of(),
                         sortOrders,
                         randomPageSize()
                     )
@@ -704,6 +709,7 @@ public void testTopNManyDescriptionAndToString() {
             10,
             List.of(BYTES_REF, BYTES_REF),
             List.of(UTF8, new FixedLengthTopNEncoder(fixedLength)),
+            List.of(),
             List.of(new TopNOperator.SortOrder(1, false, false), new TopNOperator.SortOrder(3, false, true)),
             randomPageSize()
         );
@@ -712,7 +718,7 @@ public void testTopNManyDescriptionAndToString() {
             .collect(Collectors.joining(", "));
         String tail = ", elementTypes=[BYTES_REF, BYTES_REF], encoders=[UTF8TopNEncoder, FixedLengthTopNEncoder["
             + fixedLength
-            + "]], sortOrders=["
+            + "]], partitions=[], sortOrders=["
             + sorts
             + "]]";
         assertThat(factory.describe(), equalTo("TopNOperator[count=10" + tail));
@@ -946,6 +952,7 @@ private void assertSortingOnMV(
                         topCount,
                         List.of(blockType),
                         List.of(encoder),
+                        List.of(),
                         List.of(sortOrders),
                         randomPageSize()
                     )
@@ -1076,6 +1083,7 @@ public void testRandomMultiValuesTopN() {
                 topCount,
                 elementTypes,
                 encoders,
+                List.of(),
                 uniqueOrders.stream().toList(),
                 rows
             ),
@@ -1119,6 +1127,7 @@ public void testIPSortingSingleValue() throws UnknownHostException {
                             ips.size(),
                             List.of(BYTES_REF),
                             List.of(TopNEncoder.IP),
+                            List.of(),
                             List.of(new TopNOperator.SortOrder(0, asc, randomBoolean())),
                             randomPageSize()
                         )
@@ -1245,6 +1254,7 @@ private void assertIPSortingOnMultiValues(
                             ips.size(),
                             List.of(BYTES_REF),
                             List.of(TopNEncoder.IP),
+                            List.of(),
                             List.of(new TopNOperator.SortOrder(0, asc, nullsFirst)),
                             randomPageSize()
                         )
@@ -1332,6 +1342,7 @@ public void testZeroByte() {
                         2,
                         List.of(BYTES_REF, INT),
                         List.of(TopNEncoder.UTF8, DEFAULT_UNSORTABLE),
+                        List.of(),
                         List.of(
                             new TopNOperator.SortOrder(0, true, randomBoolean()),
                             new TopNOperator.SortOrder(1, randomBoolean(), randomBoolean())
@@ -1371,6 +1382,7 @@ public void testErrorBeforeFullyDraining() {
                         topCount,
                         List.of(LONG),
                         List.of(DEFAULT_UNSORTABLE),
+                        List.of(),
                         List.of(new TopNOperator.SortOrder(0, true, randomBoolean())),
                         maxPageSize
                     )
@@ -1406,6 +1418,7 @@ public void testCloseWithoutCompleting() {
                 2,
                 List.of(INT),
                 List.of(DEFAULT_UNSORTABLE),
+                List.of(),
                 List.of(new TopNOperator.SortOrder(0, randomBoolean(), randomBoolean())),
                 randomPageSize()
             )
@@ -1429,6 +1442,7 @@ public void testRowResizes() {
                 10,
                 types,
                 encoders,
+                List.of(),
                 List.of(new TopNOperator.SortOrder(0, randomBoolean(), randomBoolean())),
                 randomPageSize()
             )

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/ExpressionWritables.java

@@ -126,6 +126,7 @@ public static List<NamedWriteableRegistry.Entry> expressions() {
         List<NamedWriteableRegistry.Entry> entries = new ArrayList<>();
         entries.addAll(ExpressionCoreWritables.expressions());
         entries.add(UnsupportedAttribute.EXPRESSION_ENTRY);
+        entries.add(Partition.ENTRY);
         entries.add(Order.ENTRY);
         return entries;
     }