Optimize MSQ realtime queries

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/LoadedSegmentDataProviderFactory.java → extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryHandlerFactory.java

@@ -25,6 +25,7 @@
 import org.apache.druid.java.util.common.concurrent.ScheduledExecutors;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.input.table.DataServerRequestDescriptor;
 import org.apache.druid.query.QueryToolChestWarehouse;
 import org.apache.druid.rpc.ServiceClientFactory;
 
@@ -33,19 +34,19 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * Creates new instances of {@link LoadedSegmentDataProvider} and manages the cancellation threadpool.
+ * Creates new instances of {@link DataServerQueryHandler} and manages the cancellation threadpool.
  */
-public class LoadedSegmentDataProviderFactory implements Closeable
+public class DataServerQueryHandlerFactory implements Closeable
 {
-  private static final Logger log = new Logger(LoadedSegmentDataProviderFactory.class);
+  private static final Logger log = new Logger(DataServerQueryHandlerFactory.class);
   private static final int DEFAULT_THREAD_COUNT = 4;
   private final CoordinatorClient coordinatorClient;
   private final ServiceClientFactory serviceClientFactory;
   private final ObjectMapper objectMapper;
   private final QueryToolChestWarehouse warehouse;
   private final ScheduledExecutorService queryCancellationExecutor;
 
-  public LoadedSegmentDataProviderFactory(
+  public DataServerQueryHandlerFactory(
       CoordinatorClient coordinatorClient,
       ServiceClientFactory serviceClientFactory,
       ObjectMapper objectMapper,
@@ -59,19 +60,21 @@ public LoadedSegmentDataProviderFactory(
     this.queryCancellationExecutor = ScheduledExecutors.fixed(DEFAULT_THREAD_COUNT, "query-cancellation-executor");
   }
 
-  public LoadedSegmentDataProvider createLoadedSegmentDataProvider(
+  public DataServerQueryHandler createDataServerQueryHandler(
       String dataSource,
-      ChannelCounters channelCounters
+      ChannelCounters channelCounters,
+      DataServerRequestDescriptor dataServerRequestDescriptor
   )
   {
-    return new LoadedSegmentDataProvider(
+    return new DataServerQueryHandler(
         dataSource,
         channelCounters,
         serviceClientFactory,
         coordinatorClient,
         objectMapper,
         warehouse,
-        queryCancellationExecutor
+        queryCancellationExecutor,
+        dataServerRequestDescriptor
     );
   }
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryResult.java

@@ -0,0 +1,60 @@
+/*
+ * 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.collect.ImmutableList;
+import org.apache.druid.java.util.common.guava.Yielder;
+import org.apache.druid.msq.input.table.RichSegmentDescriptor;
+import org.apache.druid.msq.input.table.SegmentsInputSlice;
+
+import java.util.List;
+
+/**
+ * Contains the results for a query to a dataserver. {@link #resultsYielders} contains the results fetched and
+ * {@link #segmentsInputSlice} is an {@link SegmentsInputSlice} containing the segments which have already been handed
+ * off, so that it can be fetched from deep storage.
+ */
+public class DataServerQueryResult<RowType>
+{
+
+  private final List<Yielder<RowType>> resultsYielders;
+
+  private final SegmentsInputSlice segmentsInputSlice;
+
+  public DataServerQueryResult(
+      List<Yielder<RowType>> resultsYielders,
+      List<RichSegmentDescriptor> handedOffSegments,
+      String dataSource
+  )
+  {
+    this.resultsYielders = resultsYielders;
+    this.segmentsInputSlice = new SegmentsInputSlice(dataSource, handedOffSegments, ImmutableList.of());
+  }
+
+  public List<Yielder<RowType>> getResultsYielders()
+  {
+    return resultsYielders;
+  }
+
+  public SegmentsInputSlice getHandedOffSegments()
+  {
+    return segmentsInputSlice;
+  }
+}

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

@@ -74,7 +74,7 @@ public interface WorkerContext
   DruidNode selfNode();
 
   Bouncer processorBouncer();
-  LoadedSegmentDataProviderFactory loadedSegmentDataProviderFactory();
+  DataServerQueryHandlerFactory dataServerQueryHandlerFactory();
 
   default File tempDir(int stageNumber, String id)
   {