Examples with ConnectorPlanOptimizer io.prestosql.spi.ConnectorPlanOptimizer used on opensource projects

Example 1 with ConnectorPlanOptimizer

use of io.prestosql.spi.ConnectorPlanOptimizer in project hetu-core by openlookeng.

the class ApplyConnectorOptimization method optimize.

@Override
public PlanNode optimize(PlanNode inputPlan, Session session, TypeProvider types, PlanSymbolAllocator planSymbolAllocator, PlanNodeIdAllocator idAllocator, WarningCollector warningCollector) {
    requireNonNull(inputPlan, "inputPlan is null");
    requireNonNull(session, "session is null");
    requireNonNull(types, "types is null");
    requireNonNull(idAllocator, "idAllocator is null");
    PlanNode plan = inputPlan;
    Map<CatalogName, Set<ConnectorPlanOptimizer>> connectorOptimizers = connectorOptimizersSupplier.get();
    if (connectorOptimizers.isEmpty()) {
        return plan;
    }
    // retrieve all the connectors
    ImmutableSet.Builder<CatalogName> catalogNames = ImmutableSet.builder();
    getAllCatalogNames(plan, catalogNames);
    for (CatalogName catalogName : catalogNames.build()) {
        Set<ConnectorPlanOptimizer> optimizers = connectorOptimizers.get(catalogName);
        if (optimizers == null) {
            continue;
        }
        ImmutableMap.Builder<PlanNode, ConnectorPlanNodeContext> contextMapBuilder = ImmutableMap.builder();
        buildConnectorPlanContext(plan, null, contextMapBuilder);
        Map<PlanNode, ConnectorPlanNodeContext> contextMap = contextMapBuilder.build();
        Map<PlanNode, PlanNode> updates = new HashMap<>();
        Map<String, Type> typeMap = new HashMap<>();
        for (Map.Entry<Symbol, Type> entry : types.allTypes().entrySet()) {
            typeMap.put(entry.getKey().getName(), entry.getValue());
        }
        for (PlanNode node : contextMap.keySet()) {
            ConnectorPlanNodeContext context = contextMap.get(node);
            if (!context.isClosure(catalogName) || !context.getParent().isPresent() || contextMap.get(context.getParent().get()).isClosure(catalogName)) {
                continue;
            }
            PlanNode newNode = node;
            for (ConnectorPlanOptimizer optimizer : optimizers) {
                newNode = optimizer.optimize(newNode, session.toConnectorSession(catalogName), typeMap, planSymbolAllocator, idAllocator);
            }
            if (node != newNode) {
                checkState(containsAll(ImmutableSet.copyOf(newNode.getOutputSymbols()), node.getOutputSymbols()), "the connector optimizer from %s returns a node that does not cover all output before optimization", catalogName);
                updates.put(node, newNode);
            }
        }
        Queue<PlanNode> originalNodes = new LinkedList<>(updates.keySet());
        while (!originalNodes.isEmpty()) {
            PlanNode originalNode = originalNodes.poll();
            if (!contextMap.get(originalNode).getParent().isPresent()) {
                plan = updates.get(originalNode);
                continue;
            }
            PlanNode originalParent = contextMap.get(originalNode).getParent().get();
            ImmutableList.Builder<PlanNode> newChildren = ImmutableList.builder();
            originalParent.getSources().forEach(child -> newChildren.add(updates.getOrDefault(child, child)));
            PlanNode newParent = originalParent.replaceChildren(newChildren.build());
            updates.put(originalParent, newParent);
            originalNodes.add(originalParent);
        }
    }
    return plan;
}