Dart: Async queries to realtime servers.

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

@@ -19,6 +19,7 @@
 
 package org.apache.druid.msq.dart.guice;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.jsontype.NamedType;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -35,8 +36,10 @@
 import org.apache.druid.guice.LifecycleModule;
 import org.apache.druid.guice.ManageLifecycle;
 import org.apache.druid.guice.ManageLifecycleAnnouncements;
+import org.apache.druid.guice.annotations.EscalatedGlobal;
 import org.apache.druid.guice.annotations.LoadScope;
 import org.apache.druid.guice.annotations.Self;
+import org.apache.druid.guice.annotations.Smile;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.concurrent.Execs;
@@ -50,6 +53,7 @@
 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;
@@ -58,6 +62,8 @@
 import org.apache.druid.msq.querykit.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.security.AuthorizerMapper;
 
@@ -156,6 +162,20 @@ public Outbox<ControllerMessage> createOutbox()
     {
       return new OutboxImpl<>();
     }
+
+    @Provides
+    public DartDataServerQueryHandlerFactory createDataServerQueryHandlerFactory(
+        @EscalatedGlobal ServiceClientFactory serviceClientFactory,
+        @Smile ObjectMapper smileMapper,
+        QueryToolChestWarehouse queryToolChestWarehouse
+    )
+    {
+      return new DartDataServerQueryHandlerFactory(
+          serviceClientFactory,
+          smileMapper,
+          queryToolChestWarehouse
+      );
+    }
   }
 
   @Override

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartDataServerQueryHandler.java

@@ -0,0 +1,153 @@
+/*
+ * 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.dart.worker;
+
+import com.fasterxml.jackson.databind.JavaType;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.discovery.DataServerClient;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.java.util.common.guava.Yielder;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.exec.DataServerQueryHandler;
+import org.apache.druid.msq.exec.DataServerQueryHandlerUtils;
+import org.apache.druid.msq.exec.DataServerQueryResult;
+import org.apache.druid.msq.input.table.DataServerRequestDescriptor;
+import org.apache.druid.msq.input.table.RichSegmentDescriptor;
+import org.apache.druid.query.Queries;
+import org.apache.druid.query.Query;
+import org.apache.druid.query.QueryToolChest;
+import org.apache.druid.query.QueryToolChestWarehouse;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.aggregation.MetricManipulatorFns;
+import org.apache.druid.query.context.DefaultResponseContext;
+import org.apache.druid.query.context.ResponseContext;
+import org.apache.druid.rpc.ServiceClientFactory;
+import org.apache.druid.rpc.ServiceLocation;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Dart implementation of {@link DataServerQueryHandler}. Issues queries asynchronously, with no retries.
+ */
+public class DartDataServerQueryHandler implements DataServerQueryHandler
+{
+  private final String dataSource;
+  private final ChannelCounters channelCounters;
+  private final ServiceClientFactory serviceClientFactory;
+  private final ObjectMapper objectMapper;
+  private final QueryToolChestWarehouse warehouse;
+  private final DataServerRequestDescriptor requestDescriptor;
+
+  public DartDataServerQueryHandler(
+      String dataSource,
+      ChannelCounters channelCounters,
+      ServiceClientFactory serviceClientFactory,
+      ObjectMapper objectMapper,
+      QueryToolChestWarehouse warehouse,
+      DataServerRequestDescriptor requestDescriptor
+  )
+  {
+    this.dataSource = dataSource;
+    this.channelCounters = channelCounters;
+    this.serviceClientFactory = serviceClientFactory;
+    this.objectMapper = objectMapper;
+    this.warehouse = warehouse;
+    this.requestDescriptor = requestDescriptor;
+  }
+
+  /**
+   * {@inheritDoc}
+   *
+   * This method returns immediately. The returned future resolves when the server has started sending back
+   * its response.
+   *
+   * Queries are issued once, without retries.
+   */
+  @Override
+  public <RowType, QueryType> ListenableFuture<DataServerQueryResult<RowType>> fetchRowsFromDataServer(
+      Query<QueryType> query,
+      Function<Sequence<QueryType>, Sequence<RowType>> mappingFunction,
+      Closer closer
+  )
+  {
+    final Query<QueryType> preparedQuery =
+        Queries.withSpecificSegments(
+            DataServerQueryHandlerUtils.prepareQuery(query, dataSource),
+            requestDescriptor.getSegments()
+                             .stream()
+                             .map(RichSegmentDescriptor::toPlainDescriptor)
+                             .collect(Collectors.toList())
+        );
+
+    final ServiceLocation serviceLocation =
+        ServiceLocation.fromDruidServerMetadata(requestDescriptor.getServerMetadata());
+    final DataServerClient dataServerClient = makeDataServerClient(serviceLocation);
+    final QueryToolChest<QueryType, Query<QueryType>> toolChest = warehouse.getToolChest(query);
+    final Function<QueryType, QueryType> preComputeManipulatorFn =
+        toolChest.makePreComputeManipulatorFn(query, MetricManipulatorFns.deserializing());
+    final JavaType queryResultType = toolChest.getBaseResultType();
+    final ResponseContext responseContext = new DefaultResponseContext();
+
+    return FutureUtils.transform(
+        dataServerClient.run(preparedQuery, responseContext, queryResultType, closer),
+        resultSequence -> {
+          final Yielder<RowType> yielder = DataServerQueryHandlerUtils.createYielder(
+              resultSequence.map(preComputeManipulatorFn),
+              mappingFunction,
+              channelCounters
+          );
+
+          final List<SegmentDescriptor> missingSegments =
+              DataServerQueryHandlerUtils.getMissingSegments(responseContext);
+
+          if (!missingSegments.isEmpty()) {
+            throw DruidException
+                .forPersona(DruidException.Persona.USER)
+                .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+                .build(
+                    "Segment[%s]%s not found on server[%s]. Please retry your query.",
+                    missingSegments.get(0),
+                    missingSegments.size() > 1 ? StringUtils.format(" and[%d] others", missingSegments.size() - 1) : "",
+                    serviceLocation.getHostAndPort()
+                );
+          }
+
+          return new DataServerQueryResult<>(
+              Collections.singletonList(yielder),
+              Collections.emptyList(),
+              dataSource
+          );
+        }
+    );
+  }
+
+  private DataServerClient makeDataServerClient(ServiceLocation serviceLocation)
+  {
+    return new DataServerClient(serviceClientFactory, serviceLocation, objectMapper);
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartDataServerQueryHandlerFactory.java

@@ -0,0 +1,65 @@
+/*
+ * 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.dart.worker;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.exec.DataServerQueryHandlerFactory;
+import org.apache.druid.msq.input.table.DataServerRequestDescriptor;
+import org.apache.druid.query.QueryToolChestWarehouse;
+import org.apache.druid.rpc.ServiceClientFactory;
+
+/**
+ * Factory for {@link DartDataServerQueryHandler}.
+ */
+public class DartDataServerQueryHandlerFactory implements DataServerQueryHandlerFactory
+{
+  private final ServiceClientFactory serviceClientFactory;
+  private final ObjectMapper objectMapper;
+  private final QueryToolChestWarehouse warehouse;
+
+  public DartDataServerQueryHandlerFactory(
+      ServiceClientFactory serviceClientFactory,
+      ObjectMapper objectMapper,
+      QueryToolChestWarehouse warehouse
+  )
+  {
+    this.serviceClientFactory = serviceClientFactory;
+    this.objectMapper = objectMapper;
+    this.warehouse = warehouse;
+  }
+
+  @Override
+  public DartDataServerQueryHandler createDataServerQueryHandler(
+      String dataSource,
+      ChannelCounters channelCounters,
+      DataServerRequestDescriptor requestDescriptor
+  )
+  {
+    return new DartDataServerQueryHandler(
+        dataSource,
+        channelCounters,
+        serviceClientFactory,
+        objectMapper,
+        warehouse,
+        requestDescriptor
+    );
+  }
+}

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

@@ -22,7 +22,6 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
-import org.apache.druid.client.coordinator.CoordinatorClient;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.guice.annotations.Self;
@@ -31,14 +30,12 @@
 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.exec.DataServerQueryHandlerFactory;
 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.query.DruidProcessingConfig;
 import org.apache.druid.query.QueryContext;
-import org.apache.druid.query.QueryToolChestWarehouse;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.policy.PolicyEnforcer;
 import org.apache.druid.rpc.ServiceClientFactory;
@@ -65,8 +62,7 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
   private final MemoryIntrospector memoryIntrospector;
   private final ProcessingBuffersProvider processingBuffersProvider;
   private final Outbox<ControllerMessage> outbox;
-  private final CoordinatorClient coordinatorClient;
-  private final QueryToolChestWarehouse warehouse;
+  private final DartDataServerQueryHandlerFactory dataServerQueryHandlerFactory;
   private final ServiceEmitter emitter;
 
   @Inject
@@ -84,8 +80,7 @@ public DartWorkerContextFactoryImpl(
       MemoryIntrospector memoryIntrospector,
       @Dart ProcessingBuffersProvider processingBuffersProvider,
       Outbox<ControllerMessage> outbox,
-      CoordinatorClient coordinatorClient,
-      QueryToolChestWarehouse warehouse,
+      DartDataServerQueryHandlerFactory dataServerQueryHandlerFactory,
       ServiceEmitter emitter
   )
   {
@@ -102,8 +97,7 @@ public DartWorkerContextFactoryImpl(
     this.memoryIntrospector = memoryIntrospector;
     this.processingBuffersProvider = processingBuffersProvider;
     this.outbox = outbox;
-    this.coordinatorClient = coordinatorClient;
-    this.warehouse = warehouse;
+    this.dataServerQueryHandlerFactory = dataServerQueryHandlerFactory;
     this.emitter = emitter;
   }
 
@@ -132,12 +126,7 @@ public WorkerContext build(
         outbox,
         tempDir,
         queryContext,
-        new DataServerQueryHandlerFactory(
-            coordinatorClient,
-            serviceClientFactory,
-            jsonMapper,
-            warehouse
-        ),
+        dataServerQueryHandlerFactory,
         emitter
     );
   }