From 40b3d2769a65fe445fa9f9eea8863f72dcd656f0 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Mon, 29 Dec 2025 12:02:59 -0800 Subject: [PATCH 1/3] MSQ: Use MapBinder for QueryKit implementations. This allows extensions to provide QueryKit implementations. --- .../dart/controller/sql/DartQueryMaker.java | 4 ++ .../dart/controller/sql/DartSqlEngine.java | 7 +++ .../sql/PrePlannedDartQueryMaker.java | 3 +- .../apache/druid/msq/exec/ControllerImpl.java | 3 +- .../msq/exec/QueryKitBasedMSQPlanner.java | 29 ---------- .../apache/druid/msq/guice/MSQBinders.java | 53 +++++++++++++++++++ .../druid/msq/guice/MSQIndexingModule.java | 24 +++++++++ .../druid/msq/querykit/MultiQueryKit.java | 8 +++ .../msq/querykit/WindowOperatorQueryKit.java | 10 ++-- .../msq/querykit/groupby/GroupByQueryKit.java | 5 +- .../druid/msq/querykit/scan/ScanQueryKit.java | 5 +- .../controller/http/DartSqlResourceTest.java | 2 + .../apache/druid/msq/test/MSQTestBase.java | 1 + 13 files changed, 117 insertions(+), 37 deletions(-) create mode 100644 multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQBinders.java diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index c40a59f208b8..519179e4dddb 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -41,6 +41,7 @@ import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.ControllerImpl; import org.apache.druid.msq.exec.QueryKitSpecFactory; +import org.apache.druid.msq.querykit.MultiQueryKit; import org.apache.druid.msq.exec.QueryListener; import org.apache.druid.msq.exec.ResultsContext; import org.apache.druid.msq.indexing.LegacyMSQSpec; @@ -120,6 +121,7 @@ public class DartQueryMaker implements QueryMaker private final ServerConfig serverConfig; final QueryKitSpecFactory queryKitSpecFactory; + final MultiQueryKit queryKit; public DartQueryMaker( List> 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; } diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index cbef28d98ec5..f61ed243fb37 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/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()) { diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/PrePlannedDartQueryMaker.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/PrePlannedDartQueryMaker.java index d22b7ebe028f..e6cdc9e0bb26 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/PrePlannedDartQueryMaker.java +++ b/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() diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index 0330c1f5b39b..b72e660c0f60 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/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() diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java index 15cd96b3f054..cf2ca91c8d1d 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java +++ b/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> makeQueryControllerToolKit(QueryContext queryContext, ObjectMapper jsonMapper) - { - final Map, QueryKit> kitMap = - ImmutableMap., 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() { diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQBinders.java b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQBinders.java new file mode 100644 index 000000000000..0a6b0827324a --- /dev/null +++ b/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: + *
+   * MSQBinders.queryKitBinder(binder)
+   *     .addBinding(MyCustomQuery.class)
+   *     .to(MyCustomQueryKit.class);
+   * 
+ */ + @SuppressWarnings("rawtypes") + public static MapBinder, QueryKit> queryKitBinder(Binder binder) + { + return MapBinder.newMapBinder( + binder, + new TypeLiteral<>() {}, + new TypeLiteral<>() {} + ); + } +} diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java index 30ffdf2cb6ac..e9ca68443844 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java +++ b/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,8 +88,15 @@ 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.groupby.GroupByQueryKit; +import org.apache.druid.msq.querykit.scan.ScanQueryKit; +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.msq.querykit.common.OffsetLimitStageProcessor; import org.apache.druid.msq.querykit.common.SortMergeJoinStageProcessor; import org.apache.druid.msq.querykit.groupby.GroupByPostShuffleStageProcessor; @@ -214,5 +222,21 @@ public List 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); } } diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/MultiQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/MultiQueryKit.java index 3129bbfacb9c..ef189a554126 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/MultiQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/MultiQueryKit.java @@ -20,7 +20,10 @@ 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; @@ -28,12 +31,17 @@ /** * 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> { private final Map, QueryKit> toolKitMap; + @Inject public MultiQueryKit(final Map, QueryKit> toolKitMap) { this.toolKitMap = Preconditions.checkNotNull(toolKitMap, "toolKitMap"); diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java index 90ac89d140e1..48275da4878d 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java @@ -20,7 +20,9 @@ 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.guice.annotations.Json; import org.apache.druid.frame.key.KeyColumn; import org.apache.druid.frame.key.KeyOrder; import org.apache.druid.java.util.common.ISE; @@ -56,12 +58,11 @@ public class WindowOperatorQueryKit implements QueryKit { 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, diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java index 7f9c6e623e93..4199933e500a 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java @@ -21,6 +21,8 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; +import com.google.inject.Inject; +import org.apache.druid.guice.annotations.Json; import org.apache.druid.frame.key.ClusterBy; import org.apache.druid.frame.key.KeyColumn; import org.apache.druid.frame.key.KeyOrder; @@ -58,7 +60,8 @@ public class GroupByQueryKit implements QueryKit { private final ObjectMapper jsonMapper; - public GroupByQueryKit(ObjectMapper jsonMapper) + @Inject + public GroupByQueryKit(@Json final ObjectMapper jsonMapper) { this.jsonMapper = jsonMapper; } diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java index 6382d63f7cd3..80c33f8ce262 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java @@ -20,7 +20,9 @@ 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.guice.annotations.Json; import org.apache.druid.frame.key.KeyColumn; import org.apache.druid.frame.key.KeyOrder; import org.apache.druid.java.util.common.granularity.Granularity; @@ -50,7 +52,8 @@ public class ScanQueryKit implements QueryKit { private final ObjectMapper jsonMapper; - public ScanQueryKit(final ObjectMapper jsonMapper) + @Inject + public ScanQueryKit(@Json final ObjectMapper jsonMapper) { this.jsonMapper = jsonMapper; } diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 48d4cab9450f..64335b3a76c3 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -52,6 +52,7 @@ import org.apache.druid.msq.indexing.report.MSQStatusReport; import org.apache.druid.msq.indexing.report.MSQTaskReport; import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig; +import org.apache.druid.msq.querykit.MultiQueryKit; import org.apache.druid.msq.sql.DartQueryKitSpecFactory; import org.apache.druid.msq.test.MSQTestBase; import org.apache.druid.msq.test.MSQTestControllerContext; @@ -258,6 +259,7 @@ public void register(ControllerHolder holder) StringUtils.encodeForFormat(getClass().getSimpleName() + "-controller-exec") ), new DartQueryKitSpecFactory(new TestTimelineServerView(Collections.emptyList())), + injector.getInstance(MultiQueryKit.class), new ServerConfig(), new DefaultQueryConfig(ImmutableMap.of("foo", "bar")), toolbox, diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index 8e6e38bb39b7..bc1af37f1e66 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -561,6 +561,7 @@ public String getFormatString() } ), new ExpressionModule(), + new MSQIndexingModule(), new MSQExternalDataSourceModule(), new LookylooModule(), new SegmentWranglerModule(), From 24df5988c224f76d83d3af62f2f753ca45cc593b Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Mon, 29 Dec 2025 16:32:13 -0800 Subject: [PATCH 2/3] Fix imports. --- .../druid/msq/dart/controller/sql/DartQueryMaker.java | 2 +- .../org/apache/druid/msq/guice/MSQIndexingModule.java | 10 +++++----- .../druid/msq/querykit/WindowOperatorQueryKit.java | 2 +- .../druid/msq/querykit/groupby/GroupByQueryKit.java | 2 +- .../apache/druid/msq/querykit/scan/ScanQueryKit.java | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index 519179e4dddb..c309ef0d419d 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -41,7 +41,6 @@ import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.ControllerImpl; import org.apache.druid.msq.exec.QueryKitSpecFactory; -import org.apache.druid.msq.querykit.MultiQueryKit; import org.apache.druid.msq.exec.QueryListener; import org.apache.druid.msq.exec.ResultsContext; import org.apache.druid.msq.indexing.LegacyMSQSpec; @@ -54,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; diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java index e9ca68443844..0c2680bf416e 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java @@ -92,19 +92,19 @@ 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.groupby.GroupByQueryKit; -import org.apache.druid.msq.querykit.scan.ScanQueryKit; -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.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; diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java index 48275da4878d..94ef7ac4cbf9 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java @@ -22,9 +22,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; import org.apache.druid.frame.key.ClusterBy; -import org.apache.druid.guice.annotations.Json; 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; diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java index 4199933e500a..6dacc20c6ce2 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java @@ -22,10 +22,10 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; import com.google.inject.Inject; -import org.apache.druid.guice.annotations.Json; 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; diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java index 80c33f8ce262..95f3d200b332 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java @@ -22,9 +22,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Inject; import org.apache.druid.frame.key.ClusterBy; -import org.apache.druid.guice.annotations.Json; 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; From 9e3b297db8228a023e2c1329eb32dc4bbd01e319 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Mon, 29 Dec 2025 17:05:27 -0800 Subject: [PATCH 3/3] Fix MSQTaskQueryMakerTest. --- .../java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java index 4005f08ef2e2..3e2c3b0de96c 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java @@ -46,6 +46,7 @@ import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.msq.exec.DataServerQueryHandlerFactory; +import org.apache.druid.msq.guice.MSQIndexingModule; import org.apache.druid.msq.indexing.destination.MSQTerminalStageSpecFactory; import org.apache.druid.msq.indexing.destination.SegmentGenerationTerminalStageSpecFactory; import org.apache.druid.msq.indexing.error.MSQErrorReport; @@ -219,7 +220,8 @@ public void setUp() throws Exception new LifecycleModule(), new ConfigModule(), new SegmentWranglerModule(), - new LookylooModule() + new LookylooModule(), + new MSQIndexingModule() ); Injector injector = Guice.createInjector(defaultModule, BoundFieldModule.of(this)); DruidSecondaryModule.setupJackson(injector, objectMapper);