@Override
public PlanNode visitIndexJoin(IndexJoinNode node, RewriteContext<Context> context) {
// Lookup symbols can only be passed through the probe side of an index join
Set<Symbol> probeLookupSymbols =
context
.get()
.getLookupSymbols()
.stream()
.filter(node.getProbeSource().getOutputSymbols()::contains)
.collect(toImmutableSet());
if (probeLookupSymbols.isEmpty()) {
return node;
}
PlanNode rewrittenProbeSource =
context.rewrite(
node.getProbeSource(), new Context(probeLookupSymbols, context.get().getSuccess()));
PlanNode source = node;
if (rewrittenProbeSource != node.getProbeSource()) {
source =
new IndexJoinNode(
node.getId(),
node.getType(),
rewrittenProbeSource,
node.getIndexSource(),
node.getCriteria(),
node.getProbeHashSymbol(),
node.getIndexHashSymbol());
}
return source;
}
@Override
public PlanNode visitExchange(
ExchangeNode exchange, RewriteContext<FragmentProperties> context) {
ImmutableList.Builder<SubPlan> builder = ImmutableList.builder();
if (exchange.getType() == ExchangeNode.Type.GATHER) {
context.get().setSingleNodeDistribution();
for (int i = 0; i < exchange.getSources().size(); i++) {
FragmentProperties childProperties = new FragmentProperties();
childProperties.setUnpartitionedOutput();
childProperties.setOutputLayout(exchange.getInputs().get(i));
builder.add(buildSubPlan(exchange.getSources().get(i), childProperties, context));
}
} else if (exchange.getType() == ExchangeNode.Type.REPARTITION) {
context.get().setFixedDistribution();
FragmentProperties childProperties =
new FragmentProperties()
.setPartitionedOutput(exchange.getPartitionKeys(), exchange.getHashSymbol())
.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
} else if (exchange.getType() == ExchangeNode.Type.REPARTITION_WITH_NULL_REPLICATION) {
context.get().setFixedDistribution();
FragmentProperties childProperties =
new FragmentProperties()
.setPartitionedOutput(exchange.getPartitionKeys(), exchange.getHashSymbol())
.replicateNulls()
.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
} else if (exchange.getType() == ExchangeNode.Type.REPLICATE) {
FragmentProperties childProperties = new FragmentProperties();
childProperties.setUnpartitionedOutput();
childProperties.setOutputLayout(Iterables.getOnlyElement(exchange.getInputs()));
builder.add(
buildSubPlan(
Iterables.getOnlyElement(exchange.getSources()), childProperties, context));
}
List<SubPlan> children = builder.build();
context.get().addChildren(children);
List<PlanFragmentId> childrenIds =
children
.stream()
.map(SubPlan::getFragment)
.map(PlanFragment::getId)
.collect(toImmutableList());
return new RemoteSourceNode(exchange.getId(), childrenIds, exchange.getOutputSymbols());
}
@Override
public PlanNode visitFilter(FilterNode node, RewriteContext<Context> context) {
if (node.getSource() instanceof TableScanNode) {
return planTableScan((TableScanNode) node.getSource(), node.getPredicate(), context.get());
}
return context.defaultRewrite(
node, new Context(context.get().getLookupSymbols(), context.get().getSuccess()));
}
@Override
public PlanNode visitOutput(OutputNode node, RewriteContext<FragmentProperties> context) {
context
.get()
.setSingleNodeDistribution() // TODO: add support for distributed output
.setOutputLayout(node.getOutputSymbols())
.setUnpartitionedOutput();
return context.defaultRewrite(node, context.get());
}
@Override
public PlanNode visitWindow(WindowNode node, RewriteContext<Context> context) {
if (!node.getWindowFunctions()
.values()
.stream()
.map(function -> function.getFunctionCall().getName())
.allMatch(metadata.getFunctionRegistry()::isAggregationFunction)) {
return node;
}
// Don't need this restriction if we can prove that all order by symbols are deterministically
// produced
if (!node.getOrderBy().isEmpty()) {
return node;
}
// Only RANGE frame type currently supported for aggregation functions because it guarantees
// the
// same value for each peer group.
// ROWS frame type requires the ordering to be fully deterministic (e.g. deterministically
// sorted on all columns)
if (node.getFrames()
.stream()
.map(WindowNode.Frame::getType)
.anyMatch(
type ->
type
!= WindowFrame.Type
.RANGE)) { // TODO: extract frames of type RANGE and allow optimization on
// them
return node;
}
// Lookup symbols can only be passed through if they are part of the partitioning
Set<Symbol> partitionByLookupSymbols =
context
.get()
.getLookupSymbols()
.stream()
.filter(node.getPartitionBy()::contains)
.collect(toImmutableSet());
if (partitionByLookupSymbols.isEmpty()) {
return node;
}
return context.defaultRewrite(
node, new Context(partitionByLookupSymbols, context.get().getSuccess()));
}
@Override
public PlanNode visitAggregation(AggregationNode node, RewriteContext<Context> context) {
// Lookup symbols can only be passed through if they are part of the group by columns
Set<Symbol> groupByLookupSymbols =
context
.get()
.getLookupSymbols()
.stream()
.filter(node.getGroupingKeys()::contains)
.collect(toImmutableSet());
if (groupByLookupSymbols.isEmpty()) {
return node;
}
return context.defaultRewrite(
node, new Context(groupByLookupSymbols, context.get().getSuccess()));
}
@Override
public PlanNode visitProject(ProjectNode node, RewriteContext<Context> context) {
// Rewrite the lookup symbols in terms of only the pre-projected symbols that have direct
// translations
Set<Symbol> newLookupSymbols =
context
.get()
.getLookupSymbols()
.stream()
.map(node.getAssignments()::get)
.filter(SymbolReference.class::isInstance)
.map(Symbol::from)
.collect(toImmutableSet());
if (newLookupSymbols.isEmpty()) {
return node;
}
return context.defaultRewrite(
node, new Context(newLookupSymbols, context.get().getSuccess()));
}
@Override
public PlanNode visitSort(SortNode node, RewriteContext<Context> context) {
// Sort has no bearing when building an index, so just ignore the sort
return context.rewrite(node.getSource(), context.get());
}
@Override
public PlanNode visitTableScan(TableScanNode node, RewriteContext<Context> context) {
return planTableScan(node, BooleanLiteral.TRUE_LITERAL, context.get());
}
@Override
public PlanNode visitValues(ValuesNode node, RewriteContext<FragmentProperties> context) {
context.get().setSingleNodeDistribution();
return context.defaultRewrite(node, context.get());
}
@Override
public PlanNode visitTableScan(TableScanNode node, RewriteContext<FragmentProperties> context) {
context.get().setSourceDistribution(node.getId());
return context.defaultRewrite(node, context.get());
}
@Override
public PlanNode visitMetadataDelete(
MetadataDeleteNode node, RewriteContext<FragmentProperties> context) {
context.get().setCoordinatorOnlyDistribution();
return context.defaultRewrite(node, context.get());
}
