[SPARK-32755][SQL] Maintain the order of expressions in AttributeSet and ExpressionSet

### What changes were proposed in this pull request?
This PR changes `AttributeSet` and `ExpressionSet` to maintain the insertion order of the elements. More specifically, we:
- change the underlying data structure of `AttributeSet` from `HashSet` to `LinkedHashSet` to maintain the insertion order.
- `ExpressionSet` already uses a list to keep track of the expressions, however, since it is extending Scala's immutable.Set class, operations such as map and flatMap are delegated to the immutable.Set itself. This means that the result of these operations is not an instance of ExpressionSet anymore, rather it's a implementation picked up by the parent class. We also remove this inheritance from `immutable.Set `and implement the needed methods directly. ExpressionSet has a very specific semantics and it does not make sense to extend `immutable.Set` anyway.
- change the `PlanStabilitySuite` to not sort the attributes, to be able to catch changes in the order of expressions in different runs.

### Why are the changes needed?
Expressions identity is based on the `ExprId` which is an auto-incremented number. This means that the same query can yield a query plan with different expression ids in different runs. `AttributeSet` and `ExpressionSet` internally use a `HashSet` as the underlying data structure, and therefore cannot guarantee the a fixed order of operations in different runs. This can be problematic in cases we like to check for plan changes in different runs.

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
Passes `PlanStabilitySuite` after regenerating the golden files.

Closes #29598 from dbaliafroozeh/FixOrderOfExpressions.

Authored-by: Ali Afroozeh <[email protected]>
Signed-off-by: herman <[email protected]>

Author: dbaliafroozeh

sql/catalyst/src/main/scala-2.12/org/apache/spark/sql/catalyst/expressions/ExpressionSet.scala

@@ -27,6 +27,10 @@ object ExpressionSet {
     expressions.foreach(set.add)
     set
   }
+
+  def apply(): ExpressionSet = {
+    new ExpressionSet()
+  }
 }
 
 /**
@@ -53,46 +57,102 @@ object ExpressionSet {
  * This is consistent with how we define `semanticEquals` between two expressions.
  */
 class ExpressionSet protected(
-    protected val baseSet: mutable.Set[Expression] = new mutable.HashSet,
-    protected val originals: mutable.Buffer[Expression] = new ArrayBuffer)
-  extends Set[Expression] {
+    private val baseSet: mutable.Set[Expression] = new mutable.HashSet,
+    private val originals: mutable.Buffer[Expression] = new ArrayBuffer)
+  extends Iterable[Expression] {
 
   //  Note: this class supports Scala 2.12. A parallel source tree has a 2.13 implementation.
 
   protected def add(e: Expression): Unit = {
     if (!e.deterministic) {
       originals += e
-    } else if (!baseSet.contains(e.canonicalized) ) {
+    } else if (!baseSet.contains(e.canonicalized)) {
       baseSet.add(e.canonicalized)
       originals += e
     }
   }
 
-  override def contains(elem: Expression): Boolean = baseSet.contains(elem.canonicalized)
+  protected def remove(e: Expression): Unit = {
+    if (e.deterministic) {
+      baseSet --= baseSet.filter(_ == e.canonicalized)
+      originals --= originals.filter(_.canonicalized == e.canonicalized)
+    }
+  }
+
+  def contains(elem: Expression): Boolean = baseSet.contains(elem.canonicalized)
+
+  override def filter(p: Expression => Boolean): ExpressionSet = {
+    val newBaseSet = baseSet.filter(e => p(e.canonicalized))
+    val newOriginals = originals.filter(e => p(e.canonicalized))
+    new ExpressionSet(newBaseSet, newOriginals)
+  }
+
+  override def filterNot(p: Expression => Boolean): ExpressionSet = {
+    val newBaseSet = baseSet.filterNot(e => p(e.canonicalized))
+    val newOriginals = originals.filterNot(e => p(e.canonicalized))
+    new ExpressionSet(newBaseSet, newOriginals)
+  }
 
-  override def +(elem: Expression): ExpressionSet = {
-    val newSet = new ExpressionSet(baseSet.clone(), originals.clone())
+  def +(elem: Expression): ExpressionSet = {
+    val newSet = clone()
     newSet.add(elem)
     newSet
   }
 
-  override def ++(elems: GenTraversableOnce[Expression]): ExpressionSet = {
-    val newSet = new ExpressionSet(baseSet.clone(), originals.clone())
+  def ++(elems: GenTraversableOnce[Expression]): ExpressionSet = {
+    val newSet = clone()
     elems.foreach(newSet.add)
     newSet
   }
 
-  override def -(elem: Expression): ExpressionSet = {
-    if (elem.deterministic) {
-      val newBaseSet = baseSet.clone().filterNot(_ == elem.canonicalized)
-      val newOriginals = originals.clone().filterNot(_.canonicalized == elem.canonicalized)
-      new ExpressionSet(newBaseSet, newOriginals)
-    } else {
-      new ExpressionSet(baseSet.clone(), originals.clone())
-    }
+  def -(elem: Expression): ExpressionSet = {
+    val newSet = clone()
+    newSet.remove(elem)
+    newSet
+  }
+
+  def --(elems: GenTraversableOnce[Expression]): ExpressionSet = {
+    val newSet = clone()
+    elems.foreach(newSet.remove)
+    newSet
   }
 
-  override def iterator: Iterator[Expression] = originals.iterator
+  def map(f: Expression => Expression): ExpressionSet = {
+    val newSet = new ExpressionSet()
+    this.iterator.foreach(elem => newSet.add(f(elem)))
+    newSet
+  }
+
+  def flatMap(f: Expression => Iterable[Expression]): ExpressionSet = {
+    val newSet = new ExpressionSet()
+    this.iterator.foreach(f(_).foreach(newSet.add))
+    newSet
+  }
+
+  def iterator: Iterator[Expression] = originals.iterator
+
+  def union(that: ExpressionSet): ExpressionSet = {
+    val newSet = clone()
+    that.iterator.foreach(newSet.add)
+    newSet
+  }
+
+  def subsetOf(that: ExpressionSet): Boolean = this.iterator.forall(that.contains)
+
+  def intersect(that: ExpressionSet): ExpressionSet = this.filter(that.contains)
+
+  def diff(that: ExpressionSet): ExpressionSet = this -- that
+
+  def apply(elem: Expression): Boolean = this.contains(elem)
+
+  override def equals(obj: Any): Boolean = obj match {
+    case other: ExpressionSet => this.baseSet == other.baseSet
+    case _ => false
+  }
+
+  override def hashCode(): Int = baseSet.hashCode()
+
+  override def clone(): ExpressionSet = new ExpressionSet(baseSet.clone(), originals.clone())
 
   /**
    * Returns a string containing both the post [[Canonicalize]] expressions and the original

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/AttributeSet.scala

@@ -37,7 +37,11 @@ object AttributeSet {
   val empty = apply(Iterable.empty)
 
   /** Constructs a new [[AttributeSet]] that contains a single [[Attribute]]. */
-  def apply(a: Attribute): AttributeSet = new AttributeSet(Set(new AttributeEquals(a)))
+  def apply(a: Attribute): AttributeSet = {
+    val baseSet = new mutable.LinkedHashSet[AttributeEquals]
+    baseSet += new AttributeEquals(a)
+    new AttributeSet(baseSet)
+  }
 
   /** Constructs a new [[AttributeSet]] given a sequence of [[Expression Expressions]]. */
   def apply(baseSet: Iterable[Expression]): AttributeSet = {
@@ -47,7 +51,7 @@ object AttributeSet {
   /** Constructs a new [[AttributeSet]] given a sequence of [[AttributeSet]]s. */
   def fromAttributeSets(sets: Iterable[AttributeSet]): AttributeSet = {
     val baseSet = sets.foldLeft(new mutable.LinkedHashSet[AttributeEquals]())( _ ++= _.baseSet)
-    new AttributeSet(baseSet.toSet)
+    new AttributeSet(baseSet)
   }
 }
 
@@ -62,7 +66,7 @@ object AttributeSet {
  * and also makes doing transformations hard (we always try keep older trees instead of new ones
  * when the transformation was a no-op).
  */
-class AttributeSet private (val baseSet: Set[AttributeEquals])
+class AttributeSet private (private val baseSet: mutable.LinkedHashSet[AttributeEquals])
   extends Iterable[Attribute] with Serializable {
 
   override def hashCode: Int = baseSet.hashCode()

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CostBasedJoinReorder.scala

@@ -20,7 +20,7 @@ package org.apache.spark.sql.catalyst.optimizer
 import scala.collection.mutable
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet, Expression, PredicateHelper}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet, Expression, ExpressionSet, PredicateHelper}
 import org.apache.spark.sql.catalyst.plans.{Inner, InnerLike, JoinType}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -75,18 +75,18 @@ object CostBasedJoinReorder extends Rule[LogicalPlan] with PredicateHelper {
    * Extracts items of consecutive inner joins and join conditions.
    * This method works for bushy trees and left/right deep trees.
    */
-  private def extractInnerJoins(plan: LogicalPlan): (Seq[LogicalPlan], Set[Expression]) = {
+  private def extractInnerJoins(plan: LogicalPlan): (Seq[LogicalPlan], ExpressionSet) = {
     plan match {
       case Join(left, right, _: InnerLike, Some(cond), JoinHint.NONE) =>
         val (leftPlans, leftConditions) = extractInnerJoins(left)
         val (rightPlans, rightConditions) = extractInnerJoins(right)
-        (leftPlans ++ rightPlans, splitConjunctivePredicates(cond).toSet ++
-          leftConditions ++ rightConditions)
+        (leftPlans ++ rightPlans, leftConditions ++ rightConditions ++
+          splitConjunctivePredicates(cond))
       case Project(projectList, j @ Join(_, _, _: InnerLike, Some(cond), JoinHint.NONE))
         if projectList.forall(_.isInstanceOf[Attribute]) =>
         extractInnerJoins(j)
       case _ =>
-        (Seq(plan), Set())
+        (Seq(plan), ExpressionSet())
     }
   }
 
@@ -143,15 +143,15 @@ object JoinReorderDP extends PredicateHelper with Logging {
   def search(
       conf: SQLConf,
       items: Seq[LogicalPlan],
-      conditions: Set[Expression],
+      conditions: ExpressionSet,
       output: Seq[Attribute]): LogicalPlan = {
 
     val startTime = System.nanoTime()
     // Level i maintains all found plans for i + 1 items.
     // Create the initial plans: each plan is a single item with zero cost.
     val itemIndex = items.zipWithIndex
     val foundPlans = mutable.Buffer[JoinPlanMap](itemIndex.map {
-      case (item, id) => Set(id) -> JoinPlan(Set(id), item, Set.empty, Cost(0, 0))
+      case (item, id) => Set(id) -> JoinPlan(Set(id), item, ExpressionSet(), Cost(0, 0))
     }.toMap)
 
     // Build filters from the join graph to be used by the search algorithm.
@@ -194,7 +194,7 @@ object JoinReorderDP extends PredicateHelper with Logging {
   private def searchLevel(
       existingLevels: Seq[JoinPlanMap],
       conf: SQLConf,
-      conditions: Set[Expression],
+      conditions: ExpressionSet,
       topOutput: AttributeSet,
       filters: Option[JoinGraphInfo]): JoinPlanMap = {
 
@@ -255,7 +255,7 @@ object JoinReorderDP extends PredicateHelper with Logging {
       oneJoinPlan: JoinPlan,
       otherJoinPlan: JoinPlan,
       conf: SQLConf,
-      conditions: Set[Expression],
+      conditions: ExpressionSet,
       topOutput: AttributeSet,
       filters: Option[JoinGraphInfo]): Option[JoinPlan] = {
 
@@ -329,7 +329,7 @@ object JoinReorderDP extends PredicateHelper with Logging {
   case class JoinPlan(
       itemIds: Set[Int],
       plan: LogicalPlan,
-      joinConds: Set[Expression],
+      joinConds: ExpressionSet,
       planCost: Cost) {
 
     /** Get the cost of the root node of this plan tree. */
@@ -387,7 +387,7 @@ object JoinReorderDPFilters extends PredicateHelper {
   def buildJoinGraphInfo(
       conf: SQLConf,
       items: Seq[LogicalPlan],
-      conditions: Set[Expression],
+      conditions: ExpressionSet,
       itemIndex: Seq[(LogicalPlan, Int)]): Option[JoinGraphInfo] = {
 
     if (conf.joinReorderDPStarFilter) {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -921,13 +921,13 @@ object InferFiltersFromConstraints extends Rule[LogicalPlan]
   private def getAllConstraints(
       left: LogicalPlan,
       right: LogicalPlan,
-      conditionOpt: Option[Expression]): Set[Expression] = {
+      conditionOpt: Option[Expression]): ExpressionSet = {
     val baseConstraints = left.constraints.union(right.constraints)
-      .union(conditionOpt.map(splitConjunctivePredicates).getOrElse(Nil).toSet)
+      .union(ExpressionSet(conditionOpt.map(splitConjunctivePredicates).getOrElse(Nil)))
     baseConstraints.union(inferAdditionalConstraints(baseConstraints))
   }
 
-  private def inferNewFilter(plan: LogicalPlan, constraints: Set[Expression]): LogicalPlan = {
+  private def inferNewFilter(plan: LogicalPlan, constraints: ExpressionSet): LogicalPlan = {
     val newPredicates = constraints
       .union(constructIsNotNullConstraints(constraints, plan.output))
       .filter { c =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -169,8 +169,8 @@ abstract class UnaryNode extends LogicalPlan {
    * Generates all valid constraints including an set of aliased constraints by replacing the
    * original constraint expressions with the corresponding alias
    */
-  protected def getAllValidConstraints(projectList: Seq[NamedExpression]): Set[Expression] = {
-    var allConstraints = child.constraints.asInstanceOf[Set[Expression]]
+  protected def getAllValidConstraints(projectList: Seq[NamedExpression]): ExpressionSet = {
+    var allConstraints = child.constraints
     projectList.foreach {
       case a @ Alias(l: Literal, _) =>
         allConstraints += EqualNullSafe(a.toAttribute, l)
@@ -187,7 +187,7 @@ abstract class UnaryNode extends LogicalPlan {
     allConstraints
   }
 
-  override protected lazy val validConstraints: Set[Expression] = child.constraints
+  override protected lazy val validConstraints: ExpressionSet = child.constraints
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/QueryPlanConstraints.scala

@@ -29,16 +29,14 @@ trait QueryPlanConstraints extends ConstraintHelper { self: LogicalPlan =>
    */
   lazy val constraints: ExpressionSet = {
     if (conf.constraintPropagationEnabled) {
-      ExpressionSet(
-        validConstraints
-          .union(inferAdditionalConstraints(validConstraints))
-          .union(constructIsNotNullConstraints(validConstraints, output))
-          .filter { c =>
-            c.references.nonEmpty && c.references.subsetOf(outputSet) && c.deterministic
-          }
-      )
+      validConstraints
+        .union(inferAdditionalConstraints(validConstraints))
+        .union(constructIsNotNullConstraints(validConstraints, output))
+        .filter { c =>
+          c.references.nonEmpty && c.references.subsetOf(outputSet) && c.deterministic
+        }
     } else {
-      ExpressionSet(Set.empty)
+      ExpressionSet()
     }
   }
 
@@ -50,7 +48,7 @@ trait QueryPlanConstraints extends ConstraintHelper { self: LogicalPlan =>
    *
    * See [[Canonicalize]] for more details.
    */
-  protected lazy val validConstraints: Set[Expression] = Set.empty
+  protected lazy val validConstraints: ExpressionSet = ExpressionSet()
 }
 
 trait ConstraintHelper {
@@ -60,8 +58,8 @@ trait ConstraintHelper {
    * For e.g., if an operator has constraints of the form (`a = 5`, `a = b`), this returns an
    * additional constraint of the form `b = 5`.
    */
-  def inferAdditionalConstraints(constraints: Set[Expression]): Set[Expression] = {
-    var inferredConstraints = Set.empty[Expression]
+  def inferAdditionalConstraints(constraints: ExpressionSet): ExpressionSet = {
+    var inferredConstraints = ExpressionSet()
     // IsNotNull should be constructed by `constructIsNotNullConstraints`.
     val predicates = constraints.filterNot(_.isInstanceOf[IsNotNull])
     predicates.foreach {
@@ -79,9 +77,9 @@ trait ConstraintHelper {
   }
 
   private def replaceConstraints(
-      constraints: Set[Expression],
+      constraints: ExpressionSet,
       source: Expression,
-      destination: Expression): Set[Expression] = constraints.map(_ transform {
+      destination: Expression): ExpressionSet = constraints.map(_ transform {
     case e: Expression if e.semanticEquals(source) => destination
   })
 
@@ -91,15 +89,15 @@ trait ConstraintHelper {
    * returns a constraint of the form `isNotNull(a)`
    */
   def constructIsNotNullConstraints(
-      constraints: Set[Expression],
-      output: Seq[Attribute]): Set[Expression] = {
+      constraints: ExpressionSet,
+      output: Seq[Attribute]): ExpressionSet = {
     // First, we propagate constraints from the null intolerant expressions.
-    var isNotNullConstraints: Set[Expression] = constraints.flatMap(inferIsNotNullConstraints)
+    var isNotNullConstraints = constraints.flatMap(inferIsNotNullConstraints(_))
 
     // Second, we infer additional constraints from non-nullable attributes that are part of the
     // operator's output
     val nonNullableAttributes = output.filterNot(_.nullable)
-    isNotNullConstraints ++= nonNullableAttributes.map(IsNotNull).toSet
+    isNotNullConstraints ++= nonNullableAttributes.map(IsNotNull)
 
     isNotNullConstraints -- constraints
   }