MSQ: Allow for worker gaps.

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/stage/ReadablePartition.java

@@ -59,6 +59,18 @@ public static ReadablePartition striped(final int stageNumber, final int numWork
     return new ReadablePartition(stageNumber, workerNumbers, partitionNumber);
   }
 
+  /**
+   * Returns an output partition that is striped across a set of {@code workerNumbers}.
+   */
+  public static ReadablePartition striped(
+      final int stageNumber,
+      final IntSortedSet workerNumbers,
+      final int partitionNumber
+  )
+  {
+    return new ReadablePartition(stageNumber, workerNumbers, partitionNumber);
+  }
+
   /**
    * Returns an output partition that has been collected onto a single worker.
    */

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/stage/ReadablePartitions.java

@@ -24,6 +24,7 @@
 import it.unimi.dsi.fastutil.ints.Int2IntAVLTreeMap;
 import it.unimi.dsi.fastutil.ints.Int2IntSortedMap;
 import it.unimi.dsi.fastutil.ints.IntAVLTreeSet;
+import it.unimi.dsi.fastutil.ints.IntSortedSet;
 
 import java.util.Collections;
 import java.util.List;
@@ -39,6 +40,7 @@
 @JsonSubTypes(value = {
     @JsonSubTypes.Type(name = "collected", value = CollectedReadablePartitions.class),
     @JsonSubTypes.Type(name = "striped", value = StripedReadablePartitions.class),
+    @JsonSubTypes.Type(name = "sparseStriped", value = SparseStripedReadablePartitions.class),
     @JsonSubTypes.Type(name = "combined", value = CombinedReadablePartitions.class)
 })
 public interface ReadablePartitions extends Iterable<ReadablePartition>
@@ -59,7 +61,7 @@ static ReadablePartitions empty()
   /**
    * Combines various sets of partitions into a single set.
    */
-  static CombinedReadablePartitions combine(List<ReadablePartitions> readablePartitions)
+  static ReadablePartitions combine(List<ReadablePartitions> readablePartitions)
   {
     return new CombinedReadablePartitions(readablePartitions);
   }
@@ -68,7 +70,7 @@ static CombinedReadablePartitions combine(List<ReadablePartitions> readableParti
    * Returns a set of {@code numPartitions} partitions striped across {@code numWorkers} workers: each worker contains
    * a "stripe" of each partition.
    */
-  static StripedReadablePartitions striped(
+  static ReadablePartitions striped(
       final int stageNumber,
       final int numWorkers,
       final int numPartitions
@@ -82,11 +84,36 @@ static StripedReadablePartitions striped(
     return new StripedReadablePartitions(stageNumber, numWorkers, partitionNumbers);
   }
 
+  /**
+   * Returns a set of {@code numPartitions} partitions striped across {@code workers}: each worker contains
+   * a "stripe" of each partition.
+   */
+  static ReadablePartitions striped(
+      final int stageNumber,
+      final IntSortedSet workers,
+      final int numPartitions
+  )
+  {
+    final IntAVLTreeSet partitionNumbers = new IntAVLTreeSet();
+    for (int i = 0; i < numPartitions; i++) {
+      partitionNumbers.add(i);
+    }
+
+    if (workers.lastInt() == workers.size() - 1) {
+      // Dense worker set. Use StripedReadablePartitions for compactness (send a single number rather than the
+      // entire worker set) and for backwards compatibility (older workers cannot understand
+      // SparseStripedReadablePartitions).
+      return new StripedReadablePartitions(stageNumber, workers.size(), partitionNumbers);
+    } else {
+      return new SparseStripedReadablePartitions(stageNumber, workers, partitionNumbers);
+    }
+  }
+
   /**
    * Returns a set of partitions that have been collected onto specific workers: each partition is on exactly
    * one worker.
    */
-  static CollectedReadablePartitions collected(
+  static ReadablePartitions collected(
       final int stageNumber,
       final Map<Integer, Integer> partitionToWorkerMap
   )

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/stage/SparseStripedReadablePartitions.java

@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.input.stage;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.Iterators;
+import it.unimi.dsi.fastutil.ints.IntAVLTreeSet;
+import it.unimi.dsi.fastutil.ints.IntSortedSet;
+import org.apache.druid.msq.input.SlicerUtils;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * Set of partitions striped across a sparse set of {@code workers}. Each worker contains a "stripe" of each partition.
+ *
+ * @see StripedReadablePartitions dense version, where workers from [0..N) are all used.
+ */
+public class SparseStripedReadablePartitions implements ReadablePartitions
+{
+  private final int stageNumber;
+  private final IntSortedSet workers;
+  private final IntSortedSet partitionNumbers;
+
+  /**
+   * Constructor. Most callers should use {@link ReadablePartitions#striped(int, int, int)} instead, which takes
+   * a partition count rather than a set of partition numbers.
+   */
+  public SparseStripedReadablePartitions(
+      final int stageNumber,
+      final IntSortedSet workers,
+      final IntSortedSet partitionNumbers
+  )
+  {
+    this.stageNumber = stageNumber;
+    this.workers = workers;
+    this.partitionNumbers = partitionNumbers;
+  }
+
+  @JsonCreator
+  private SparseStripedReadablePartitions(
+      @JsonProperty("stageNumber") final int stageNumber,
+      @JsonProperty("workers") final Set<Integer> workers,
+      @JsonProperty("partitionNumbers") final Set<Integer> partitionNumbers
+  )
+  {
+    this(stageNumber, new IntAVLTreeSet(workers), new IntAVLTreeSet(partitionNumbers));
+  }
+
+  @Override
+  public Iterator<ReadablePartition> iterator()
+  {
+    return Iterators.transform(
+        partitionNumbers.iterator(),
+        partitionNumber -> ReadablePartition.striped(stageNumber, workers, partitionNumber)
+    );
+  }
+
+  @Override
+  public List<ReadablePartitions> split(final int maxNumSplits)
+  {
+    final List<ReadablePartitions> retVal = new ArrayList<>();
+
+    for (List<Integer> entries : SlicerUtils.makeSlicesStatic(partitionNumbers.iterator(), maxNumSplits)) {
+      if (!entries.isEmpty()) {
+        retVal.add(new SparseStripedReadablePartitions(stageNumber, workers, new IntAVLTreeSet(entries)));
+      }
+    }
+
+    return retVal;
+  }
+
+  @JsonProperty
+  int getStageNumber()
+  {
+    return stageNumber;
+  }
+
+  @JsonProperty
+  IntSortedSet getWorkers()
+  {
+    return workers;
+  }
+
+  @JsonProperty
+  IntSortedSet getPartitionNumbers()
+  {
+    return partitionNumbers;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    SparseStripedReadablePartitions that = (SparseStripedReadablePartitions) o;
+    return stageNumber == that.stageNumber
+           && Objects.equals(workers, that.workers)
+           && Objects.equals(partitionNumbers, that.partitionNumbers);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(stageNumber, workers, partitionNumbers);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "StripedReadablePartitions{" +
+           "stageNumber=" + stageNumber +
+           ", workers=" + workers +
+           ", partitionNumbers=" + partitionNumbers +
+           '}';
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/kernel/controller/ControllerStageTracker.java

@@ -403,7 +403,7 @@ void addPartialKeyInformationForWorker(
       throw new ISE("Stage does not gather result key statistics");
     }
 
-    if (workerNumber < 0 || workerNumber >= workerCount) {
+    if (!workerInputs.workers().contains(workerNumber)) {
       throw new IAE("Invalid workerNumber [%s]", workerNumber);
     }
 
@@ -522,7 +522,7 @@ void mergeClusterByStatisticsCollectorForTimeChunk(
       throw new ISE("Stage does not gather result key statistics");
     }
 
-    if (workerNumber < 0 || workerNumber >= workerCount) {
+    if (!workerInputs.workers().contains(workerNumber)) {
       throw new IAE("Invalid workerNumber [%s]", workerNumber);
     }
 
@@ -656,7 +656,7 @@ void mergeClusterByStatisticsCollectorForAllTimeChunks(
       throw new ISE("Stage does not gather result key statistics");
     }
 
-    if (workerNumber < 0 || workerNumber >= workerCount) {
+    if (!workerInputs.workers().contains(workerNumber)) {
       throw new IAE("Invalid workerNumber [%s]", workerNumber);
     }
 
@@ -763,7 +763,7 @@ void setClusterByPartitionBoundaries(ClusterByPartitions clusterByPartitions)
     this.resultPartitionBoundaries = clusterByPartitions;
     this.resultPartitions = ReadablePartitions.striped(
         stageDef.getStageNumber(),
-        workerCount,
+        workerInputs.workers(),
         clusterByPartitions.size()
     );
 
@@ -788,7 +788,7 @@ void setDoneReadingInputForWorker(final int workerNumber)
       throw DruidException.defensive("Cannot setDoneReadingInput for stage[%s], it is not sorting", stageDef.getId());
     }
 
-    if (workerNumber < 0 || workerNumber >= workerCount) {
+    if (!workerInputs.workers().contains(workerNumber)) {
       throw new IAE("Invalid workerNumber[%s] for stage[%s]", workerNumber, stageDef.getId());
     }
 
@@ -830,7 +830,7 @@ void setDoneReadingInputForWorker(final int workerNumber)
   @SuppressWarnings("unchecked")
   boolean setResultsCompleteForWorker(final int workerNumber, final Object resultObject)
   {
-    if (workerNumber < 0 || workerNumber >= workerCount) {
+    if (!workerInputs.workers().contains(workerNumber)) {
       throw new IAE("Invalid workerNumber [%s]", workerNumber);
     }
 
@@ -947,14 +947,18 @@ private void generateResultPartitionsAndBoundariesWithoutKeyStatistics()
         resultPartitionBoundaries = maybeResultPartitionBoundaries.valueOrThrow();
         resultPartitions = ReadablePartitions.striped(
             stageNumber,
-            workerCount,
+            workerInputs.workers(),
             resultPartitionBoundaries.size()
         );
-      } else if (shuffleSpec.kind() == ShuffleKind.MIX) {
-        resultPartitionBoundaries = ClusterByPartitions.oneUniversalPartition();
-        resultPartitions = ReadablePartitions.striped(stageNumber, workerCount, shuffleSpec.partitionCount());
       } else {
-        resultPartitions = ReadablePartitions.striped(stageNumber, workerCount, shuffleSpec.partitionCount());
+        if (shuffleSpec.kind() == ShuffleKind.MIX) {
+          resultPartitionBoundaries = ClusterByPartitions.oneUniversalPartition();
+        }
+        resultPartitions = ReadablePartitions.striped(
+            stageNumber,
+            workerInputs.workers(),
+            shuffleSpec.partitionCount()
+        );
       }
     } else {
       // No reshuffling: retain partitioning from nonbroadcast inputs.