[SPARK-28715][SQL] Introduce collectInPlanAndSubqueries and subqueriesAll in QueryPlan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala

@@ -188,6 +188,22 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] extends TreeNode[PlanT
     })
   }
 
+  /**
+   * Returns a sequence containing the result of applying a partial function to all elements in this
+   * plan, also considering all the plans in its (nested) subqueries
+   */
+  def collectInPlanAndSubqueries[B](f: PartialFunction[PlanType, B]): Seq[B] =
+    (this +: subqueriesAll).flatMap(_.collect(f))
+
+  /**
+   * Returns a sequence containing the subqueries in this plan, also including the (nested)
+   * subquries in its children
+   */
+  def subqueriesAll: Seq[PlanType] = {
+    val subqueries = this.flatMap(_.subqueries)
+    subqueries ++ subqueries.flatMap(_.subqueriesAll)
+  }
+
   override protected def innerChildren: Seq[QueryPlan[_]] = subqueries
 
   /**

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala

@@ -39,4 +39,45 @@ class QueryPlanSuite extends SparkFunSuite {
     assert(mappedOrigin == Origin.apply(Some(0), Some(0)))
   }
 
+  test("collectInPlanAndSubqueries") {
+    val a: NamedExpression = AttributeReference("a", IntegerType)()
+    val plan =
+      Union(
+        Seq(
+          Project(
+            Seq(a),
+            Filter(
+              ListQuery(Project(
+                Seq(a),
+                Filter(
+                  ListQuery(Project(
+                    Seq(a),
+                    UnresolvedRelation(TableIdentifier("t", None))
+                  )),
+                  UnresolvedRelation(TableIdentifier("t", None))
+                )
+              )),
+              UnresolvedRelation(TableIdentifier("t", None))
+            )
+          ),
+          Project(
+            Seq(a),
+            Filter(
+              ListQuery(Project(
+                Seq(a),
+                UnresolvedRelation(TableIdentifier("t", None))
+              )),
+              UnresolvedRelation(TableIdentifier("t", None))
+            )
+          )
+        )
+      )
+
+    val countRelationsInPlan = plan.collect({ case _: UnresolvedRelation => 1 }).sum
+    val countRelationsInPlanAndSubqueries =
+      plan.collectInPlanAndSubqueries({ case _: UnresolvedRelation => 1 }).sum
+
+    assert(countRelationsInPlan == 2)
+    assert(countRelationsInPlanAndSubqueries == 5)
+  }
 }