MSQ: Use MapBinder for QueryKit implementations.

multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java

@@ -53,6 +53,7 @@
 import org.apache.druid.msq.indexing.report.MSQResultsReport;
 import org.apache.druid.msq.indexing.report.MSQStatusReport;
 import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
+import org.apache.druid.msq.querykit.MultiQueryKit;
 import org.apache.druid.msq.sql.MSQTaskQueryMaker;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
@@ -120,6 +121,7 @@ public class DartQueryMaker implements QueryMaker
   private final ServerConfig serverConfig;
 
   final QueryKitSpecFactory queryKitSpecFactory;
+  final MultiQueryKit queryKit;
 
   public DartQueryMaker(
       List<Entry<Integer, String>> fieldMapping,
@@ -129,6 +131,7 @@ public DartQueryMaker(
       DartControllerConfig controllerConfig,
       ExecutorService controllerExecutor,
       QueryKitSpecFactory queryKitSpecFactory,
+      MultiQueryKit queryKit,
       ServerConfig serverConfig
   )
   {
@@ -139,6 +142,7 @@ public DartQueryMaker(
     this.controllerConfig = controllerConfig;
     this.controllerExecutor = controllerExecutor;
     this.queryKitSpecFactory = queryKitSpecFactory;
+    this.queryKit = queryKit;
     this.serverConfig = serverConfig;
   }
 

multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java

@@ -41,6 +41,7 @@
 import org.apache.druid.msq.dart.guice.DartControllerConfig;
 import org.apache.druid.msq.exec.QueryKitSpecFactory;
 import org.apache.druid.msq.indexing.error.CancellationReason;
+import org.apache.druid.msq.querykit.MultiQueryKit;
 import org.apache.druid.msq.sql.DartQueryKitSpecFactory;
 import org.apache.druid.msq.sql.MSQTaskSqlEngine;
 import org.apache.druid.query.DefaultQueryConfig;
@@ -82,6 +83,7 @@ public class DartSqlEngine implements SqlEngine
   private final ExecutorService controllerExecutor;
   private final ServerConfig serverConfig;
   private final QueryKitSpecFactory queryKitSpecFactory;
+  private final MultiQueryKit queryKit;
   private final DefaultQueryConfig dartQueryConfig;
   private final SqlToolbox toolbox;
   private final DartSqlClients sqlClients;
@@ -92,6 +94,7 @@ public DartSqlEngine(
       DartControllerRegistry controllerRegistry,
       DartControllerConfig controllerConfig,
       DartQueryKitSpecFactory queryKitSpecFactory,
+      MultiQueryKit queryKit,
       ServerConfig serverConfig,
       @Dart DefaultQueryConfig dartQueryConfig,
       SqlToolbox toolbox,
@@ -104,6 +107,7 @@ public DartSqlEngine(
         controllerConfig,
         Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s"),
         queryKitSpecFactory,
+        queryKit,
         serverConfig,
         dartQueryConfig,
         toolbox,
@@ -117,6 +121,7 @@ public DartSqlEngine(
       DartControllerConfig controllerConfig,
       ExecutorService controllerExecutor,
       QueryKitSpecFactory queryKitSpecFactory,
+      MultiQueryKit queryKit,
       ServerConfig serverConfig,
       DefaultQueryConfig dartQueryConfig,
       SqlToolbox toolbox,
@@ -128,6 +133,7 @@ public DartSqlEngine(
     this.controllerConfig = controllerConfig;
     this.controllerExecutor = controllerExecutor;
     this.queryKitSpecFactory = queryKitSpecFactory;
+    this.queryKit = queryKit;
     this.serverConfig = serverConfig;
     this.dartQueryConfig = dartQueryConfig;
     this.toolbox = toolbox;
@@ -216,6 +222,7 @@ public QueryMaker buildQueryMakerForSelect(RelRoot relRoot, PlannerContext plann
         controllerConfig,
         controllerExecutor,
         queryKitSpecFactory,
+        queryKit,
         serverConfig
     );
     if (plannerContext.queryContext().isPrePlanned()) {

multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/PrePlannedDartQueryMaker.java

@@ -125,8 +125,7 @@ private QueryDefMSQSpec buildMSQSpec(
         querySpec.getQuery(),
         plannerContext.getJsonMapper(),
         dartQueryMaker.queryKitSpecFactory.makeQueryKitSpec(
-            QueryKitBasedMSQPlanner
-                .makeQueryControllerToolKit(querySpec.getContext(), plannerContext.getJsonMapper()),
+            dartQueryMaker.queryKit,
             dartQueryId,
             querySpec.getTuningConfig(),
             querySpec.getContext()

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

@@ -155,6 +155,7 @@
 import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig;
 import org.apache.druid.msq.kernel.controller.ControllerStagePhase;
 import org.apache.druid.msq.kernel.controller.WorkerInputs;
+import org.apache.druid.msq.querykit.MultiQueryKit;
 import org.apache.druid.msq.querykit.QueryKitUtils;
 import org.apache.druid.msq.shuffle.input.DurableStorageInputChannelFactory;
 import org.apache.druid.msq.shuffle.input.WorkerInputChannelFactory;
@@ -715,7 +716,7 @@ private QueryDefinition initializeQueryDefAndState()
           legacyQuery,
           context.jsonMapper(),
           queryKitSpecFactory.makeQueryKitSpec(
-              QueryKitBasedMSQPlanner.makeQueryControllerToolKit(querySpec.getContext(), context.jsonMapper()),
+              context.injector().getInstance(MultiQueryKit.class),
               context.queryId(),
               querySpec.getTuningConfig(),
               querySpec.getContext()

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

@@ -36,30 +36,20 @@
 import org.apache.druid.msq.kernel.QueryDefinition;
 import org.apache.druid.msq.kernel.QueryDefinitionBuilder;
 import org.apache.druid.msq.kernel.StageDefinition;
-import org.apache.druid.msq.querykit.MultiQueryKit;
-import org.apache.druid.msq.querykit.QueryKit;
 import org.apache.druid.msq.querykit.QueryKitSpec;
 import org.apache.druid.msq.querykit.QueryKitUtils;
 import org.apache.druid.msq.querykit.ShuffleSpecFactory;
-import org.apache.druid.msq.querykit.WindowOperatorQueryKit;
-import org.apache.druid.msq.querykit.groupby.GroupByQueryKit;
 import org.apache.druid.msq.querykit.results.ExportResultsStageProcessor;
 import org.apache.druid.msq.querykit.results.QueryResultStageProcessor;
-import org.apache.druid.msq.querykit.scan.ScanQueryKit;
 import org.apache.druid.msq.util.MSQTaskQueryMakerUtils;
 import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryContext;
-import org.apache.druid.query.groupby.GroupByQuery;
-import org.apache.druid.query.operator.WindowOperatorQuery;
-import org.apache.druid.query.scan.ScanQuery;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.sql.calcite.planner.ColumnMappings;
 import org.apache.druid.sql.http.ResultFormat;
 import org.apache.druid.storage.ExportStorageProvider;
 
-import java.util.Map;
-
 public class QueryKitBasedMSQPlanner
 {
   private final MSQSpec querySpec;
@@ -91,25 +81,6 @@ public QueryKitBasedMSQPlanner(
     this.queryKitSpec = queryKitSpec;
   }
 
-  @SuppressWarnings("rawtypes")
-  public static QueryKit<Query<?>> makeQueryControllerToolKit(QueryContext queryContext, ObjectMapper jsonMapper)
-  {
-    final Map<Class<? extends Query>, QueryKit> kitMap =
-        ImmutableMap.<Class<? extends Query>, QueryKit>builder()
-                    .put(ScanQuery.class, new ScanQueryKit(jsonMapper))
-                    .put(GroupByQuery.class, new GroupByQueryKit(jsonMapper))
-                    .put(
-                        WindowOperatorQuery.class,
-                        new WindowOperatorQueryKit(
-                            jsonMapper,
-                            MultiStageQueryContext.isWindowFunctionOperatorTransformationEnabled(queryContext)
-                        )
-                    )
-                    .build();
-
-    return new MultiQueryKit(kitMap);
-  }
-
   @SuppressWarnings("unchecked")
   public QueryDefinition makeQueryDefinition()
   {

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

@@ -0,0 +1,53 @@
+/*
+ * 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.guice;
+
+import com.google.inject.Binder;
+import com.google.inject.TypeLiteral;
+import com.google.inject.multibindings.MapBinder;
+import org.apache.druid.msq.querykit.QueryKit;
+import org.apache.druid.query.Query;
+
+/**
+ * Utility class for MSQ-related Guice bindings.
+ */
+public class MSQBinders
+{
+  /**
+   * Creates a MapBinder for QueryKit implementations. Extensions can use this
+   * to register their own QueryKit implementations for custom query types.
+   *
+   * Example usage:
+   * <pre>
+   * MSQBinders.queryKitBinder(binder)
+   *     .addBinding(MyCustomQuery.class)
+   *     .to(MyCustomQueryKit.class);
+   * </pre>
+   */
+  @SuppressWarnings("rawtypes")
+  public static MapBinder<Class<? extends Query>, QueryKit> queryKitBinder(Binder binder)
+  {
+    return MapBinder.newMapBinder(
+        binder,
+        new TypeLiteral<>() {},
+        new TypeLiteral<>() {}
+    );
+  }
+}

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

@@ -24,6 +24,7 @@
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Binder;
+import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.msq.counters.ChannelCounters;
 import org.apache.druid.msq.counters.CounterSnapshotsSerializer;
@@ -87,16 +88,23 @@
 import org.apache.druid.msq.input.table.TableInputSpec;
 import org.apache.druid.msq.kernel.NilExtraInfoHolder;
 import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.msq.querykit.MultiQueryKit;
 import org.apache.druid.msq.querykit.RestrictedInputNumberDataSource;
+import org.apache.druid.msq.querykit.WindowOperatorQueryKit;
 import org.apache.druid.msq.querykit.WindowOperatorQueryStageProcessor;
 import org.apache.druid.msq.querykit.common.OffsetLimitStageProcessor;
 import org.apache.druid.msq.querykit.common.SortMergeJoinStageProcessor;
 import org.apache.druid.msq.querykit.groupby.GroupByPostShuffleStageProcessor;
 import org.apache.druid.msq.querykit.groupby.GroupByPreShuffleStageProcessor;
+import org.apache.druid.msq.querykit.groupby.GroupByQueryKit;
 import org.apache.druid.msq.querykit.results.ExportResultsStageProcessor;
 import org.apache.druid.msq.querykit.results.QueryResultStageProcessor;
+import org.apache.druid.msq.querykit.scan.ScanQueryKit;
 import org.apache.druid.msq.querykit.scan.ScanQueryStageProcessor;
 import org.apache.druid.msq.util.PassthroughAggregatorFactory;
+import org.apache.druid.query.groupby.GroupByQuery;
+import org.apache.druid.query.operator.WindowOperatorQuery;
+import org.apache.druid.query.scan.ScanQuery;
 
 import java.util.Collections;
 import java.util.List;
@@ -214,5 +222,21 @@ public List<? extends Module> getJacksonModules()
   @Override
   public void configure(Binder binder)
   {
+    binder.bind(MultiQueryKit.class).in(LazySingleton.class);
+
+    MSQBinders.queryKitBinder(binder)
+              .addBinding(ScanQuery.class)
+              .to(ScanQueryKit.class);
+    binder.bind(ScanQueryKit.class).in(LazySingleton.class);
+
+    MSQBinders.queryKitBinder(binder)
+              .addBinding(GroupByQuery.class)
+              .to(GroupByQueryKit.class);
+    binder.bind(GroupByQueryKit.class).in(LazySingleton.class);
+
+    MSQBinders.queryKitBinder(binder)
+              .addBinding(WindowOperatorQuery.class)
+              .to(WindowOperatorQueryKit.class);
+    binder.bind(WindowOperatorQueryKit.class).in(LazySingleton.class);
   }
 }

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

@@ -20,20 +20,28 @@
 package org.apache.druid.msq.querykit;
 
 import com.google.common.base.Preconditions;
+import com.google.inject.Binder;
+import com.google.inject.Inject;
 import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.msq.guice.MSQBinders;
 import org.apache.druid.msq.kernel.QueryDefinition;
 import org.apache.druid.query.Query;
 
 import java.util.Map;
 
 /**
  * Delegates to other {@link QueryKit} implementations based on the class of the {@link Query}.
+ *
+ * The map of query types to QueryKit implementations is populated via Guice MapBinder.
+ * Extensions can register additional QueryKit implementations using
+ * {@link MSQBinders#queryKitBinder(Binder)}.
  */
 @SuppressWarnings("rawtypes")
 public class MultiQueryKit implements QueryKit<Query<?>>
 {
   private final Map<Class<? extends Query>, QueryKit> toolKitMap;
 
+  @Inject
   public MultiQueryKit(final Map<Class<? extends Query>, QueryKit> toolKitMap)
   {
     this.toolKitMap = Preconditions.checkNotNull(toolKitMap, "toolKitMap");

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

@@ -20,9 +20,11 @@
 package org.apache.druid.msq.querykit;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Inject;
 import org.apache.druid.frame.key.ClusterBy;
 import org.apache.druid.frame.key.KeyColumn;
 import org.apache.druid.frame.key.KeyOrder;
+import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.java.util.common.logger.Logger;
@@ -56,12 +58,11 @@ public class WindowOperatorQueryKit implements QueryKit<WindowOperatorQuery>
 {
   private static final Logger log = new Logger(WindowOperatorQueryKit.class);
   private final ObjectMapper jsonMapper;
-  private final boolean isOperatorTransformationEnabled;
 
-  public WindowOperatorQueryKit(ObjectMapper jsonMapper, boolean isOperatorTransformationEnabled)
+  @Inject
+  public WindowOperatorQueryKit(@Json final ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
-    this.isOperatorTransformationEnabled = isOperatorTransformationEnabled;
   }
 
   @Override
@@ -88,6 +89,9 @@ public QueryDefinition makeQueryDefinition(
                                                           .getFinalStageDefinition()
                                                           .getSignature();
 
+    final boolean isOperatorTransformationEnabled =
+        MultiStageQueryContext.isWindowFunctionOperatorTransformationEnabled(originalQuery.context());
+
     final WindowStages windowStages = new WindowStages(
         originalQuery,
         jsonMapper,

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

@@ -21,9 +21,11 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Preconditions;
+import com.google.inject.Inject;
 import org.apache.druid.frame.key.ClusterBy;
 import org.apache.druid.frame.key.KeyColumn;
 import org.apache.druid.frame.key.KeyOrder;
+import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.java.util.common.granularity.Granularity;
@@ -58,7 +60,8 @@ public class GroupByQueryKit implements QueryKit<GroupByQuery>
 {
   private final ObjectMapper jsonMapper;
 
-  public GroupByQueryKit(ObjectMapper jsonMapper)
+  @Inject
+  public GroupByQueryKit(@Json final ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }

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

@@ -20,9 +20,11 @@
 package org.apache.druid.msq.querykit.scan;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Inject;
 import org.apache.druid.frame.key.ClusterBy;
 import org.apache.druid.frame.key.KeyColumn;
 import org.apache.druid.frame.key.KeyOrder;
+import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.msq.input.stage.StageInputSpec;
 import org.apache.druid.msq.kernel.MixShuffleSpec;
@@ -50,7 +52,8 @@ public class ScanQueryKit implements QueryKit<ScanQuery>
 {
   private final ObjectMapper jsonMapper;
 
-  public ScanQueryKit(final ObjectMapper jsonMapper)
+  @Inject
+  public ScanQueryKit(@Json final ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }