This is an automated email from the ASF dual-hosted git repository.
gianm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 89d0e9b5e84 feat: Pluggable DataSourcePlanners for MSQ. (#20253)
89d0e9b5e84 is described below
commit 89d0e9b5e84942c0dfe00a164719bea956d008c3
Author: Gian Merlino <[email protected]>
AuthorDate: Fri Sep 4 16:06:00 2026 -0700
feat: Pluggable DataSourcePlanners for MSQ. (#20253)
This patch adds a DataSourcePlanner interface that can be used to
plug in DataSourcePlan logic, and migrates DataSourcePlan to use
this interface rather than hardcoding all the logic.
---
.../org/apache/druid/msq/guice/MSQBinders.java | 19 +
.../apache/druid/msq/guice/MSQIndexingModule.java | 3 +
.../apache/druid/msq/querykit/DataSourcePlan.java | 644 +--------------------
.../druid/msq/querykit/DataSourcePlanner.java | 53 ++
.../druid/msq/querykit/DataSourcePlanners.java | 91 +++
.../apache/druid/msq/querykit/QueryKitSpec.java | 21 +-
.../datasource/DataSourcePlannerUtils.java | 86 +++
.../datasource/ExternalDataSourcePlanner.java | 64 ++
.../datasource/FilteredDataSourcePlanner.java | 67 +++
.../datasource/InlineDataSourcePlanner.java | 58 ++
.../querykit/datasource/JoinDataSourcePlanner.java | 322 +++++++++++
.../datasource/LookupDataSourcePlanner.java | 56 ++
.../datasource/QueryDataSourcePlanner.java | 85 +++
.../datasource/RestrictedDataSourcePlanner.java | 58 ++
.../datasource/TableDataSourcePlanner.java | 70 +++
.../datasource/UnionDataSourcePlanner.java | 86 +++
.../datasource/UnnestDataSourcePlanner.java | 77 +++
.../druid/msq/sql/DartQueryKitSpecFactory.java | 12 +-
.../druid/msq/sql/MSQTaskQueryKitSpecFactory.java | 12 +-
.../dart/controller/http/DartSqlResourceTest.java | 3 +-
.../druid/msq/querykit/DataSourcePlannerTest.java | 168 ++++++
21 files changed, 1419 insertions(+), 636 deletions(-)
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
index e3cd3e4028e..12841662094 100644
--- 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
@@ -29,7 +29,9 @@ import org.apache.druid.msq.dart.Dart;
import org.apache.druid.msq.input.InputSlice;
import org.apache.druid.msq.input.InputSliceReaderProvider;
import org.apache.druid.msq.input.InputSpecSlicerProvider;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
import org.apache.druid.msq.querykit.QueryKit;
+import org.apache.druid.query.DataSource;
import org.apache.druid.query.Query;
import java.lang.annotation.Annotation;
@@ -60,6 +62,23 @@ public class MSQBinders
);
}
+ /**
+ * Creates a MapBinder for {@link DataSourcePlanner} implementations, keyed
on the exact {@link DataSource} class
+ * they plan. Extensions can use this to make their own {@link DataSource}
types planable by {@link QueryKit}.
+ *
+ * Example usage:
+ * <pre>
+ * MSQBinders.dataSourcePlannerBinder(binder)
+ * .addBinding(MyCustomDataSource.class)
+ * .to(MyCustomDataSourcePlanner.class);
+ * </pre>
+ */
+ @SuppressWarnings("rawtypes")
+ public static MapBinder<Class<? extends DataSource>, DataSourcePlanner>
dataSourcePlannerBinder(Binder binder)
+ {
+ return MapBinder.newMapBinder(binder, new TypeLiteral<>() {}, new
TypeLiteral<>() {});
+ }
+
/**
* Bind an {@link InputSpecSlicerProvider} for use on a controller. The
annotation should be
* {@link IndexingService} for providers used by tasks, or {@link Dart} for
providers used by Dart.
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 8781fa5bb0f..8e472e6e7d8 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
@@ -254,6 +254,9 @@ public class MSQIndexingModule implements DruidModule
.to(WindowOperatorQueryKit.class);
binder.bind(WindowOperatorQueryKit.class).in(LazySingleton.class);
+ // Create the binder no matter what, to ensure we can at least get an
empty map.
+ MSQBinders.dataSourcePlannerBinder(binder);
+
MSQBinders.inputSpecSlicerProviderBinder(binder, IndexingService.class)
.addBinding()
.to(IndexerTableInputSpecSlicerProvider.class)
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java
index d3d4d86405b..c782a60c1a0 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java
@@ -20,81 +20,28 @@
package org.apache.druid.msq.querykit;
import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
-import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
-import it.unimi.dsi.fastutil.ints.IntSets;
-import org.apache.druid.frame.key.ClusterBy;
-import org.apache.druid.frame.key.KeyColumn;
import org.apache.druid.java.util.common.IAE;
-import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.UOE;
-import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.msq.exec.Limits;
import org.apache.druid.msq.input.InputSpec;
-import org.apache.druid.msq.input.external.ExternalInputSpec;
-import org.apache.druid.msq.input.inline.InlineInputSpec;
-import org.apache.druid.msq.input.lookup.LookupInputSpec;
import org.apache.druid.msq.input.stage.StageInputSpec;
-import org.apache.druid.msq.input.table.TableInputSpec;
-import org.apache.druid.msq.kernel.HashShuffleSpec;
-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.kernel.StageDefinitionBuilder;
-import org.apache.druid.msq.querykit.common.SortMergeJoinStageProcessor;
import org.apache.druid.query.DataSource;
-import org.apache.druid.query.FilteredDataSource;
-import org.apache.druid.query.InlineDataSource;
-import org.apache.druid.query.JoinAlgorithm;
-import org.apache.druid.query.JoinDataSource;
-import org.apache.druid.query.LookupDataSource;
import org.apache.druid.query.QueryContext;
-import org.apache.druid.query.QueryDataSource;
-import org.apache.druid.query.RestrictedDataSource;
-import org.apache.druid.query.SegmentDescriptor;
-import org.apache.druid.query.TableDataSource;
-import org.apache.druid.query.UnionDataSource;
-import org.apache.druid.query.UnnestDataSource;
-import org.apache.druid.query.planning.JoinDataSourceAnalysis;
-import org.apache.druid.query.planning.PreJoinableClause;
import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
-import org.apache.druid.query.spec.MultipleSpecificSegmentSpec;
import org.apache.druid.query.spec.QuerySegmentSpec;
-import org.apache.druid.query.spec.SpecificSegmentSpec;
-import org.apache.druid.segment.column.ColumnHolder;
-import org.apache.druid.segment.column.RowSignature;
-import org.apache.druid.segment.join.JoinConditionAnalysis;
-import org.apache.druid.sql.calcite.external.ExternalDataSource;
-import org.apache.druid.sql.calcite.parser.DruidSqlInsert;
-import org.joda.time.Interval;
import javax.annotation.Nullable;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
import java.util.List;
-import java.util.Map;
import java.util.Optional;
-import java.util.stream.Collectors;
/**
* Plan for getting data from a {@link DataSource}. Used by {@link QueryKit}
implementations.
*/
public class DataSourcePlan
{
- /**
- * A map with {@link DruidSqlInsert#SQL_INSERT_SEGMENT_GRANULARITY} set to
null, so we can clear it from the context
- * of subqueries.
- */
- private static final Map<String, Object> CONTEXT_MAP_NO_SEGMENT_GRANULARITY
= new HashMap<>();
- private static final Logger log = new Logger(DataSourcePlan.class);
-
- static {
-
CONTEXT_MAP_NO_SEGMENT_GRANULARITY.put(DruidSqlInsert.SQL_INSERT_SEGMENT_GRANULARITY,
null);
- }
-
private final DataSource newDataSource;
private final List<InputSpec> inputSpecs;
private final IntSet broadcastInputs;
@@ -102,7 +49,7 @@ public class DataSourcePlan
@Nullable
private final QueryDefinitionBuilder subQueryDefBuilder;
- DataSourcePlan(
+ public DataSourcePlan(
final DataSource newDataSource,
final List<InputSpec> inputSpecs,
final IntSet broadcastInputs,
@@ -121,7 +68,7 @@ public class DataSourcePlan
}
}
- DataSourcePlan withDataSource(DataSource newDataSource)
+ public DataSourcePlan withDataSource(DataSource newDataSource)
{
return new DataSourcePlan(newDataSource, inputSpecs, broadcastInputs,
subQueryDefBuilder);
}
@@ -146,96 +93,21 @@ public class DataSourcePlan
final boolean broadcast
)
{
-
- if (dataSource instanceof TableDataSource) {
- return forTable(
- (TableDataSource) dataSource,
- querySegmentSpec,
- broadcast
- );
- } else if (dataSource instanceof RestrictedDataSource) {
- return forRestricted(
- (RestrictedDataSource) dataSource,
- querySegmentSpec,
- broadcast
- );
- } else if (dataSource instanceof ExternalDataSource) {
- checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
- return forExternal((ExternalDataSource) dataSource, broadcast);
- } else if (dataSource instanceof InlineDataSource) {
- checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
- return forInline((InlineDataSource) dataSource, broadcast);
- } else if (dataSource instanceof LookupDataSource) {
- return forLookup((LookupDataSource) dataSource, broadcast);
- } else if (dataSource instanceof FilteredDataSource) {
- return forFilteredDataSource(
- queryKitSpec,
- queryContext,
- (FilteredDataSource) dataSource,
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
- } else if (dataSource instanceof UnnestDataSource) {
- return forUnnest(
- queryKitSpec,
- queryContext,
- (UnnestDataSource) dataSource,
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
- } else if (dataSource instanceof QueryDataSource) {
- checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
- return forQuery(
- queryKitSpec,
- (QueryDataSource) dataSource,
- minStageNumber,
- broadcast
- );
- } else if (dataSource instanceof UnionDataSource) {
- return forUnion(
- queryKitSpec,
- queryContext,
- (UnionDataSource) dataSource,
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
- } else if (dataSource instanceof JoinDataSource) {
- JoinDataSource joinDataSource = (JoinDataSource) dataSource;
- final JoinAlgorithm preferredJoinAlgorithm =
joinDataSource.getJoinAlgorithm();
- final JoinAlgorithm deducedJoinAlgorithm = deduceJoinAlgorithm(
- preferredJoinAlgorithm,
- joinDataSource
- );
-
- switch (deducedJoinAlgorithm) {
- case BROADCAST:
- return forBroadcastHashJoin(
- queryKitSpec,
- queryContext,
- joinDataSource,
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
-
- case SORT_MERGE:
- return forSortMergeJoin(
- queryKitSpec,
- joinDataSource,
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
-
- default:
- throw new UOE("Cannot handle join algorithm [%s]",
deducedJoinAlgorithm);
- }
- } else {
+ //noinspection rawtypes
+ final DataSourcePlanner planner =
queryKitSpec.getDataSourcePlanners().getPlanner(dataSource.getClass());
+ if (planner == null) {
throw new UOE("Cannot handle dataSource [%s]", dataSource);
}
+
+ //noinspection unchecked
+ return planner.planDataSource(
+ queryKitSpec,
+ queryContext,
+ dataSource,
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ );
}
/**
@@ -286,50 +158,6 @@ public class DataSourcePlan
return Optional.ofNullable(subQueryDefBuilder);
}
- /**
- * Contains the logic that deduces the join algorithm to be used. Ideally,
this should reside while planning the
- * native query, however we don't have the resources and the structure in
place (when adding this function) to do so.
- * Therefore, this is done while planning the MSQ query
- * It takes into account the algorithm specified by "sqlJoinAlgorithm" in
the query context and the join condition
- * that is present in the query.
- */
- private static JoinAlgorithm deduceJoinAlgorithm(JoinAlgorithm
preferredJoinAlgorithm, JoinDataSource joinDataSource)
- {
- JoinAlgorithm deducedJoinAlgorithm;
- if (JoinAlgorithm.BROADCAST.equals(preferredJoinAlgorithm)) {
- deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
- } else if (canUseSortMergeJoin(joinDataSource.getConditionAnalysis())) {
- deducedJoinAlgorithm = JoinAlgorithm.SORT_MERGE;
- } else {
- deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
- }
-
- if (deducedJoinAlgorithm != preferredJoinAlgorithm) {
- log.debug(
- "User wanted to plan join [%s] as [%s], however the join will be
executed as [%s]",
- joinDataSource,
- preferredJoinAlgorithm.toString(),
- deducedJoinAlgorithm.toString()
- );
- }
-
- return deducedJoinAlgorithm;
- }
-
- /**
- * Checks if the sortMerge algorithm can execute a particular join condition.
- * <p>
- * One check: join condition on two tables "table1" and "table2" is of the
form
- * table1.columnA = table2.columnA && table1.columnB = table2.columnB && ....
- */
- private static boolean canUseSortMergeJoin(JoinConditionAnalysis
joinConditionAnalysis)
- {
- return joinConditionAnalysis
- .getEquiConditions()
- .stream()
- .allMatch(equality -> equality.getLeftExpr().isIdentifier());
- }
-
/**
* Whether this datasource must be processed by a single worker. True if,
and only if, all inputs are broadcast.
*/
@@ -337,446 +165,4 @@ public class DataSourcePlan
{
return broadcastInputs.size() == inputSpecs.size();
}
-
- private static DataSourcePlan forTable(
- final TableDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final boolean broadcast
- )
- {
- final List<SegmentDescriptor> segments;
- if (querySegmentSpec instanceof MultipleSpecificSegmentSpec) {
- segments = ((MultipleSpecificSegmentSpec)
querySegmentSpec).getDescriptors();
- } else if (querySegmentSpec instanceof SpecificSegmentSpec) {
- segments = List.of(((SpecificSegmentSpec)
querySegmentSpec).getDescriptor());
- } else {
- segments = null;
- }
- List<Interval> intervals = querySegmentSpec.getIntervals();
- return new DataSourcePlan(
- (broadcast && dataSource.isGlobal()) ? dataSource : new
InputNumberDataSource(0),
- List.of(new TableInputSpec(dataSource.getName(), intervals, segments)),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- null
- );
- }
-
- private static DataSourcePlan forRestricted(
- final RestrictedDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final boolean broadcast
- )
- {
- DataSource restricted = (broadcast && dataSource.isGlobal())
- ? dataSource
- : new RestrictedInputNumberDataSource(0,
dataSource.getPolicy());
- return forTable(dataSource.getBase(), querySegmentSpec,
broadcast).withDataSource(restricted);
- }
-
- private static DataSourcePlan forExternal(
- final ExternalDataSource dataSource,
- final boolean broadcast
- )
- {
- return new DataSourcePlan(
- dataSource,
- Collections.singletonList(
- new ExternalInputSpec(
- dataSource.getInputSource(),
- dataSource.getInputFormat(),
- dataSource.getSignature()
- )
- ),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- null
- );
- }
-
- private static DataSourcePlan forInline(
- final InlineDataSource dataSource,
- final boolean broadcast
- )
- {
- return new DataSourcePlan(
- dataSource,
- Collections.singletonList(new InlineInputSpec(dataSource)),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- null
- );
- }
-
- private static DataSourcePlan forLookup(
- final LookupDataSource dataSource,
- final boolean broadcast
- )
- {
- return new DataSourcePlan(
- dataSource,
- Collections.singletonList(new
LookupInputSpec(dataSource.getLookupName())),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- null
- );
- }
-
- private static DataSourcePlan forQuery(
- final QueryKitSpec queryKitSpec,
- final QueryDataSource dataSource,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- final QueryDefinition subQueryDef =
queryKitSpec.getQueryKit().makeQueryDefinition(
- queryKitSpec,
- // Subqueries ignore SQL_INSERT_SEGMENT_GRANULARITY, even if set in
the context. It's only used for the
- // outermost query, and setting it for the subquery makes us
erroneously add bucketing where it doesn't belong.
-
dataSource.getQuery().withOverriddenContext(CONTEXT_MAP_NO_SEGMENT_GRANULARITY),
- ShuffleSpecFactories.globalSortWithTargetPartitions(),
- minStageNumber
- );
-
- final int stageNumber =
subQueryDef.getFinalStageDefinition().getStageNumber();
-
- return new DataSourcePlan(
- new InputNumberDataSource(0),
- Collections.singletonList(new StageInputSpec(stageNumber)),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- QueryDefinition.builder(subQueryDef)
- );
- }
-
- private static DataSourcePlan forFilteredDataSource(
- final QueryKitSpec queryKitSpec,
- final QueryContext queryContext,
- final FilteredDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- final DataSourcePlan basePlan = forDataSource(
- queryKitSpec,
- queryContext,
- dataSource.getBase(),
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
-
- DataSource newDataSource = basePlan.getNewDataSource();
-
- final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
- newDataSource = FilteredDataSource.create(newDataSource,
dataSource.getFilter());
- return new DataSourcePlan(
- newDataSource,
- inputSpecs,
- basePlan.getBroadcastInputs(),
- basePlan.getSubQueryDefBuilder().orElse(null)
- );
-
- }
-
- /**
- * Build a plan for Unnest data source
- */
- private static DataSourcePlan forUnnest(
- final QueryKitSpec queryKitSpec,
- final QueryContext queryContext,
- final UnnestDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- // Find the plan for base data source by recursing
- final DataSourcePlan basePlan = forDataSource(
- queryKitSpec,
- queryContext,
- dataSource.getBase(),
- querySegmentSpec,
- minStageNumber,
- broadcast
- );
- DataSource newDataSource = basePlan.getNewDataSource();
-
- final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
-
- // Create the new data source using the data source from the base plan
- newDataSource = UnnestDataSource.create(
- newDataSource,
- dataSource.getVirtualColumn(),
- dataSource.getUnnestFilter()
- );
- // The base data source can be a join and might already have broadcast
inputs
- // Need to set the broadcast inputs from the basePlan
- return new DataSourcePlan(
- newDataSource,
- inputSpecs,
- basePlan.getBroadcastInputs(),
- basePlan.getSubQueryDefBuilder().orElse(null)
- );
- }
-
- private static DataSourcePlan forUnion(
- final QueryKitSpec queryKitSpec,
- final QueryContext queryContext,
- final UnionDataSource unionDataSource,
- final QuerySegmentSpec querySegmentSpec,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- // This is done to prevent loss of generality since MSQ can plan any type
of DataSource.
- List<DataSource> children = unionDataSource.getChildren();
-
- final QueryDefinitionBuilder subqueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
- final List<DataSource> newChildren = new ArrayList<>();
- final List<InputSpec> inputSpecs = new ArrayList<>();
- final IntSet broadcastInputs = new IntOpenHashSet();
-
- for (DataSource child : children) {
- DataSourcePlan childDataSourcePlan = forDataSource(
- queryKitSpec,
- queryContext,
- child,
- querySegmentSpec,
- Math.max(minStageNumber, subqueryDefBuilder.getNextStageNumber()),
- broadcast
- );
-
- int shift = inputSpecs.size();
-
-
newChildren.add(shiftInputNumbers(childDataSourcePlan.getNewDataSource(),
shift));
- inputSpecs.addAll(childDataSourcePlan.getInputSpecs());
-
childDataSourcePlan.getSubQueryDefBuilder().ifPresent(subqueryDefBuilder::addAll);
- childDataSourcePlan.getBroadcastInputs().forEach(inp ->
broadcastInputs.add(inp + shift));
- }
- return new DataSourcePlan(
- new UnionDataSource(newChildren),
- inputSpecs,
- broadcastInputs,
- subqueryDefBuilder
- );
- }
-
- /**
- * Build a plan for broadcast hash-join.
- */
- private static DataSourcePlan forBroadcastHashJoin(
- final QueryKitSpec queryKitSpec,
- final QueryContext queryContext,
- final JoinDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- final QueryDefinitionBuilder subQueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
- final JoinDataSourceAnalysis analysis =
dataSource.getJoinAnalysisForDataSource();
-
- final DataSourcePlan basePlan = forDataSource(
- queryKitSpec,
- queryContext,
- analysis.getBaseDataSource(),
- querySegmentSpec,
- Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
- broadcast
- );
-
- DataSource newDataSource = basePlan.getNewDataSource();
- final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
- final IntSet broadcastInputs = new
IntOpenHashSet(basePlan.getBroadcastInputs());
- basePlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
-
- for (int i = 0; i < analysis.getPreJoinableClauses().size(); i++) {
- final PreJoinableClause clause = analysis.getPreJoinableClauses().get(i);
- final DataSourcePlan clausePlan = forDataSource(
- queryKitSpec,
- queryContext,
- clause.getDataSource(),
- new MultipleIntervalSegmentSpec(Intervals.ONLY_ETERNITY),
- Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
- true // Always broadcast right-hand side of the join.
- );
-
- // Shift all input numbers in the clausePlan.
- final int shift = inputSpecs.size();
-
- newDataSource = JoinDataSource.create(
- newDataSource,
- shiftInputNumbers(clausePlan.getNewDataSource(), shift),
- clause.getPrefix(),
- clause.getCondition(),
- clause.getJoinType(),
- // First JoinDataSource (i == 0) involves the base table, so we need
to propagate the base table filter.
- i == 0 ? analysis.getJoinBaseTableFilter().orElse(null) : null,
- dataSource.getJoinableFactoryWrapper(),
- clause.getJoinAlgorithm()
- );
- inputSpecs.addAll(clausePlan.getInputSpecs());
- clausePlan.getBroadcastInputs().intStream().forEach(n ->
broadcastInputs.add(n + shift));
- clausePlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
- }
-
- return new DataSourcePlan(newDataSource, inputSpecs, broadcastInputs,
subQueryDefBuilder);
- }
-
- /**
- * Build a plan for sort-merge join.
- */
- private static DataSourcePlan forSortMergeJoin(
- final QueryKitSpec queryKitSpec,
- final JoinDataSource dataSource,
- final QuerySegmentSpec querySegmentSpec,
- final int minStageNumber,
- final boolean broadcast
- )
- {
- checkQuerySegmentSpecIsEternity(dataSource, querySegmentSpec);
-
SortMergeJoinStageProcessor.validateCondition(dataSource.getConditionAnalysis());
-
- // Partition by keys given by the join condition.
- final List<List<KeyColumn>> partitionKeys =
SortMergeJoinStageProcessor.toKeyColumns(
-
SortMergeJoinStageProcessor.validateCondition(dataSource.getConditionAnalysis())
- );
-
- final QueryDefinitionBuilder subQueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
-
- // Plan the left input.
- // We're confident that we can cast dataSource.getLeft() to
QueryDataSource, because DruidJoinQueryRel creates
- // subqueries when the join algorithm is sortMerge.
- final DataSourcePlan leftPlan = forQuery(
- queryKitSpec,
- (QueryDataSource) dataSource.getLeft(),
- Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
- false
- );
- leftPlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
-
- // Plan the right input.
- // We're confident that we can cast dataSource.getRight() to
QueryDataSource, because DruidJoinQueryRel creates
- // subqueries when the join algorithm is sortMerge.
- final DataSourcePlan rightPlan = forQuery(
- queryKitSpec,
- (QueryDataSource) dataSource.getRight(),
- Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
- false
- );
- rightPlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
-
- // Build up the left stage.
- final StageDefinitionBuilder leftBuilder =
subQueryDefBuilder.getStageBuilder(
- ((StageInputSpec)
Iterables.getOnlyElement(leftPlan.getInputSpecs())).getStageNumber()
- );
-
- final List<KeyColumn> leftPartitionKey = partitionKeys.get(0);
- leftBuilder.shuffleSpec(new HashShuffleSpec(new
ClusterBy(leftPartitionKey, 0), 1, true));
-
leftBuilder.signature(QueryKitUtils.sortableSignature(leftBuilder.getSignature(),
leftPartitionKey));
- leftBuilder.maxWorkerCount(Limits.MAX_WORKERS);
-
- // Build up the right stage.
- final StageDefinitionBuilder rightBuilder =
subQueryDefBuilder.getStageBuilder(
- ((StageInputSpec)
Iterables.getOnlyElement(rightPlan.getInputSpecs())).getStageNumber()
- );
-
- final List<KeyColumn> rightPartitionKey = partitionKeys.get(1);
- rightBuilder.shuffleSpec(new HashShuffleSpec(new
ClusterBy(rightPartitionKey, 0), 1, true));
-
rightBuilder.signature(QueryKitUtils.sortableSignature(rightBuilder.getSignature(),
rightPartitionKey));
- rightBuilder.maxWorkerCount(Limits.MAX_WORKERS);
-
- // Compute join signature.
- final RowSignature.Builder joinSignatureBuilder = RowSignature.builder();
-
- for (String leftColumn : leftBuilder.getSignature().getColumnNames()) {
- joinSignatureBuilder.add(leftColumn,
leftBuilder.getSignature().getColumnType(leftColumn).orElse(null));
- }
-
- for (String rightColumn : rightBuilder.getSignature().getColumnNames()) {
- joinSignatureBuilder.add(
- dataSource.getRightPrefix() + rightColumn,
- rightBuilder.getSignature().getColumnType(rightColumn).orElse(null)
- );
- }
-
- // Build up the join stage.
- final int stageNumber = Math.max(minStageNumber,
subQueryDefBuilder.getNextStageNumber());
-
- subQueryDefBuilder.add(
- StageDefinition.builder(stageNumber)
- .inputs(
- ImmutableList.of(
-
Iterables.getOnlyElement(leftPlan.getInputSpecs()),
-
Iterables.getOnlyElement(rightPlan.getInputSpecs())
- )
- )
- .maxWorkerCount(Limits.MAX_WORKERS)
- .signature(joinSignatureBuilder.build())
- .processor(
- new SortMergeJoinStageProcessor(
- dataSource.getRightPrefix(),
- dataSource.getConditionAnalysis(),
- dataSource.getJoinType()
- )
- )
- );
-
- return new DataSourcePlan(
- new InputNumberDataSource(0),
- Collections.singletonList(new StageInputSpec(stageNumber)),
- broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
- subQueryDefBuilder
- );
- }
-
- private static DataSource shiftInputNumbers(final DataSource dataSource,
final int shift)
- {
- if (shift < 0) {
- throw new IAE("Shift must be >= 0");
- } else if (shift == 0) {
- return dataSource;
- } else {
- if (dataSource instanceof InputNumberDataSource) {
- return new InputNumberDataSource(((InputNumberDataSource)
dataSource).getInputNumber() + shift);
- } else {
- return dataSource.withChildren(
- dataSource.getChildren()
- .stream()
- .map(child -> shiftInputNumbers(child, shift))
- .collect(Collectors.toList())
- );
- }
- }
- }
-
- private static List<Interval> querySegmentSpecIntervals(final
QuerySegmentSpec querySegmentSpec)
- {
- if (querySegmentSpec instanceof MultipleIntervalSegmentSpec) {
- return querySegmentSpec.getIntervals();
- } else {
- throw new UOE("Cannot handle querySegmentSpec type [%s]",
querySegmentSpec.getClass().getName());
- }
- }
-
- /**
- * Verify that the provided {@link QuerySegmentSpec} is a {@link
MultipleIntervalSegmentSpec} with
- * interval {@link Intervals#ETERNITY}. If not, throw an {@link
UnsupportedOperationException}.
- * <p>
- * See {@link
org.apache.druid.sql.calcite.rel.DruidQuery#canUseIntervalFiltering(DataSource)}.
- */
- private static void checkQuerySegmentSpecIsEternity(
- final DataSource dataSource,
- final QuerySegmentSpec querySegmentSpec
- )
- {
- final boolean querySegmentSpecIsEternity =
- querySegmentSpec instanceof MultipleIntervalSegmentSpec
- && querySegmentSpec.getIntervals().equals(Intervals.ONLY_ETERNITY);
-
- if (!querySegmentSpecIsEternity) {
- throw new UOE(
- "Cannot filter datasource [%s] using [%s]",
- dataSource.getClass().getName(),
- ColumnHolder.TIME_COLUMN_NAME
- );
- }
- }
}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanner.java
new file mode 100644
index 00000000000..c907e1160ae
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanner.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.querykit;
+
+import com.google.inject.Binder;
+import org.apache.druid.msq.guice.MSQBinders;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+/**
+ * Builds a {@link DataSourcePlan} for a particular class of {@link
DataSource}. Register with
+ * {@link MSQBinders#dataSourcePlannerBinder(Binder)}.
+ */
+public interface DataSourcePlanner<T extends DataSource>
+{
+ /**
+ * Same contract as {@link DataSourcePlan#forDataSource}, for the one
datasource type this planner is registered
+ * against.
+ *
+ * @param queryKitSpec reference for recursive planning
+ * @param queryContext query context
+ * @param dataSource datasource to plan
+ * @param querySegmentSpec intervals for mandatory pruning. The returned
plan must be filtered to this interval.
+ * @param minStageNumber starting stage number for subqueries
+ * @param broadcast whether the plan should broadcast data for this
datasource
+ */
+ DataSourcePlan planDataSource(
+ QueryKitSpec queryKitSpec,
+ QueryContext queryContext,
+ T dataSource,
+ QuerySegmentSpec querySegmentSpec,
+ int minStageNumber,
+ boolean broadcast
+ );
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanners.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanners.java
new file mode 100644
index 00000000000..39fcd44ea3c
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlanners.java
@@ -0,0 +1,91 @@
+/*
+ * 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.querykit;
+
+import com.google.inject.Inject;
+import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.msq.querykit.datasource.ExternalDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.FilteredDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.InlineDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.JoinDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.LookupDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.QueryDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.RestrictedDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.TableDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.UnionDataSourcePlanner;
+import org.apache.druid.msq.querykit.datasource.UnnestDataSourcePlanner;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.FilteredDataSource;
+import org.apache.druid.query.GlobalTableDataSource;
+import org.apache.druid.query.InlineDataSource;
+import org.apache.druid.query.JoinDataSource;
+import org.apache.druid.query.LookupDataSource;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.RestrictedDataSource;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.UnionDataSource;
+import org.apache.druid.query.UnnestDataSource;
+import org.apache.druid.sql.calcite.external.ExternalDataSource;
+
+import java.util.Map;
+
+/**
+ * Provider of {@link DataSourcePlanner}.
+ */
+@LazySingleton
+@SuppressWarnings("rawtypes")
+public class DataSourcePlanners
+{
+ private static final TableDataSourcePlanner TABLE_PLANNER = new
TableDataSourcePlanner();
+
+ /**
+ * Planners for builtin {@link DataSource} types.
+ */
+ private static final Map<Class<? extends DataSource>, DataSourcePlanner>
BUILTIN =
+ Map.ofEntries(
+ Map.entry(TableDataSource.class, TABLE_PLANNER),
+ Map.entry(GlobalTableDataSource.class, TABLE_PLANNER),
+ Map.entry(RestrictedDataSource.class, new
RestrictedDataSourcePlanner()),
+ Map.entry(ExternalDataSource.class, new ExternalDataSourcePlanner()),
+ Map.entry(InlineDataSource.class, new InlineDataSourcePlanner()),
+ Map.entry(LookupDataSource.class, new LookupDataSourcePlanner()),
+ Map.entry(FilteredDataSource.class, new FilteredDataSourcePlanner()),
+ Map.entry(UnnestDataSource.class, new UnnestDataSourcePlanner()),
+ Map.entry(QueryDataSource.class, new QueryDataSourcePlanner()),
+ Map.entry(UnionDataSource.class, new UnionDataSourcePlanner()),
+ Map.entry(JoinDataSource.class, new JoinDataSourcePlanner())
+ );
+
+ private final Map<Class<? extends DataSource>, DataSourcePlanner> planners;
+
+ @Inject
+ public DataSourcePlanners(Map<Class<? extends DataSource>,
DataSourcePlanner> planners)
+ {
+ this.planners = planners;
+ }
+
+ @SuppressWarnings("unchecked")
+ public <T extends DataSource> DataSourcePlanner<T> getPlanner(Class<T>
dataSourceClass)
+ {
+ // Check extension planners first, so extensions can override builtin
planners.
+ final DataSourcePlanner extensionPlanner = planners.get(dataSourceClass);
+ return extensionPlanner != null ? extensionPlanner :
BUILTIN.get(dataSourceClass);
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitSpec.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitSpec.java
index 4454026fca9..6ace9c4c990 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitSpec.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitSpec.java
@@ -20,27 +20,32 @@
package org.apache.druid.msq.querykit;
import org.apache.druid.msq.kernel.QueryDefinition;
+import org.apache.druid.query.DataSource;
import org.apache.druid.query.Query;
/**
- * Container for {@link QueryKit} plus the queryId that we want to build.
+ * Container for {@link QueryKit}, {@link DataSourcePlanners}, and the queryId
that we want to build.
*/
public class QueryKitSpec
{
private final QueryKit<Query<?>> queryKit;
private final String queryId;
+ private final DataSourcePlanners dataSourcePlanners;
/**
- * @param queryKit kit that is used to translate native
subqueries; i.e.,
- * {@link
org.apache.druid.query.QueryDataSource}. Typically a {@link MultiQueryKit}.
- * @param queryId queryId of the resulting {@link
QueryDefinition}
+ * @param queryKit kit that is used to translate native
subqueries; i.e.,
+ * {@link org.apache.druid.query.QueryDataSource}.
Typically a {@link MultiQueryKit}.
+ * @param dataSourcePlanners planners for datasource types that {@link
DataSourcePlan} does not handle itself
+ * @param queryId queryId of the resulting {@link QueryDefinition}
*/
public QueryKitSpec(
QueryKit<Query<?>> queryKit,
+ DataSourcePlanners dataSourcePlanners,
String queryId
)
{
this.queryId = queryId;
+ this.dataSourcePlanners = dataSourcePlanners;
this.queryKit = queryKit;
}
@@ -52,6 +57,14 @@ public class QueryKitSpec
return queryKit;
}
+ /**
+ * Registered {@link DataSourcePlanner}, keyed on the exact {@link
DataSource} class each one plans.
+ */
+ public DataSourcePlanners getDataSourcePlanners()
+ {
+ return dataSourcePlanners;
+ }
+
/**
* Query ID to use when building {@link QueryDefinition}.
*/
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/DataSourcePlannerUtils.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/DataSourcePlannerUtils.java
new file mode 100644
index 00000000000..59e05a5e1af
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/DataSourcePlannerUtils.java
@@ -0,0 +1,86 @@
+/*
+ * 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.querykit.datasource;
+
+import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.UOE;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.segment.column.ColumnHolder;
+
+import java.util.stream.Collectors;
+
+/**
+ * Utilities shared by {@link DataSourcePlanner} implementations.
+ */
+public class DataSourcePlannerUtils
+{
+ /**
+ * Shift every {@link InputNumberDataSource} in the provided datasource tree
upwards by the given amount. Used when
+ * merging the inputs of multiple child plans into a single plan.
+ */
+ public static DataSource shiftInputNumbers(final DataSource dataSource,
final int shift)
+ {
+ if (shift < 0) {
+ throw new IAE("Shift must be >= 0");
+ } else if (shift == 0) {
+ return dataSource;
+ } else {
+ if (dataSource instanceof InputNumberDataSource) {
+ return new InputNumberDataSource(((InputNumberDataSource)
dataSource).getInputNumber() + shift);
+ } else {
+ return dataSource.withChildren(
+ dataSource.getChildren()
+ .stream()
+ .map(child -> shiftInputNumbers(child, shift))
+ .collect(Collectors.toList())
+ );
+ }
+ }
+ }
+
+ /**
+ * Verify that the provided {@link QuerySegmentSpec} is a {@link
MultipleIntervalSegmentSpec} with
+ * interval {@link Intervals#ETERNITY}. If not, throw an {@link
UnsupportedOperationException}.
+ * <p>
+ * See {@link
org.apache.druid.sql.calcite.rel.DruidQuery#canUseIntervalFiltering(DataSource)}.
+ */
+ public static void checkQuerySegmentSpecIsEternity(
+ final DataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec
+ )
+ {
+ final boolean querySegmentSpecIsEternity =
+ querySegmentSpec instanceof MultipleIntervalSegmentSpec
+ && querySegmentSpec.getIntervals().equals(Intervals.ONLY_ETERNITY);
+
+ if (!querySegmentSpecIsEternity) {
+ throw new UOE(
+ "Cannot filter datasource [%s] using [%s]",
+ dataSource.getClass().getName(),
+ ColumnHolder.TIME_COLUMN_NAME
+ );
+ }
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/ExternalDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/ExternalDataSourcePlanner.java
new file mode 100644
index 00000000000..ca7c5f39d7b
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/ExternalDataSourcePlanner.java
@@ -0,0 +1,64 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.msq.input.external.ExternalInputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.sql.calcite.external.ExternalDataSource;
+
+import java.util.Collections;
+
+/**
+ * Planner for {@link ExternalDataSource}.
+ */
+public class ExternalDataSourcePlanner implements
DataSourcePlanner<ExternalDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final ExternalDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ DataSourcePlannerUtils.checkQuerySegmentSpecIsEternity(dataSource,
querySegmentSpec);
+
+ return new DataSourcePlan(
+ dataSource,
+ Collections.singletonList(
+ new ExternalInputSpec(
+ dataSource.getInputSource(),
+ dataSource.getInputFormat(),
+ dataSource.getSignature()
+ )
+ ),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ null
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/FilteredDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/FilteredDataSourcePlanner.java
new file mode 100644
index 00000000000..28b9c5ec50c
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/FilteredDataSourcePlanner.java
@@ -0,0 +1,67 @@
+/*
+ * 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.querykit.datasource;
+
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.FilteredDataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Planner for {@link FilteredDataSource}. Plans the base datasource, then
reapplies the filter.
+ */
+public class FilteredDataSourcePlanner implements
DataSourcePlanner<FilteredDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final FilteredDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ final DataSourcePlan basePlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ dataSource.getBase(),
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ );
+
+ final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
+ final DataSource newDataSource =
FilteredDataSource.create(basePlan.getNewDataSource(), dataSource.getFilter());
+ return new DataSourcePlan(
+ newDataSource,
+ inputSpecs,
+ basePlan.getBroadcastInputs(),
+ basePlan.getSubQueryDefBuilder().orElse(null)
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/InlineDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/InlineDataSourcePlanner.java
new file mode 100644
index 00000000000..fc7d26dca70
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/InlineDataSourcePlanner.java
@@ -0,0 +1,58 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.msq.input.inline.InlineInputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.InlineDataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+import java.util.Collections;
+
+/**
+ * Planner for {@link InlineDataSource}.
+ */
+public class InlineDataSourcePlanner implements
DataSourcePlanner<InlineDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final InlineDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ DataSourcePlannerUtils.checkQuerySegmentSpecIsEternity(dataSource,
querySegmentSpec);
+
+ return new DataSourcePlan(
+ dataSource,
+ Collections.singletonList(new InlineInputSpec(dataSource)),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ null
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/JoinDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/JoinDataSourcePlanner.java
new file mode 100644
index 00000000000..40a40fe1235
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/JoinDataSourcePlanner.java
@@ -0,0 +1,322 @@
+/*
+ * 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.querykit.datasource;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.frame.key.ClusterBy;
+import org.apache.druid.frame.key.KeyColumn;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.exec.Limits;
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.input.stage.StageInputSpec;
+import org.apache.druid.msq.kernel.HashShuffleSpec;
+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.kernel.StageDefinitionBuilder;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.msq.querykit.QueryKitUtils;
+import org.apache.druid.msq.querykit.common.SortMergeJoinStageProcessor;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.JoinAlgorithm;
+import org.apache.druid.query.JoinDataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.planning.JoinDataSourceAnalysis;
+import org.apache.druid.query.planning.PreJoinableClause;
+import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.join.JoinConditionAnalysis;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Planner for {@link JoinDataSource}. Dispatches to broadcast hash-join or
sort-merge join, based on
+ * {@link #deduceJoinAlgorithm}.
+ */
+public class JoinDataSourcePlanner implements DataSourcePlanner<JoinDataSource>
+{
+ private static final Logger log = new Logger(JoinDataSourcePlanner.class);
+
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final JoinDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ final JoinAlgorithm preferredJoinAlgorithm = dataSource.getJoinAlgorithm();
+ final JoinAlgorithm deducedJoinAlgorithm =
deduceJoinAlgorithm(preferredJoinAlgorithm, dataSource);
+
+ return switch (deducedJoinAlgorithm) {
+ case BROADCAST -> planBroadcastHashJoin(
+ queryKitSpec,
+ queryContext,
+ dataSource,
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ );
+ case SORT_MERGE -> planSortMergeJoin(
+ queryKitSpec,
+ queryContext,
+ dataSource,
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ );
+ };
+ }
+
+ /**
+ * Build a plan for broadcast hash-join.
+ */
+ private static DataSourcePlan planBroadcastHashJoin(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final JoinDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ final QueryDefinitionBuilder subQueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
+ final JoinDataSourceAnalysis analysis =
dataSource.getJoinAnalysisForDataSource();
+
+ final DataSourcePlan basePlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ analysis.getBaseDataSource(),
+ querySegmentSpec,
+ Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
+ broadcast
+ );
+
+ DataSource newDataSource = basePlan.getNewDataSource();
+ final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
+ final IntSet broadcastInputs = new
IntOpenHashSet(basePlan.getBroadcastInputs());
+ basePlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
+
+ for (int i = 0; i < analysis.getPreJoinableClauses().size(); i++) {
+ final PreJoinableClause clause = analysis.getPreJoinableClauses().get(i);
+ final DataSourcePlan clausePlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ clause.getDataSource(),
+ new MultipleIntervalSegmentSpec(Intervals.ONLY_ETERNITY),
+ Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
+ true // Always broadcast right-hand side of the join.
+ );
+
+ // Shift all input numbers in the clausePlan.
+ final int shift = inputSpecs.size();
+
+ newDataSource = JoinDataSource.create(
+ newDataSource,
+
DataSourcePlannerUtils.shiftInputNumbers(clausePlan.getNewDataSource(), shift),
+ clause.getPrefix(),
+ clause.getCondition(),
+ clause.getJoinType(),
+ // First JoinDataSource (i == 0) involves the base table, so we need
to propagate the base table filter.
+ i == 0 ? analysis.getJoinBaseTableFilter().orElse(null) : null,
+ dataSource.getJoinableFactoryWrapper(),
+ clause.getJoinAlgorithm()
+ );
+ inputSpecs.addAll(clausePlan.getInputSpecs());
+ clausePlan.getBroadcastInputs().intStream().forEach(n ->
broadcastInputs.add(n + shift));
+ clausePlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
+ }
+
+ return new DataSourcePlan(newDataSource, inputSpecs, broadcastInputs,
subQueryDefBuilder);
+ }
+
+ /**
+ * Build a plan for sort-merge join.
+ */
+ private static DataSourcePlan planSortMergeJoin(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final JoinDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ DataSourcePlannerUtils.checkQuerySegmentSpecIsEternity(dataSource,
querySegmentSpec);
+
SortMergeJoinStageProcessor.validateCondition(dataSource.getConditionAnalysis());
+
+ // Partition by keys given by the join condition.
+ final List<List<KeyColumn>> partitionKeys =
SortMergeJoinStageProcessor.toKeyColumns(
+
SortMergeJoinStageProcessor.validateCondition(dataSource.getConditionAnalysis())
+ );
+
+ final QueryDefinitionBuilder subQueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
+
+ // Plan the left input.
+ // We're confident that we can cast dataSource.getLeft() to
QueryDataSource, because DruidJoinQueryRel creates
+ // subqueries when the join algorithm is sortMerge.
+ final DataSourcePlan leftPlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ (QueryDataSource) dataSource.getLeft(),
+ querySegmentSpec,
+ Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
+ false
+ );
+ leftPlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
+
+ // Plan the right input.
+ // We're confident that we can cast dataSource.getRight() to
QueryDataSource, because DruidJoinQueryRel creates
+ // subqueries when the join algorithm is sortMerge.
+ final DataSourcePlan rightPlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ (QueryDataSource) dataSource.getRight(),
+ querySegmentSpec,
+ Math.max(minStageNumber, subQueryDefBuilder.getNextStageNumber()),
+ false
+ );
+ rightPlan.getSubQueryDefBuilder().ifPresent(subQueryDefBuilder::addAll);
+
+ // Build up the left stage.
+ final StageDefinitionBuilder leftBuilder =
subQueryDefBuilder.getStageBuilder(
+ ((StageInputSpec)
Iterables.getOnlyElement(leftPlan.getInputSpecs())).getStageNumber()
+ );
+
+ final List<KeyColumn> leftPartitionKey = partitionKeys.get(0);
+ leftBuilder.shuffleSpec(new HashShuffleSpec(new
ClusterBy(leftPartitionKey, 0), 1, true));
+
leftBuilder.signature(QueryKitUtils.sortableSignature(leftBuilder.getSignature(),
leftPartitionKey));
+ leftBuilder.maxWorkerCount(Limits.MAX_WORKERS);
+
+ // Build up the right stage.
+ final StageDefinitionBuilder rightBuilder =
subQueryDefBuilder.getStageBuilder(
+ ((StageInputSpec)
Iterables.getOnlyElement(rightPlan.getInputSpecs())).getStageNumber()
+ );
+
+ final List<KeyColumn> rightPartitionKey = partitionKeys.get(1);
+ rightBuilder.shuffleSpec(new HashShuffleSpec(new
ClusterBy(rightPartitionKey, 0), 1, true));
+
rightBuilder.signature(QueryKitUtils.sortableSignature(rightBuilder.getSignature(),
rightPartitionKey));
+ rightBuilder.maxWorkerCount(Limits.MAX_WORKERS);
+
+ // Compute join signature.
+ final RowSignature.Builder joinSignatureBuilder = RowSignature.builder();
+
+ for (final String leftColumn :
leftBuilder.getSignature().getColumnNames()) {
+ joinSignatureBuilder.add(leftColumn,
leftBuilder.getSignature().getColumnType(leftColumn).orElse(null));
+ }
+
+ for (final String rightColumn :
rightBuilder.getSignature().getColumnNames()) {
+ joinSignatureBuilder.add(
+ dataSource.getRightPrefix() + rightColumn,
+ rightBuilder.getSignature().getColumnType(rightColumn).orElse(null)
+ );
+ }
+
+ // Build up the join stage.
+ final int stageNumber = Math.max(minStageNumber,
subQueryDefBuilder.getNextStageNumber());
+
+ subQueryDefBuilder.add(
+ StageDefinition.builder(stageNumber)
+ .inputs(
+ ImmutableList.of(
+
Iterables.getOnlyElement(leftPlan.getInputSpecs()),
+
Iterables.getOnlyElement(rightPlan.getInputSpecs())
+ )
+ )
+ .maxWorkerCount(Limits.MAX_WORKERS)
+ .signature(joinSignatureBuilder.build())
+ .processor(
+ new SortMergeJoinStageProcessor(
+ dataSource.getRightPrefix(),
+ dataSource.getConditionAnalysis(),
+ dataSource.getJoinType()
+ )
+ )
+ );
+
+ return new DataSourcePlan(
+ new InputNumberDataSource(0),
+ Collections.singletonList(new StageInputSpec(stageNumber)),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ subQueryDefBuilder
+ );
+ }
+
+ /**
+ * Contains the logic that deduces the join algorithm to be used. Ideally,
this should reside while planning the
+ * native query, however we don't have the resources and the structure in
place (when adding this function) to do so.
+ * Therefore, this is done while planning the MSQ query
+ * It takes into account the algorithm specified by "sqlJoinAlgorithm" in
the query context and the join condition
+ * that is present in the query.
+ */
+ private static JoinAlgorithm deduceJoinAlgorithm(
+ final JoinAlgorithm preferredJoinAlgorithm,
+ final JoinDataSource joinDataSource
+ )
+ {
+ final JoinAlgorithm deducedJoinAlgorithm;
+ if (JoinAlgorithm.BROADCAST.equals(preferredJoinAlgorithm)) {
+ deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
+ } else if (canUseSortMergeJoin(joinDataSource.getConditionAnalysis())) {
+ deducedJoinAlgorithm = JoinAlgorithm.SORT_MERGE;
+ } else {
+ deducedJoinAlgorithm = JoinAlgorithm.BROADCAST;
+ }
+
+ if (deducedJoinAlgorithm != preferredJoinAlgorithm) {
+ log.debug(
+ "User wanted to plan join [%s] as [%s], however the join will be
executed as [%s]",
+ joinDataSource,
+ preferredJoinAlgorithm.toString(),
+ deducedJoinAlgorithm.toString()
+ );
+ }
+
+ return deducedJoinAlgorithm;
+ }
+
+ /**
+ * Checks if the sortMerge algorithm can execute a particular join condition.
+ * <p>
+ * One check: join condition on two tables "table1" and "table2" is of the
form
+ * table1.columnA = table2.columnA && table1.columnB = table2.columnB && ....
+ */
+ private static boolean canUseSortMergeJoin(final JoinConditionAnalysis
joinConditionAnalysis)
+ {
+ return joinConditionAnalysis
+ .getEquiConditions()
+ .stream()
+ .allMatch(equality -> equality.getLeftExpr().isIdentifier());
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/LookupDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/LookupDataSourcePlanner.java
new file mode 100644
index 00000000000..80a00782310
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/LookupDataSourcePlanner.java
@@ -0,0 +1,56 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.msq.input.lookup.LookupInputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.LookupDataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+import java.util.Collections;
+
+/**
+ * Planner for {@link LookupDataSource}.
+ */
+public class LookupDataSourcePlanner implements
DataSourcePlanner<LookupDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final LookupDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ return new DataSourcePlan(
+ dataSource,
+ Collections.singletonList(new
LookupInputSpec(dataSource.getLookupName())),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ null
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/QueryDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/QueryDataSourcePlanner.java
new file mode 100644
index 00000000000..c33374c61f5
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/QueryDataSourcePlanner.java
@@ -0,0 +1,85 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.msq.input.stage.StageInputSpec;
+import org.apache.druid.msq.kernel.QueryDefinition;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.msq.querykit.ShuffleSpecFactories;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.sql.calcite.parser.DruidSqlInsert;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Planner for {@link QueryDataSource}. Plans the subquery as a separate set
of stages.
+ */
+public class QueryDataSourcePlanner implements
DataSourcePlanner<QueryDataSource>
+{
+ /**
+ * A map with {@link DruidSqlInsert#SQL_INSERT_SEGMENT_GRANULARITY} set to
null, so we can clear it from the context
+ * of subqueries.
+ */
+ private static final Map<String, Object> CONTEXT_MAP_NO_SEGMENT_GRANULARITY
= new HashMap<>();
+
+ static {
+
CONTEXT_MAP_NO_SEGMENT_GRANULARITY.put(DruidSqlInsert.SQL_INSERT_SEGMENT_GRANULARITY,
null);
+ }
+
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final QueryDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ DataSourcePlannerUtils.checkQuerySegmentSpecIsEternity(dataSource,
querySegmentSpec);
+
+ final QueryDefinition subQueryDef =
queryKitSpec.getQueryKit().makeQueryDefinition(
+ queryKitSpec,
+ // Subqueries ignore SQL_INSERT_SEGMENT_GRANULARITY, even if set in
the context. It's only used for the
+ // outermost query, and setting it for the subquery makes us
erroneously add bucketing where it doesn't belong.
+
dataSource.getQuery().withOverriddenContext(CONTEXT_MAP_NO_SEGMENT_GRANULARITY),
+ ShuffleSpecFactories.globalSortWithTargetPartitions(),
+ minStageNumber
+ );
+
+ final int stageNumber =
subQueryDef.getFinalStageDefinition().getStageNumber();
+
+ return new DataSourcePlan(
+ new InputNumberDataSource(0),
+ Collections.singletonList(new StageInputSpec(stageNumber)),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ QueryDefinition.builder(subQueryDef)
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/RestrictedDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/RestrictedDataSourcePlanner.java
new file mode 100644
index 00000000000..a4d2e699389
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/RestrictedDataSourcePlanner.java
@@ -0,0 +1,58 @@
+/*
+ * 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.querykit.datasource;
+
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.msq.querykit.RestrictedInputNumberDataSource;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.RestrictedDataSource;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+/**
+ * Planner for {@link RestrictedDataSource}. Plans the base table, then
reapplies the policy.
+ */
+public class RestrictedDataSourcePlanner implements
DataSourcePlanner<RestrictedDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final RestrictedDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ final DataSource restricted = (broadcast && dataSource.isGlobal())
+ ? dataSource
+ : new RestrictedInputNumberDataSource(0,
dataSource.getPolicy());
+ return DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ dataSource.getBase(),
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ ).withDataSource(restricted);
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/TableDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/TableDataSourcePlanner.java
new file mode 100644
index 00000000000..68d46da8451
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/TableDataSourcePlanner.java
@@ -0,0 +1,70 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.msq.input.table.TableInputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.spec.MultipleSpecificSegmentSpec;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.query.spec.SpecificSegmentSpec;
+import org.joda.time.Interval;
+
+import java.util.List;
+
+/**
+ * Planner for {@link TableDataSource}.
+ */
+public class TableDataSourcePlanner implements
DataSourcePlanner<TableDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final TableDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ final List<SegmentDescriptor> segments;
+ if (querySegmentSpec instanceof MultipleSpecificSegmentSpec) {
+ segments = ((MultipleSpecificSegmentSpec)
querySegmentSpec).getDescriptors();
+ } else if (querySegmentSpec instanceof SpecificSegmentSpec) {
+ segments = List.of(((SpecificSegmentSpec)
querySegmentSpec).getDescriptor());
+ } else {
+ segments = null;
+ }
+ final List<Interval> intervals = querySegmentSpec.getIntervals();
+ return new DataSourcePlan(
+ (broadcast && dataSource.isGlobal()) ? dataSource : new
InputNumberDataSource(0),
+ List.of(new TableInputSpec(dataSource.getName(), intervals, segments)),
+ broadcast ? IntOpenHashSet.of(0) : IntSets.emptySet(),
+ null
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnionDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnionDataSourcePlanner.java
new file mode 100644
index 00000000000..edc1c62d890
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnionDataSourcePlanner.java
@@ -0,0 +1,86 @@
+/*
+ * 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.querykit.datasource;
+
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSet;
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.kernel.QueryDefinition;
+import org.apache.druid.msq.kernel.QueryDefinitionBuilder;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.UnionDataSource;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Planner for {@link UnionDataSource}. Plans each child, then concatenates
their inputs.
+ */
+public class UnionDataSourcePlanner implements
DataSourcePlanner<UnionDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final UnionDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ // This is done to prevent loss of generality since MSQ can plan any type
of DataSource.
+ final List<DataSource> children = dataSource.getChildren();
+
+ final QueryDefinitionBuilder subqueryDefBuilder =
QueryDefinition.builder(queryKitSpec.getQueryId());
+ final List<DataSource> newChildren = new ArrayList<>();
+ final List<InputSpec> inputSpecs = new ArrayList<>();
+ final IntSet broadcastInputs = new IntOpenHashSet();
+
+ for (final DataSource child : children) {
+ final DataSourcePlan childDataSourcePlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ child,
+ querySegmentSpec,
+ Math.max(minStageNumber, subqueryDefBuilder.getNextStageNumber()),
+ broadcast
+ );
+
+ final int shift = inputSpecs.size();
+
+
newChildren.add(DataSourcePlannerUtils.shiftInputNumbers(childDataSourcePlan.getNewDataSource(),
shift));
+ inputSpecs.addAll(childDataSourcePlan.getInputSpecs());
+
childDataSourcePlan.getSubQueryDefBuilder().ifPresent(subqueryDefBuilder::addAll);
+ childDataSourcePlan.getBroadcastInputs().forEach(inp ->
broadcastInputs.add(inp + shift));
+ }
+
+ return new DataSourcePlan(
+ new UnionDataSource(newChildren),
+ inputSpecs,
+ broadcastInputs,
+ subqueryDefBuilder
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnnestDataSourcePlanner.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnnestDataSourcePlanner.java
new file mode 100644
index 00000000000..6288a5401ad
--- /dev/null
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/datasource/UnnestDataSourcePlanner.java
@@ -0,0 +1,77 @@
+/*
+ * 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.querykit.datasource;
+
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.querykit.DataSourcePlan;
+import org.apache.druid.msq.querykit.DataSourcePlanner;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.UnnestDataSource;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Planner for {@link UnnestDataSource}. Plans the base datasource, then
reapplies the unnest.
+ */
+public class UnnestDataSourcePlanner implements
DataSourcePlanner<UnnestDataSource>
+{
+ @Override
+ public DataSourcePlan planDataSource(
+ final QueryKitSpec queryKitSpec,
+ final QueryContext queryContext,
+ final UnnestDataSource dataSource,
+ final QuerySegmentSpec querySegmentSpec,
+ final int minStageNumber,
+ final boolean broadcast
+ )
+ {
+ // Find the plan for base data source by recursing
+ final DataSourcePlan basePlan = DataSourcePlan.forDataSource(
+ queryKitSpec,
+ queryContext,
+ dataSource.getBase(),
+ querySegmentSpec,
+ minStageNumber,
+ broadcast
+ );
+
+ final List<InputSpec> inputSpecs = new
ArrayList<>(basePlan.getInputSpecs());
+
+ // Create the new data source using the data source from the base plan
+ final DataSource newDataSource = UnnestDataSource.create(
+ basePlan.getNewDataSource(),
+ dataSource.getVirtualColumn(),
+ dataSource.getUnnestFilter()
+ );
+
+ // The base data source can be a join and might already have broadcast
inputs
+ // Need to set the broadcast inputs from the basePlan
+ return new DataSourcePlan(
+ newDataSource,
+ inputSpecs,
+ basePlan.getBroadcastInputs(),
+ basePlan.getSubQueryDefBuilder().orElse(null)
+ );
+ }
+}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/DartQueryKitSpecFactory.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/DartQueryKitSpecFactory.java
index 44ade1f09e0..6c4d45a2c1a 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/DartQueryKitSpecFactory.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/DartQueryKitSpecFactory.java
@@ -19,8 +19,10 @@
package org.apache.druid.msq.sql;
+import com.google.inject.Inject;
import org.apache.druid.msq.exec.QueryKitSpecFactory;
import org.apache.druid.msq.indexing.MSQTuningConfig;
+import org.apache.druid.msq.querykit.DataSourcePlanners;
import org.apache.druid.msq.querykit.QueryKit;
import org.apache.druid.msq.querykit.QueryKitSpec;
import org.apache.druid.query.Query;
@@ -28,6 +30,14 @@ import org.apache.druid.query.QueryContext;
public class DartQueryKitSpecFactory implements QueryKitSpecFactory
{
+ private final DataSourcePlanners dataSourcePlanners;
+
+ @Inject
+ public DartQueryKitSpecFactory(final DataSourcePlanners dataSourcePlanners)
+ {
+ this.dataSourcePlanners = dataSourcePlanners;
+ }
+
@Override
public QueryKitSpec makeQueryKitSpec(
final QueryKit<Query<?>> queryKit,
@@ -36,6 +46,6 @@ public class DartQueryKitSpecFactory implements
QueryKitSpecFactory
final QueryContext queryContext
)
{
- return new QueryKitSpec(queryKit, queryId);
+ return new QueryKitSpec(queryKit, dataSourcePlanners, queryId);
}
}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryKitSpecFactory.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryKitSpecFactory.java
index b1fe9c03157..ed72734a7f0 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryKitSpecFactory.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryKitSpecFactory.java
@@ -19,8 +19,10 @@
package org.apache.druid.msq.sql;
+import com.google.inject.Inject;
import org.apache.druid.msq.exec.QueryKitSpecFactory;
import org.apache.druid.msq.indexing.MSQTuningConfig;
+import org.apache.druid.msq.querykit.DataSourcePlanners;
import org.apache.druid.msq.querykit.QueryKit;
import org.apache.druid.msq.querykit.QueryKitSpec;
import org.apache.druid.query.Query;
@@ -28,6 +30,14 @@ import org.apache.druid.query.QueryContext;
public class MSQTaskQueryKitSpecFactory implements QueryKitSpecFactory
{
+ private final DataSourcePlanners dataSourcePlanners;
+
+ @Inject
+ public MSQTaskQueryKitSpecFactory(final DataSourcePlanners
dataSourcePlanners)
+ {
+ this.dataSourcePlanners = dataSourcePlanners;
+ }
+
@Override
public QueryKitSpec makeQueryKitSpec(
QueryKit<Query<?>> queryKit,
@@ -36,6 +46,6 @@ public class MSQTaskQueryKitSpecFactory implements
QueryKitSpecFactory
QueryContext queryContext
)
{
- return new QueryKitSpec(queryKit, queryId);
+ return new QueryKitSpec(queryKit, dataSourcePlanners, queryId);
}
}
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 3f602511d9b..58b4265430b 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
@@ -54,6 +54,7 @@ import org.apache.druid.msq.indexing.error.MSQFaultUtils;
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.DataSourcePlanners;
import org.apache.druid.msq.querykit.MultiQueryKit;
import org.apache.druid.msq.sql.DartQueryKitSpecFactory;
import org.apache.druid.msq.test.MSQTestBase;
@@ -268,7 +269,7 @@ public class DartSqlResourceTest extends MSQTestBase
)
)
),
- new DartQueryKitSpecFactory(),
+ new DartQueryKitSpecFactory(new DataSourcePlanners(Map.of())),
injector.getInstance(MultiQueryKit.class),
new ServerConfig(),
new DefaultQueryConfig(ImmutableMap.of("foo", "bar")),
diff --git
a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/DataSourcePlannerTest.java
b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/DataSourcePlannerTest.java
new file mode 100644
index 00000000000..133637cf067
--- /dev/null
+++
b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/DataSourcePlannerTest.java
@@ -0,0 +1,168 @@
+/*
+ * 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.querykit;
+
+import com.google.common.collect.Iterables;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.UOE;
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.input.inline.InlineInputSpec;
+import org.apache.druid.msq.input.table.TableInputSpec;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.InlineDataSource;
+import org.apache.druid.query.LeafDataSource;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
+import org.apache.druid.query.spec.QuerySegmentSpec;
+import org.apache.druid.segment.column.RowSignature;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class DataSourcePlannerTest
+{
+ private static final QuerySegmentSpec ETERNITY =
+ new MultipleIntervalSegmentSpec(List.of(Intervals.ETERNITY));
+
+ private static final InputSpec PLANNED_SPEC =
+ new InlineInputSpec(InlineDataSource.fromIterable(List.of(),
RowSignature.empty()));
+
+ @Test
+ public void testUnhandledDataSourceThrowsWithNoPlanners()
+ {
+ Assertions.assertThrows(
+ UOE.class,
+ () -> plan(new TestDataSource(), Map.of())
+ );
+ }
+
+ @Test
+ public void testRegisteredPlannerHandlesItsDataSource()
+ {
+ final DataSourcePlan plan = plan(
+ new TestDataSource(),
+ Map.of(TestDataSource.class, new TestDataSourcePlanner<>())
+ );
+
+ Assertions.assertEquals(List.of(PLANNED_SPEC), plan.getInputSpecs());
+ }
+
+ @Test
+ public void testPlannerDoesNotApplyToSubclassOfItsDataSource()
+ {
+ Assertions.assertThrows(
+ UOE.class,
+ () -> plan(new TestDataSourceSubclass(), Map.of(TestDataSource.class,
new TestDataSourcePlanner<>()))
+ );
+ }
+
+ @Test
+ public void testRegisteredPlannerOverridesBuiltin()
+ {
+ final DataSourcePlan plan = plan(
+ new TableDataSource("foo"),
+ Map.of(TableDataSource.class, new TestDataSourcePlanner<>())
+ );
+
+ Assertions.assertEquals(List.of(PLANNED_SPEC), plan.getInputSpecs());
+ }
+
+ @Test
+ public void testBuiltinPlannerHandlesTable()
+ {
+ final DataSourcePlan plan = plan(new TableDataSource("foo"), Map.of());
+
+ Assertions.assertInstanceOf(TableInputSpec.class,
Iterables.getOnlyElement(plan.getInputSpecs()));
+ }
+
+ @SuppressWarnings("rawtypes")
+ private static DataSourcePlan plan(
+ final DataSource dataSource,
+ final Map<Class<? extends DataSource>, DataSourcePlanner> planners
+ )
+ {
+ return DataSourcePlan.forDataSource(
+ new QueryKitSpec(null, new DataSourcePlanners(planners), "queryId"),
+ QueryContext.empty(),
+ dataSource,
+ ETERNITY,
+ 0,
+ false
+ );
+ }
+
+ private static class TestDataSource extends LeafDataSource
+ {
+ @Override
+ public Set<String> getTableNames()
+ {
+ return Set.of();
+ }
+
+ @Override
+ public boolean isCacheable(boolean isBroker)
+ {
+ return false;
+ }
+
+ @Override
+ public boolean isGlobal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean isProcessable()
+ {
+ return true;
+ }
+
+ @Override
+ public byte[] getCacheKey()
+ {
+ return null;
+ }
+ }
+
+ private static class TestDataSourceSubclass extends TestDataSource
+ {
+ }
+
+ private static class TestDataSourcePlanner<T extends DataSource> implements
DataSourcePlanner<T>
+ {
+ @Override
+ public DataSourcePlan planDataSource(
+ QueryKitSpec queryKitSpec,
+ QueryContext queryContext,
+ T dataSource,
+ QuerySegmentSpec querySegmentSpec,
+ int minStageNumber,
+ boolean broadcast
+ )
+ {
+ return new DataSourcePlan(dataSource, List.of(PLANNED_SPEC),
IntSets.emptySet(), null);
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
