[VL] RAS: Group reduction support

.github/workflows/velox_docker.yml

@@ -24,6 +24,7 @@ on:
       - 'gluten-celeborn/common'
       - 'gluten-celeborn/package'
       - 'gluten-celeborn/velox'
+      - 'gluten-ras/**'
       - 'gluten-core/**'
       - 'gluten-data/**'
       - 'gluten-delta/**'

gluten-ras/common/src/main/scala/org/apache/gluten/ras/memo/Memo.scala

@@ -126,8 +126,20 @@ object Memo {
       extends MemoLike[T] {
       private val ras = parent.ras
 
+      // TODO: Traverse up the tree to do more merges.
       private def prepareInsert(node: T): Prepare[T] = {
-        assert(!ras.isGroupLeaf(node))
+        if (ras.isGroupLeaf(node)) {
+          val group = parent.memoTable.allGroups()(ras.planModel.getGroupId(node))
+          val residentCluster = group.clusterKey()
+
+          if (residentCluster == targetCluster) {
+            return Prepare.cluster(parent, targetCluster)
+          }
+          // The resident cluster of group leaf is not the same with target cluster.
+          // Merge.
+          parent.memoTable.mergeClusters(residentCluster, targetCluster)
+          return Prepare.cluster(parent, targetCluster)
+        }
 
         val childrenPrepares =
           ras.planModel.childrenOf(node).map(child => parent.prepareInsert(child))
@@ -155,8 +167,6 @@ object Memo {
         }
         // The new node already memorized to memo, but in the different cluster.
         // Merge the two clusters.
-        //
-        // TODO: Traverse up the tree to do more merges.
         parent.memoTable.mergeClusters(cachedCluster, targetCluster)
         Prepare.tree(parent, targetCluster, childrenPrepares)
       }

gluten-ras/common/src/test/scala/org/apache/gluten/ras/RasSuite.scala

@@ -172,6 +172,32 @@ abstract class RasSuite extends AnyFunSuite {
     assert(optimized == Unary(23, Unary(23, Leaf(70))))
   }
 
+  test(s"Group reduction") {
+    object RemoveUnary extends RasRule[TestNode] {
+      override def shift(node: TestNode): Iterable[TestNode] = node match {
+        case Unary(cost, child) => List(child)
+        case other => List.empty
+      }
+
+      override def shape(): Shape[TestNode] = Shapes.fixedHeight(1)
+    }
+
+    val ras =
+      Ras[TestNode](
+        PlanModelImpl,
+        CostModelImpl,
+        MetadataModelImpl,
+        PropertyModelImpl,
+        ExplainImpl,
+        RasRule.Factory.reuse(List(RemoveUnary)))
+        .withNewConfig(_ => conf)
+    val plan = Unary(60, Unary(90, Leaf(70)))
+    val planner = ras.newPlanner(plan)
+    val optimized = planner.plan()
+
+    assert(optimized == Leaf(70))
+  }
+
   test(s"Unary node insertion") {
     object InsertUnary2 extends RasRule[TestNode] {
       override def shift(node: TestNode): Iterable[TestNode] = node match {