public PlanNode extract() { return extract(getNode(rootGroup)); }
@Override public PlanNode optimize(PlanNode plan, Session session, TypeProvider types, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator, WarningCollector warningCollector) { // only disable new rules if we have legacy rules to fall back to if (!SystemSessionProperties.isNewOptimizerEnabled(session) && !legacyRules.isEmpty()) { for (PlanOptimizer optimizer : legacyRules) { plan = optimizer.optimize(plan, session, symbolAllocator.getTypes(), symbolAllocator, idAllocator, warningCollector); } return plan; } Memo memo = new Memo(idAllocator, plan); Lookup lookup = Lookup.from(planNode -> Stream.of(memo.resolve(planNode))); Matcher matcher = new PlanNodeMatcher(lookup); Duration timeout = SystemSessionProperties.getOptimizerTimeout(session); Context context = new Context(memo, lookup, idAllocator, symbolAllocator, System.nanoTime(), timeout.toMillis(), session, warningCollector); exploreGroup(memo.getRootGroup(), context, matcher); return memo.extract(); }
@Test public void testInitialization() { PlanNode plan = node(node()); Memo memo = new Memo(idAllocator, plan); assertEquals(memo.getGroupCount(), 2); assertMatchesStructure(plan, memo.extract()); }
@Test public void testReplaceSubtree() { PlanNode plan = node(node(node())); Memo memo = new Memo(idAllocator, plan); assertEquals(memo.getGroupCount(), 3); // replace child of root node with subtree PlanNode transformed = node(node()); memo.replace(getChildGroup(memo, memo.getRootGroup()), transformed, "rule"); assertEquals(memo.getGroupCount(), 3); assertMatchesStructure(memo.extract(), node(plan.getId(), transformed)); }
@Test public void testRemoveNode() { PlanNode z = node(); PlanNode y = node(z); PlanNode x = node(y); Memo memo = new Memo(idAllocator, x); assertEquals(memo.getGroupCount(), 3); int yGroup = getChildGroup(memo, memo.getRootGroup()); memo.replace(yGroup, memo.getNode(yGroup).getSources().get(0), "rule"); assertEquals(memo.getGroupCount(), 2); assertMatchesStructure( memo.extract(), node(x.getId(), node(z.getId()))); }
@Test public void testReplaceNode() { PlanNode z = node(); PlanNode y = node(z); PlanNode x = node(y); Memo memo = new Memo(idAllocator, x); assertEquals(memo.getGroupCount(), 3); // replace child of root node with another node, retaining child's child int yGroup = getChildGroup(memo, memo.getRootGroup()); GroupReference zRef = (GroupReference) getOnlyElement(memo.getNode(yGroup).getSources()); PlanNode transformed = node(zRef); memo.replace(yGroup, transformed, "rule"); assertEquals(memo.getGroupCount(), 3); assertMatchesStructure(memo.extract(), node(x.getId(), node(transformed.getId(), z))); }
@Test public void testInsertNode() { PlanNode z = node(); PlanNode x = node(z); Memo memo = new Memo(idAllocator, x); assertEquals(memo.getGroupCount(), 2); int zGroup = getChildGroup(memo, memo.getRootGroup()); PlanNode y = node(memo.getNode(zGroup)); memo.replace(zGroup, y, "rule"); assertEquals(memo.getGroupCount(), 3); assertMatchesStructure( memo.extract(), node(x.getId(), node(y.getId(), node(z.getId())))); }
@Test public void testReplaceNonLeafSubtree() { PlanNode w = node(); PlanNode z = node(w); PlanNode y = node(z); PlanNode x = node(y); Memo memo = new Memo(idAllocator, x); assertEquals(memo.getGroupCount(), 4); int yGroup = getChildGroup(memo, memo.getRootGroup()); int zGroup = getChildGroup(memo, yGroup); PlanNode rewrittenW = memo.getNode(zGroup).getSources().get(0); PlanNode newZ = node(rewrittenW); PlanNode newY = node(newZ); memo.replace(yGroup, newY, "rule"); assertEquals(memo.getGroupCount(), 4); assertMatchesStructure( memo.extract(), node(x.getId(), node(newY.getId(), node(newZ.getId(), node(w.getId()))))); }
@Test public void testMultipleReferences() { PlanNode z = node(); PlanNode y = node(z); PlanNode x = node(y); Memo memo = new Memo(idAllocator, x); assertEquals(memo.getGroupCount(), 3); int yGroup = getChildGroup(memo, memo.getRootGroup()); PlanNode rewrittenZ = memo.getNode(yGroup).getSources().get(0); PlanNode y1 = node(rewrittenZ); PlanNode y2 = node(rewrittenZ); PlanNode newX = node(y1, y2); memo.replace(memo.getRootGroup(), newX, "rule"); assertEquals(memo.getGroupCount(), 4); assertMatchesStructure( memo.extract(), node(newX.getId(), node(y1.getId(), node(z.getId())), node(y2.getId(), node(z.getId())))); }