MSQ: Use VSF, give stages control over inputs.

This patch integrates MSQ with virtual storage. It also refactors how MSQ
reads inputs to give stages more control over how inputs are read and merged.
In particular, stages are now able to fully control merging logic.

The main changes:

1) Integrate with virtual storage: merge the two DataSegmentProvider impls
   (Dart and Task) into DataSegmentProviderImpl that relies on SegmentManager.

2) Give stages control over input merging: rework InputSliceReader to return
   ReadablePartitions directly, without embedding any merging logic. Break out
   StandardPartitionReader as a separate class.

Other changes:

1) Move ReadableInput to the querykit package. It is no longer specific to the
   MSQ framework.

2) Remove StandardStageProcessor, refactoring dependent code to not require it.

3) Remove ExternalColumnSelectorFactory wrapper. Type casting is now handled
   directly by RowBasedColumnSelectorFactory.

4) Include full query context in worker context, rather than just a subset.

Includes apache#18871.

Author: gianm

indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java

@@ -54,9 +54,10 @@ public SegmentCacheManagerFactory(
 
   public SegmentCacheManager manufacturate(File storageDir)
   {
-    final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig().withLocations(
-        Collections.singletonList(new StorageLocationConfig(storageDir, null, null))
-    );
+    final SegmentLoaderConfig loaderConfig =
+        new SegmentLoaderConfig()
+            .setLocations(Collections.singletonList(new StorageLocationConfig(storageDir, null, null)))
+            .setVirtualStorage(true, true);
     final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
     return new SegmentLocalCacheManager(
         storageLocations,

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

@@ -52,19 +52,20 @@
 import org.apache.druid.msq.dart.controller.http.DartQueryInfo;
 import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
 import org.apache.druid.msq.dart.controller.sql.DartSqlEngine;
-import org.apache.druid.msq.dart.worker.DartDataSegmentProvider;
 import org.apache.druid.msq.dart.worker.DartDataServerQueryHandlerFactory;
 import org.apache.druid.msq.dart.worker.DartWorkerContextFactory;
 import org.apache.druid.msq.dart.worker.DartWorkerContextFactoryImpl;
 import org.apache.druid.msq.dart.worker.DartWorkerRunner;
 import org.apache.druid.msq.dart.worker.http.DartWorkerResource;
+import org.apache.druid.msq.exec.DataSegmentProviderImpl;
 import org.apache.druid.msq.exec.MemoryIntrospector;
-import org.apache.druid.msq.querykit.DataSegmentProvider;
+import org.apache.druid.msq.input.table.DataSegmentProvider;
 import org.apache.druid.msq.rpc.ResourcePermissionMapper;
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryToolChestWarehouse;
 import org.apache.druid.rpc.ServiceClientFactory;
 import org.apache.druid.server.DruidNode;
+import org.apache.druid.server.SegmentManager;
 import org.apache.druid.server.security.AuthorizerMapper;
 
 import java.io.File;
@@ -104,11 +105,6 @@ public void configure(Binder binder)
             .to(DartWorkerContextFactoryImpl.class)
             .in(LazySingleton.class);
 
-      binder.bind(DataSegmentProvider.class)
-            .annotatedWith(Dart.class)
-            .to(DartDataSegmentProvider.class)
-            .in(LazySingleton.class);
-
       binder.bind(ResourcePermissionMapper.class)
             .annotatedWith(Dart.class)
             .to(DartResourcePermissionMapper.class);
@@ -139,6 +135,14 @@ public DartWorkerRunner createWorkerRunner(
       );
     }
 
+    @Provides
+    @LazySingleton
+    @Dart
+    public DataSegmentProvider createDataSegmentProvider(final SegmentManager segmentManager)
+    {
+      return new DataSegmentProviderImpl(segmentManager, null);
+    }
+
     @Provides
     @Dart
     public MessageRelayMonitor createMessageRelayMonitor(

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

@@ -30,7 +30,7 @@
 import org.apache.druid.msq.exec.WorkerMemoryParameters;
 import org.apache.druid.msq.exec.WorkerStorageParameters;
 import org.apache.druid.msq.kernel.StageId;
-import org.apache.druid.msq.querykit.DataSegmentProvider;
+import org.apache.druid.msq.input.table.DataSegmentProvider;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.segment.IndexIO;

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

@@ -42,7 +42,7 @@
 import org.apache.druid.msq.exec.WorkerMemoryParameters;
 import org.apache.druid.msq.exec.WorkerStorageParameters;
 import org.apache.druid.msq.kernel.WorkOrder;
-import org.apache.druid.msq.querykit.DataSegmentProvider;
+import org.apache.druid.msq.input.table.DataSegmentProvider;
 import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryContext;

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

@@ -33,7 +33,7 @@
 import org.apache.druid.msq.exec.MemoryIntrospector;
 import org.apache.druid.msq.exec.ProcessingBuffersProvider;
 import org.apache.druid.msq.exec.WorkerContext;
-import org.apache.druid.msq.querykit.DataSegmentProvider;
+import org.apache.druid.msq.input.table.DataSegmentProvider;
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.groupby.GroupingEngine;

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

@@ -41,8 +41,9 @@
 import org.apache.druid.data.input.impl.DimensionsSpec;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.frame.FrameType;
-import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
+import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory;
 import org.apache.druid.frame.channel.ReadableConcatFrameChannel;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
 import org.apache.druid.frame.key.ClusterBy;
 import org.apache.druid.frame.key.ClusterByPartition;
 import org.apache.druid.frame.key.ClusterByPartitions;
@@ -92,8 +93,7 @@
 import org.apache.druid.msq.counters.ChannelCounters;
 import org.apache.druid.msq.counters.CounterSnapshots;
 import org.apache.druid.msq.counters.CounterSnapshotsTree;
-import org.apache.druid.msq.indexing.InputChannelFactory;
-import org.apache.druid.msq.indexing.InputChannelsImpl;
+import org.apache.druid.msq.exec.std.StandardPartitionReader;
 import org.apache.druid.msq.indexing.LegacyMSQSpec;
 import org.apache.druid.msq.indexing.MSQControllerTask;
 import org.apache.druid.msq.indexing.MSQSpec;
@@ -142,14 +142,12 @@
 import org.apache.druid.msq.input.inline.InlineInputSpecSlicer;
 import org.apache.druid.msq.input.lookup.LookupInputSpec;
 import org.apache.druid.msq.input.lookup.LookupInputSpecSlicer;
-import org.apache.druid.msq.input.stage.InputChannels;
 import org.apache.druid.msq.input.stage.StageInputSpec;
 import org.apache.druid.msq.input.stage.StageInputSpecSlicer;
 import org.apache.druid.msq.input.table.TableInputSpec;
 import org.apache.druid.msq.kernel.QueryDefinition;
 import org.apache.druid.msq.kernel.StageDefinition;
 import org.apache.druid.msq.kernel.StageId;
-import org.apache.druid.msq.kernel.StagePartition;
 import org.apache.druid.msq.kernel.WorkOrder;
 import org.apache.druid.msq.kernel.controller.ControllerQueryKernel;
 import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig;
@@ -2797,31 +2795,25 @@ private void startQueryResultsReader()
       final FrameProcessorExecutor resultReaderExec = createResultReaderExec(queryId());
       resultReaderExec.registerCancellationId(RESULT_READER_CANCELLATION_ID);
 
-      ReadableConcatFrameChannel resultsChannel = null;
+      ReadableFrameChannel resultsChannel = null;
 
       try {
-        final InputChannels inputChannels = new InputChannelsImpl(
+        final StandardPartitionReader partitionReader = new StandardPartitionReader(
             queryDef,
-            queryKernel.getResultPartitionsForStage(finalStageId),
             inputChannelFactory,
             FrameWriterSpec.fromContext(querySpec.getContext()),
-            () -> ArenaMemoryAllocator.createOnHeap(5_000_000),
             resultReaderExec,
             RESULT_READER_CANCELLATION_ID,
-            null
+            null,
+            new ArenaMemoryAllocatorFactory(MultiStageQueryContext.getFrameSize(querySpec.getContext()))
         );
 
         resultsChannel = ReadableConcatFrameChannel.open(
             StreamSupport.stream(queryKernel.getResultPartitionsForStage(finalStageId).spliterator(), false)
                          .map(
                              readablePartition -> {
                                try {
-                                 return inputChannels.openChannel(
-                                     new StagePartition(
-                                         queryKernel.getStageDefinition(finalStageId).getId(),
-                                         readablePartition.getPartitionNumber()
-                                     )
-                                 );
+                                 return partitionReader.openChannel(readablePartition);
                                }
                                catch (IOException e) {
                                  throw new RuntimeException(e);
@@ -2852,7 +2844,7 @@ private void startQueryResultsReader()
       }
       catch (Throwable e) {
         // There was some issue setting up the result reader. Shut down the results channel and stop the executor.
-        final ReadableConcatFrameChannel finalResultsChannel = resultsChannel;
+        final ReadableFrameChannel finalResultsChannel = resultsChannel;
         throw CloseableUtils.closeAndWrapInCatch(
             e,
             () -> CloseableUtils.closeAll(

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

@@ -0,0 +1,88 @@
+/*
+ * 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.exec;
+
+import com.google.common.base.Preconditions;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.msq.counters.CounterNames;
+import org.apache.druid.msq.counters.CounterTracker;
+import org.apache.druid.msq.indexing.CountingReadableFrameChannel;
+import org.apache.druid.msq.input.InputSlice;
+import org.apache.druid.msq.input.stage.StageInputSlice;
+import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.msq.kernel.WorkOrder;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Wrapper around {@link InputChannelFactory} that increments counters as data is read.
+ */
+public class CountingInputChannelFactory implements InputChannelFactory
+{
+  private final InputChannelFactory baseFactory;
+  private final WorkOrder workOrder;
+  private final CounterTracker counterTracker;
+
+  public CountingInputChannelFactory(
+      InputChannelFactory baseFactory,
+      WorkOrder workOrder,
+      CounterTracker counterTracker
+  )
+  {
+    this.baseFactory = Preconditions.checkNotNull(baseFactory, "baseFactory");
+    this.workOrder = Preconditions.checkNotNull(workOrder, "workOrder");
+    this.counterTracker = Preconditions.checkNotNull(counterTracker, "counterTracker");
+  }
+
+  @Override
+  public ReadableFrameChannel openChannel(StageId stageId, int workerNumber, int partitionNumber) throws IOException
+  {
+    return new CountingReadableFrameChannel(
+        baseFactory.openChannel(stageId, workerNumber, partitionNumber),
+        counterTracker.channel(getCounterNameForStage(stageId.getStageNumber())),
+        partitionNumber
+    );
+  }
+
+  /**
+   * Returns the counter name based on the input number (position in {@link WorkOrder#getInputs()}) for a stage.
+   * If multiple match, or if none matches, uses the size of the input array (which ends up not corresponding to
+   * any input number, allowing detection of these situations when looking at counters).
+   */
+  private String getCounterNameForStage(int stageNumber)
+  {
+    final List<InputSlice> inputs = workOrder.getInputs();
+    Integer matchingSlice = null;
+    for (int i = 0; i < inputs.size(); i++) {
+      final InputSlice slice = inputs.get(i);
+      if (slice instanceof StageInputSlice && ((StageInputSlice) slice).getStageNumber() == stageNumber) {
+        if (matchingSlice == null) {
+          matchingSlice = i;
+        } else {
+          matchingSlice = null;
+          break;
+        }
+      }
+    }
+
+    return CounterNames.inputChannel(matchingSlice == null ? inputs.size() : matchingSlice);
+  }
+}