Add Util.restrictOutputs()
Introduce a new utility function to make identity projections when
necessary to drop some outputs of a node.  Will be heavily used in the
forthcoming output pruning rules.
package com.facebook.presto.sql.planner.iterative.rule;
import com.facebook.presto.sql.planner.DependencyExtractor;
+import com.facebook.presto.sql.planner.PlanNodeIdAllocator;
import com.facebook.presto.sql.planner.Symbol;
+import com.facebook.presto.sql.planner.plan.Assignments;
import com.facebook.presto.sql.planner.plan.PlanNode;
+import com.facebook.presto.sql.planner.plan.ProjectNode;
import com.facebook.presto.sql.tree.Expression;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Sets;
@@ -26,6 +29,8 @@ import java.util.List;
import java.util.Optional;
import java.util.Set;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
class Util
{
private Util()
@@ -62,4 +67,24 @@ class Util
parent.replaceChildren(
child.getSources())));
}
+
+ /**
+ * @return If the node has outputs not in permittedOutputs, returns an identity projection containing only those node outputs also in permittedOutputs.
+ */
+ public static Optional<PlanNode> restrictOutputs(PlanNodeIdAllocator idAllocator, PlanNode node, Set<Symbol> permittedOutputs)
+ {
+ List<Symbol> restrictedOutputs = node.getOutputSymbols().stream()
+ .filter(permittedOutputs::contains)
+ .collect(toImmutableList());
+
+ if (restrictedOutputs.size() == node.getOutputSymbols().size()) {
+ return Optional.empty();
+ }
+
+ return Optional.of(
+ new ProjectNode(
+ idAllocator.getNextId(),
+ node,
+ Assignments.identity(restrictedOutputs)));
+ }
}
