feat: Make segmentLoadAheadCount able to be configured at worker task level in addition to context (apache#19559)

* Make segmentLoadAheadCount able to be configured at worker task level in addition to context

* fixups based on review

Author: capistrant

docs/querying/dart.md

@@ -73,6 +73,7 @@ For Historicals, you can set the following configs:
 |---|---|---|
 | `druid.msq.dart.worker.concurrentQueries` | Maximum number of query workers that can run concurrently on a Historical. We recommend leaving this config at the default value. If need to change this value, set it to a value equal to or larger than `druid.msq.dart.controller.concurrentQueries` on your Brokers. If you don't, queries can get stuck waiting for each other. Don't set it to a value higher than the number of merge buffers. | Equal to the number of merge buffers |
 | `druid.msq.dart.worker.heapFraction` | Maximum amount of heap available for use across all Dart queries as a decimal. | 0.35 (35% of heap) |
+| `druid.msq.dart.worker.segmentLoadAheadCount` | Number of segments a worker will prefetch ahead of processing them. This lets tiers with different hardware tune prefetch independently. It acts as a worker-local default: a value supplied in the query context (`segmentLoadAheadCount`) always takes precedence when it is set. Non-positive values are treated as unset. | Unset (uses twice the number of processing threads) |
 
 
 ## Run a Dart query

multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerConfig.java

@@ -21,6 +21,9 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.msq.exec.MemoryIntrospector;
+import org.apache.druid.msq.querykit.ReadableInputQueue;
+
+import javax.annotation.Nullable;
 
 /**
  * Runtime configuration for workers (which run on Historicals).
@@ -41,6 +44,21 @@ public class DartWorkerConfig
   @JsonProperty("heapFraction")
   private double heapFraction = DEFAULT_HEAP_FRACTION;
 
+  /**
+   * Worker-local value for the segment load-ahead count used to size segment prefetch in {@link ReadableInputQueue}.
+   * <p>
+   * Defaults to null (unset), which leaves the value to query context (client or controller-default supplied) and the
+   * built-in {@code 2 * threadCount} fallback. When set to a positive value, it acts as a worker-local default used
+   * only when the query context does not supply a value.
+   * <p>
+   * This is per-worker-process configuration, set independently on each worker. It lets workers with different
+   * hardware (for example, separate tiers with more or less memory and storage bandwidth) tune their own prefetch
+   * depth, rather than relying solely on a single cluster-wide query-context default.
+   */
+  @JsonProperty("segmentLoadAheadCount")
+  @Nullable
+  private Integer segmentLoadAheadCount = null;
+
   public int getConcurrentQueries()
   {
     return concurrentQueries;
@@ -50,4 +68,10 @@ public double getHeapFraction()
   {
     return heapFraction;
   }
+
+  @Nullable
+  public Integer getSegmentLoadAheadCount()
+  {
+    return segmentLoadAheadCount;
+  }
 }

multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java

@@ -30,6 +30,7 @@
 import org.apache.druid.messages.server.Outbox;
 import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
 import org.apache.druid.msq.dart.controller.messages.PostCounters;
+import org.apache.druid.msq.dart.guice.DartWorkerConfig;
 import org.apache.druid.msq.exec.ControllerClient;
 import org.apache.druid.msq.exec.DataServerQueryHandlerFactory;
 import org.apache.druid.msq.exec.FrameContext;
@@ -57,6 +58,7 @@
 import org.apache.druid.utils.CloseableUtils;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.util.List;
 
@@ -92,6 +94,13 @@ public class DartWorkerContext implements WorkerContext
   private final ServiceEmitter emitter;
   private final int threadCount;
 
+  /**
+   * Worker-local segment load-ahead count from {@link DartWorkerConfig#getSegmentLoadAheadCount()}, or null if unset.
+   * Used as the default in {@link #segmentLoadAheadCount(WorkOrder)} when the query context does not supply a value.
+   */
+  @Nullable
+  private final Integer segmentLoadAheadCountConfig;
+
   /**
    * Lazy initialized upon call to {@link #frameContext(WorkOrder)}.
    */
@@ -109,6 +118,7 @@ public class DartWorkerContext implements WorkerContext
       final Injector injector,
       final DartWorkerClient workerClient,
       final DruidProcessingConfig processingConfig,
+      final DartWorkerConfig workerConfig,
       final SegmentWrangler segmentWrangler,
       final SegmentManager segmentManager,
       final VirtualStorageManager virtualStorageManager,
@@ -148,6 +158,9 @@ public class DartWorkerContext implements WorkerContext
     final int baseThreadCount = processingConfig.getNumThreads();
     final Integer maxThreads = MultiStageQueryContext.getMaxThreads(queryContext);
     this.threadCount = (maxThreads != null && maxThreads > 0) ? Math.min(baseThreadCount, maxThreads) : baseThreadCount;
+
+    // Worker-local segment load-ahead config from this worker's DartWorkerConfig.
+    this.segmentLoadAheadCountConfig = workerConfig.getSegmentLoadAheadCount();
   }
 
   @Override
@@ -277,6 +290,32 @@ public int threadCount()
     return threadCount;
   }
 
+  @Override
+  public int segmentLoadAheadCount(final WorkOrder workOrder)
+  {
+    final Integer fromContext = MultiStageQueryContext.getSegmentLoadAheadCount(workOrder.getWorkerContext());
+    return resolveSegmentLoadAheadCount(fromContext, segmentLoadAheadCountConfig, threadCount);
+  }
+
+  /**
+   * Determine which of the three potential sources of segment load ahead count to use.
+   * <p>
+   * Precedence is: a value supplied in the query context wins when set; otherwise the worker-local config is used
+   * when set to a positive value; lastly we fall back to {@code 2 * threadCount}.
+   */
+  static int resolveSegmentLoadAheadCount(
+      @Nullable final Integer fromContext,
+      @Nullable final Integer workerConfig,
+      final int threadCount
+  )
+  {
+    if (fromContext != null) {
+      return fromContext;
+    }
+    final boolean hasWorkerConfig = workerConfig != null && workerConfig > 0;
+    return hasWorkerConfig ? workerConfig : threadCount * 2;
+  }
+
   @Override
   public boolean includeAllCounters()
   {

multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContextFactoryImpl.java

@@ -31,6 +31,7 @@
 import org.apache.druid.messages.server.Outbox;
 import org.apache.druid.msq.dart.Dart;
 import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
+import org.apache.druid.msq.dart.guice.DartWorkerConfig;
 import org.apache.druid.msq.exec.MemoryIntrospector;
 import org.apache.druid.msq.exec.ProcessingBuffersProvider;
 import org.apache.druid.msq.exec.WorkerContext;
@@ -60,6 +61,7 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
   private final Injector injector;
   private final ServiceClientFactory serviceClientFactory;
   private final DruidProcessingConfig processingConfig;
+  private final DartWorkerConfig workerConfig;
   private final SegmentWrangler segmentWrangler;
   private final SegmentManager segmentManager;
   private final VirtualStorageManager virtualStorageManager;
@@ -80,6 +82,7 @@ public DartWorkerContextFactoryImpl(
       Injector injector,
       @EscalatedGlobal ServiceClientFactory serviceClientFactory,
       DruidProcessingConfig processingConfig,
+      DartWorkerConfig workerConfig,
       SegmentWrangler segmentWrangler,
       SegmentManager segmentManager,
       VirtualStorageManager virtualStorageManager,
@@ -99,6 +102,7 @@ public DartWorkerContextFactoryImpl(
     this.injector = injector;
     this.serviceClientFactory = serviceClientFactory;
     this.processingConfig = processingConfig;
+    this.workerConfig = workerConfig;
     this.segmentWrangler = segmentWrangler;
     this.coordinatorClient = coordinatorClient;
     this.segmentManager = segmentManager;
@@ -128,6 +132,7 @@ public WorkerContext build(
         injector,
         createWorkerClient(queryId),
         processingConfig,
+        workerConfig,
         segmentWrangler,
         segmentManager,
         virtualStorageManager,

multi-stage-query/src/main/java/org/apache/druid/msq/exec/ExecutionContext.java

@@ -91,6 +91,13 @@ public interface ExecutionContext
    */
   int threadCount();
 
+  /**
+   * Effective segment load-ahead count for {@link #workOrder()}, resolved by
+   * {@link WorkerContext#segmentLoadAheadCount(WorkOrder)}. Used to size the segment prefetch in
+   * {@link org.apache.druid.msq.querykit.ReadableInputQueue}.
+   */
+  int segmentLoadAheadCount();
+
   /**
    * Cancellation ID that must be provided to {@link FrameProcessorExecutor} when running work.
    */

multi-stage-query/src/main/java/org/apache/druid/msq/exec/ExecutionContextImpl.java

@@ -50,6 +50,7 @@ public class ExecutionContextImpl implements ExecutionContext
   private final FrameContext frameContext;
   private final CounterTracker counters;
   private final int maxOutstandingProcessors;
+  private final int segmentLoadAheadCount;
   private final String cancellationId;
   private final RunWorkOrderListener listener;
   private final Set<String> intermediateOutputChannelFactoryNames = Sets.newConcurrentHashSet();
@@ -65,6 +66,7 @@ public class ExecutionContextImpl implements ExecutionContext
       final FrameContext frameContext,
       final CounterTracker counters,
       final int maxOutstandingProcessors,
+      final int segmentLoadAheadCount,
       final String cancellationId,
       final RunWorkOrderListener listener
   )
@@ -79,6 +81,7 @@ public class ExecutionContextImpl implements ExecutionContext
     this.frameContext = frameContext;
     this.counters = counters;
     this.maxOutstandingProcessors = maxOutstandingProcessors;
+    this.segmentLoadAheadCount = segmentLoadAheadCount;
     this.cancellationId = cancellationId;
     this.listener = listener;
   }
@@ -147,6 +150,12 @@ public int threadCount()
     return maxOutstandingProcessors;
   }
 
+  @Override
+  public int segmentLoadAheadCount()
+  {
+    return segmentLoadAheadCount;
+  }
+
   @Override
   public String cancellationId()
   {

multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java

@@ -392,6 +392,7 @@ private ExecutionContext makeExecutionContext()
         frameContext,
         counterTracker,
         workerContext.threadCount(),
+        workerContext.segmentLoadAheadCount(workOrder),
         cancellationId,
         listener
     );

multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java

@@ -109,6 +109,20 @@ public interface WorkerContext extends Closeable
    */
   int threadCount();
 
+  /**
+   * Effective number of segments to load ahead of when they are needed while processing the given {@code workOrder},
+   * used to size the segment prefetch in {@link org.apache.druid.msq.querykit.ReadableInputQueue}.
+   *
+   * The default honors {@link MultiStageQueryContext#CTX_SEGMENT_LOAD_AHEAD_COUNT} from the work order's context
+   * (set by the controller from client and broker-default context), and otherwise falls back to
+   * {@code 2 * threadCount()}. Implementations may override to layer in worker-local configuration.
+   */
+  default int segmentLoadAheadCount(WorkOrder workOrder)
+  {
+    final Integer fromContext = MultiStageQueryContext.getSegmentLoadAheadCount(workOrder.getWorkerContext());
+    return fromContext != null ? fromContext : threadCount() * 2;
+  }
+
   /**
    * Fetch node info about self.
    */

multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java

@@ -51,7 +51,6 @@
 import org.apache.druid.msq.input.stage.StageInputSlice;
 import org.apache.druid.msq.input.table.SegmentsInputSlice;
 import org.apache.druid.msq.kernel.StageDefinition;
-import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.filter.SegmentPruner;
 import org.apache.druid.query.planning.ExecutionVertex;
@@ -294,12 +293,10 @@ private ReadableInputQueue makeBaseInputQueue(
     }
 
     final List<PhysicalInputSlice> filteredSlices = filterBaseInput(physicalInputSlices);
-    final Integer segmentLoadAheadCount =
-        MultiStageQueryContext.getSegmentLoadAheadCount(context.workOrder().getWorkerContext());
     return new ReadableInputQueue(
         new StandardPartitionReader(context),
         filteredSlices,
-        segmentLoadAheadCount != null ? segmentLoadAheadCount : context.threadCount() * 2
+        context.segmentLoadAheadCount()
     );
   }
 

multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java

@@ -278,7 +278,10 @@ public class MultiStageQueryContext
   public static final String CTX_MAX_THREADS = "maxThreads";
 
   /**
-   * Maximum number of segments to load ahead of them being needed. Used when setting up {@link ReadableInputQueue}.
+   * Number of segments to load ahead of them being needed. Used when setting up {@link ReadableInputQueue}.
+   * <p>
+   * A worker may be configured with a local default for this value. When this context value is set, it always wins;
+   * the worker-local default applies only when this context value is absent.
    */
   public static final String CTX_SEGMENT_LOAD_AHEAD_COUNT = "segmentLoadAheadCount";