[CALCITE-7030] Enhance TopologicalOrderIterator to support BOTTOM_UP

Author: NobiGo

core/src/main/java/org/apache/calcite/plan/hep/HepPlanner.java

@@ -58,7 +58,6 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -467,30 +466,16 @@ private Iterator<HepRelVertex> getGraphIterator(
       return DepthFirstIterator.of(graph, start).iterator();
 
     case TOP_DOWN:
-      assert start == root;
-      // see above
-/*
-        collectGarbage();
-*/
-      return TopologicalOrderIterator.of(graph).iterator();
-
     case BOTTOM_UP:
-    default:
       assert start == root;
-
       // see above
 /*
         collectGarbage();
 */
-
-      // TODO jvs 4-Apr-2006:  enhance TopologicalOrderIterator
-      // to support reverse walk.
-      final List<HepRelVertex> list = new ArrayList<>();
-      for (HepRelVertex vertex : TopologicalOrderIterator.of(graph)) {
-        list.add(vertex);
-      }
-      Collections.reverse(list);
-      return list.iterator();
+      return TopologicalOrderIterator.of(graph, programState.matchOrder).iterator();
+    default:
+      throw new
+          UnsupportedOperationException("Unsupported match order: " + programState.matchOrder);
     }
   }
 

core/src/main/java/org/apache/calcite/util/graph/DefaultDirectedGraph.java

@@ -61,6 +61,7 @@ public static <V, E extends DefaultEdge> DefaultDirectedGraph<V, E> create(
     return new DefaultDirectedGraph<>(edgeFactory);
   }
 
+  @Deprecated
   public String toStringUnordered() {
     return "graph("
         + "vertices: " + vertexMap.keySet()

core/src/main/java/org/apache/calcite/util/graph/TopologicalOrderIterator.java

@@ -16,6 +16,8 @@
  */
 package org.apache.calcite.util.graph;
 
+import org.apache.calcite.plan.hep.HepMatchOrder;
+
 import org.checkerframework.checker.initialization.qual.UnderInitialization;
 import org.checkerframework.checker.nullness.qual.RequiresNonNull;
 
@@ -38,10 +40,17 @@ public class TopologicalOrderIterator<V, E extends DefaultEdge>
     implements Iterator<V> {
   final Map<V, int[]> countMap = new HashMap<>();
   final List<V> empties = new ArrayList<>();
+  final HepMatchOrder hepMatchOrder;
   private final DefaultDirectedGraph<V, E> graph;
 
   public TopologicalOrderIterator(DirectedGraph<V, E> graph) {
+    this(graph, HepMatchOrder.TOP_DOWN);
+  }
+
+  public TopologicalOrderIterator(DirectedGraph<V, E> graph, HepMatchOrder hepMatchOrder) {
+    assert hepMatchOrder == HepMatchOrder.TOP_DOWN || hepMatchOrder == HepMatchOrder.BOTTOM_UP;
     this.graph = (DefaultDirectedGraph<V, E>) graph;
+    this.hepMatchOrder = hepMatchOrder;
     populate(countMap, empties);
   }
 
@@ -50,21 +59,40 @@ public static <V, E extends DefaultEdge> Iterable<V> of(
     return () -> new TopologicalOrderIterator<>(graph);
   }
 
+  public static <V, E extends DefaultEdge> Iterable<V> of(
+      final DirectedGraph<V, E> graph, HepMatchOrder hepMatchOrder) {
+    return () -> new TopologicalOrderIterator<>(graph, hepMatchOrder);
+  }
+
   @RequiresNonNull("graph")
   private void populate(
       @UnderInitialization TopologicalOrderIterator<V, E> this,
       Map<V, int[]> countMap, List<V> empties) {
     for (V v : graph.vertexMap.keySet()) {
       countMap.put(v, new int[] {0});
     }
-    for (DefaultDirectedGraph.VertexInfo<V, E> info
-        : graph.vertexMap.values()) {
-      for (E edge : info.outEdges) {
-        //noinspection SuspiciousMethodCalls
-        final int[] ints =
-            requireNonNull(countMap.get(edge.target),
-                () -> "no value for " + edge.target);
-        ++ints[0];
+    if (hepMatchOrder == HepMatchOrder.TOP_DOWN) {
+      for (DefaultDirectedGraph.VertexInfo<V, E> info
+          : graph.vertexMap.values()) {
+        for (E edge : info.outEdges) {
+          //noinspection SuspiciousMethodCalls
+          final int[] ints =
+              requireNonNull(countMap.get(edge.target),
+                  () -> "no value for " + edge.target);
+          ++ints[0];
+        }
+      }
+    }
+    if (hepMatchOrder == HepMatchOrder.BOTTOM_UP) {
+      for (DefaultDirectedGraph.VertexInfo<V, E> info
+          : graph.vertexMap.values()) {
+        for (E edge : info.outEdges) {
+          //noinspection SuspiciousMethodCalls
+          final int[] ints =
+              requireNonNull(countMap.get(edge.source),
+                  () -> "no value for " + edge.source);
+          ++ints[0];
+        }
       }
     }
     for (Map.Entry<V, int[]> entry : countMap.entrySet()) {
@@ -84,15 +112,30 @@ private void populate(
     DefaultDirectedGraph.VertexInfo<V, E> vertexInfo =
         requireNonNull(graph.vertexMap.get(v),
             () -> "no vertex " + v);
-    for (E o : vertexInfo.outEdges) {
-      //noinspection unchecked
-      final V target = (V) o.target;
-      int[] ints =
-          requireNonNull(countMap.get(target),
-              () -> "no counts found for target " + target);
-      if (--ints[0] == 0) {
-        countMap.remove(target);
-        empties.add(target);
+    if (hepMatchOrder == HepMatchOrder.TOP_DOWN) {
+      for (E o : vertexInfo.outEdges) {
+        //noinspection unchecked
+        final V target = (V) o.target;
+        int[] ints =
+            requireNonNull(countMap.get(target),
+                () -> "no counts found for target " + target);
+        if (--ints[0] == 0) {
+          countMap.remove(target);
+          empties.add(target);
+        }
+      }
+    }
+    if (hepMatchOrder == HepMatchOrder.BOTTOM_UP) {
+      for (E o : vertexInfo.inEdges) {
+        //noinspection unchecked
+        final V source = (V) o.source;
+        int[] ints =
+            requireNonNull(countMap.get(source),
+                () -> "no counts found for source " + source);
+        if (--ints[0] == 0) {
+          countMap.remove(source);
+          empties.add(source);
+        }
       }
     }
     return v;

core/src/test/java/org/apache/calcite/util/graph/DirectedGraphTest.java

@@ -15,6 +15,8 @@
  * limitations under the License.
  */
 package org.apache.calcite.util.graph;
+import org.apache.calcite.plan.hep.HepMatchOrder;
+
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
@@ -172,6 +174,19 @@ private <V> List<List<V>> paths(DirectedGraph<V, DefaultEdge> g,
     assertThat(list, hasToString("[A, B, E, C, F, D]"));
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7030">[CALCITE-7030]
+   * Enhance TopologicalOrderIterator to support BOTTOM_UP</a>. */
+  @Test void testTopologicalOrderWithBottomUpIterator() {
+    final DefaultDirectedGraph<String, DefaultEdge> graph = createDag();
+    final List<String> list = new ArrayList<>();
+    for (String s : TopologicalOrderIterator.of(graph, HepMatchOrder.BOTTOM_UP)) {
+      list.add(s);
+    }
+    assertThat(list, hasToString("[D, F, C, B, E, A]"));
+  }
+
   private DefaultDirectedGraph<String, DefaultEdge> createDag() {
     //    D         F
     //    ^         ^