feat: Enhance connector optimizer to work on plan with multiple connectors (#26246)

## Description

Previously, connector optimizers in Presto were limited to optimizing
subplans containing tables from only a single connector. For example, a
Hive connector optimizer could only process subplans that exclusively
contained Hive table scans, preventing optimization opportunities for
federated queries that span multiple data sources.

This PR extends the connector optimization framework to support
cross-connector optimizations for federated queries. A new method
`getSupportedConnectorIds()` has been added to the
`ConnectorPlanOptimizer` interface that allows optimizers to declare
which combinations of connectors they can optimize together.

* When `getSupportedConnectorIds()` returns an empty list, the optimizer
maintains the existing behavior - only processing subplans from the
connector it's registered with.
    
* When `getSupportedConnectorIds()` returns a non-empty list, the
optimizer will be applied to subplans that contain table scans from
**exactly** the specified set of connectors. For example, if an
optimizer specifies `[hive, thrift]`, it will: process subplans with
both Hive and Thrift tables, not subplans with only Hive tables, nor
subplans with Hive, Thrift, and MySQL tables

## Motivation and Context
To have connector optimizer to work for sub plans with multiple
connectors involved.

## Impact
To have connector optimizer to work for sub plans with multiple
connectors involved.

## Test Plan
Unit tests

## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== RELEASE NOTES ==

General Changes
* Add new feature to connector optimizer so that it can work for sub plans with multiple connectors
```

Author: feilong-liu

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ApplyConnectorOptimization.java

@@ -54,6 +54,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Queue;
@@ -133,77 +134,97 @@ public PlanOptimizerResult optimize(PlanNode plan, Session session, TypeProvider
         // In order to preserve the fixpoint, we will "pretend" the newly added C2 table scan is part of C1's job to maintain.
         for (ConnectorId connectorId : connectorIds.build()) {
             Set<ConnectorPlanOptimizer> optimizers = connectorOptimizers.get(connectorId);
-            if (optimizers == null) {
+            if (optimizers == null || optimizers.isEmpty()) {
                 continue;
             }
 
+            ImmutableMap.Builder<List<ConnectorId>, Set<ConnectorPlanOptimizer>> optimizersWithConnectorRange = ImmutableMap.builder();
+            List<ConnectorId> currentConnectors = null;
+            ImmutableSet.Builder<ConnectorPlanOptimizer> currentGroup = null;
+            for (ConnectorPlanOptimizer optimizer : optimizers) {
+                List<ConnectorId> supportedConnectors = optimizer.getSupportedConnectorIds().isEmpty()
+                        ? ImmutableList.of(connectorId)
+                        : optimizer.getSupportedConnectorIds();
+
+                if (!supportedConnectors.equals(currentConnectors)) {
+                    if (currentGroup != null) {
+                        optimizersWithConnectorRange.put(currentConnectors, currentGroup.build());
+                    }
+                    currentConnectors = supportedConnectors;
+                    currentGroup = ImmutableSet.builder();
+                }
+                currentGroup.add(optimizer);
+            }
+            optimizersWithConnectorRange.put(currentConnectors, currentGroup.build());
+
             ImmutableMap.Builder<PlanNode, ConnectorPlanNodeContext> contextMapBuilder = ImmutableMap.builder();
             buildConnectorPlanNodeContext(plan, null, contextMapBuilder);
             Map<PlanNode, ConnectorPlanNodeContext> contextMap = contextMapBuilder.build();
+            for (Map.Entry<List<ConnectorId>, Set<ConnectorPlanOptimizer>> entry : optimizersWithConnectorRange.build().entrySet()) {
+                // keep track of changed nodes; the keys are original nodes and the values are the new nodes
+                Map<PlanNode, PlanNode> updates = new HashMap<>();
+
+                // process connector optimizers
+                for (PlanNode node : contextMap.keySet()) {
+                    // For a subtree with root `node` to be a max closure, the following conditions must hold:
+                    //    * The subtree with root `node` is a closure.
+                    //    * `node` has no parent, or the subtree with root as `node`'s parent is not a closure.
+                    ConnectorPlanNodeContext context = contextMap.get(node);
+                    if (!context.isClosure(connectorId, session, entry.getKey()) ||
+                            !context.getParent().isPresent() ||
+                            contextMap.get(context.getParent().get()).isClosure(connectorId, session, entry.getKey())) {
+                        continue;
+                    }
 
-            // keep track of changed nodes; the keys are original nodes and the values are the new nodes
-            Map<PlanNode, PlanNode> updates = new HashMap<>();
-
-            // process connector optimizers
-            for (PlanNode node : contextMap.keySet()) {
-                // For a subtree with root `node` to be a max closure, the following conditions must hold:
-                //    * The subtree with root `node` is a closure.
-                //    * `node` has no parent, or the subtree with root as `node`'s parent is not a closure.
-                ConnectorPlanNodeContext context = contextMap.get(node);
-                if (!context.isClosure(connectorId, session) ||
-                        !context.getParent().isPresent() ||
-                        contextMap.get(context.getParent().get()).isClosure(connectorId, session)) {
-                    continue;
-                }
-
-                PlanNode newNode = node;
+                    PlanNode newNode = node;
 
-                // the returned node is still a max closure (only if there is no new connector added, which does happen but ignored here)
-                for (ConnectorPlanOptimizer optimizer : optimizers) {
-                    long start = System.nanoTime();
-                    newNode = optimizer.optimize(newNode, session.toConnectorSession(connectorId), variableAllocator, idAllocator);
-                    if (enableVerboseRuntimeStats || trackOptimizerRuntime(session, optimizer)) {
-                        session.getRuntimeStats().addMetricValue(String.format("optimizer%sTimeNanos", getOptimizerNameForLog(optimizer)), NANO, System.nanoTime() - start);
+                    // the returned node is still a max closure (only if there is no new connector added, which does happen but ignored here)
+                    for (ConnectorPlanOptimizer optimizer : entry.getValue()) {
+                        long start = System.nanoTime();
+                        newNode = optimizer.optimize(newNode, session.toConnectorSession(connectorId), variableAllocator, idAllocator);
+                        if (enableVerboseRuntimeStats || trackOptimizerRuntime(session, optimizer)) {
+                            session.getRuntimeStats().addMetricValue(String.format("optimizer%sTimeNanos", getOptimizerNameForLog(optimizer)), NANO, System.nanoTime() - start);
+                        }
                     }
-                }
 
-                if (node != newNode) {
-                    // the optimizer has allocated a new PlanNode
-                    checkState(
-                            containsAll(ImmutableSet.copyOf(newNode.getOutputVariables()), node.getOutputVariables()),
-                            "the connector optimizer from %s returns a node that does not cover all output before optimization",
-                            connectorId);
+                    if (node != newNode) {
+                        // the optimizer has allocated a new PlanNode
+                        checkState(
+                                containsAll(ImmutableSet.copyOf(newNode.getOutputVariables()), node.getOutputVariables()),
+                                "the connector optimizer from %s returns a node that does not cover all output before optimization",
+                                connectorId);
 
-                    updates.put(node, newNode);
-                }
-            }
-            // up to this point, we have a set of updated nodes; need to recursively update their parents
-
-            // alter the plan with a bottom-up approach (but does not have to be strict bottom-up to guarantee the correctness of the algorithm)
-            // use "original nodes" to keep track of the plan structure and "updates" to keep track of the new nodes
-            Queue<PlanNode> originalNodes = new LinkedList<>(updates.keySet());
-            while (!originalNodes.isEmpty()) {
-                PlanNode originalNode = originalNodes.poll();
-
-                if (!contextMap.get(originalNode).getParent().isPresent()) {
-                    // originalNode must be the root; update the plan
-                    plan = updates.get(originalNode);
-                    continue;
+                        updates.put(node, newNode);
+                    }
                 }
+                // up to this point, we have a set of updated nodes; need to recursively update their parents
+
+                // alter the plan with a bottom-up approach (but does not have to be strict bottom-up to guarantee the correctness of the algorithm)
+                // use "original nodes" to keep track of the plan structure and "updates" to keep track of the new nodes
+                Queue<PlanNode> originalNodes = new LinkedList<>(updates.keySet());
+                while (!originalNodes.isEmpty()) {
+                    PlanNode originalNode = originalNodes.poll();
+
+                    if (!contextMap.get(originalNode).getParent().isPresent()) {
+                        // originalNode must be the root; update the plan
+                        plan = updates.get(originalNode);
+                        continue;
+                    }
 
-                PlanNode originalParent = contextMap.get(originalNode).getParent().get();
+                    PlanNode originalParent = contextMap.get(originalNode).getParent().get();
 
-                // need to create a new parent given the child has changed; the new parent needs to point to the new child.
-                // if a node has been updated, it will occur in `updates`; otherwise, just use the original node
-                ImmutableList.Builder<PlanNode> newChildren = ImmutableList.builder();
-                originalParent.getSources().forEach(child -> newChildren.add(updates.getOrDefault(child, child)));
-                PlanNode newParent = originalParent.replaceChildren(newChildren.build());
+                    // need to create a new parent given the child has changed; the new parent needs to point to the new child.
+                    // if a node has been updated, it will occur in `updates`; otherwise, just use the original node
+                    ImmutableList.Builder<PlanNode> newChildren = ImmutableList.builder();
+                    originalParent.getSources().forEach(child -> newChildren.add(updates.getOrDefault(child, child)));
+                    PlanNode newParent = originalParent.replaceChildren(newChildren.build());
 
-                // mark the new parent as updated
-                updates.put(originalParent, newParent);
+                    // mark the new parent as updated
+                    updates.put(originalParent, newParent);
 
-                // enqueue the parent node in order to recursively update its ancestors
-                originalNodes.add(originalParent);
+                    // enqueue the parent node in order to recursively update its ancestors
+                    originalNodes.add(originalParent);
+                }
             }
         }
 
@@ -306,17 +327,16 @@ public Set<Class<? extends PlanNode>> getReachablePlanNodeTypes()
             return reachablePlanNodeTypes;
         }
 
-        boolean isClosure(ConnectorId connectorId, Session session)
+        boolean isClosure(ConnectorId connectorId, Session session, List<ConnectorId> supportedConnectorId)
         {
             // check if all children can reach the only connector
             boolean includeValuesNode = isIncludeValuesNodeInConnectorOptimizer(session);
             Set<ConnectorId> connectorIds = includeValuesNode ? reachableConnectors.stream().filter(x -> !x.equals(EMPTY_CONNECTOR_ID)).collect(toImmutableSet()) : reachableConnectors;
-            if (connectorIds.size() != 1 || !connectorIds.contains(connectorId)) {
-                return false;
+            if (connectorIds.contains(connectorId) && new HashSet<>(supportedConnectorId).containsAll(connectorIds) && supportedConnectorId.size() == connectorIds.size()) {
+                // check if all children are accessible by connectors
+                return containsAll(CONNECTOR_ACCESSIBLE_PLAN_NODES, reachablePlanNodeTypes);
             }
-
-            // check if all children are accessible by connectors
-            return containsAll(CONNECTOR_ACCESSIBLE_PLAN_NODES, reachablePlanNodeTypes);
+            return false;
         }
     }