[SPARK-35703][SQL] Relax constraint for bucket join and remove HashClusteredDistribution

### What changes were proposed in this pull request?

This PR proposes the following:
1. Introducing a new trait `ShuffleSpec` which is used in `EnsureRequirements` when the node has more than one children and serves two purposes: 1) compare all children and check if they are compatible w.r.t partitioning & distribution, 2) create a new partitioning to re-shuffle the other side in case they are not compatible.
2. Remove `HashClusteredDistribution` and replace its usages with `ClusteredDistribution`.

Under the new mechanism, when `EnsureRequirements` check whether shuffles are required for a plan node with >1 children, it does the following:
1. check each child of the node and see if it can satisfy the corresponding required distribution.
2. check if all children of the node are compatible with each other w.r.t their partitioning and distribution
3. if 2) fails, choose the best shuffle spec (in terms of shuffle parallelism) that can be used to repartition the other children, so that they will all have compatible partitioning

### Why are the changes needed?

Spark currently only allow bucket join when the set of cluster keys from output partitioning _exactly match_ the set of join keys from the required distribution. For instance, in the following:
```sql
SELECT * FROM A JOIN B ON A.c1 = B.c1 AND A.c2 = B.c2
```
bucket join will only be triggered if both `A` and `B` are bucketed on columns `c1` and `c2`, in which case Spark will avoid shuffling both sides of the join.

The above requirement, however, is too strict, as shuffle can also be avoided if both `A` and `B` are bucketed on either column `c1` or `c2`. That is, if all rows that have the same value in column `c1` are clustered into the same partition, then all rows have the same values in column `c1` and `c2` are also clustered into the same partition.

In order to allow this, we'll need to change the logic of deciding whether two sides of a join operator are "co-partitioned". Currently, this is done by checking each side's output partitioning against its required distribution separately, using `Partitioning.satisfies` method. Since `HashClusteredDistribution` requires a `HashPartitioning` to have the exact match on the cluster keys, this can be done in isolation without looking at the other side's output partitioning and required distribution.

However, the approach is no longer valid if we are going to relax the above constraint, as we need to compare the output partitioning and required distribution **on both sides**. For instance, in the above example, if `A` is bucketed on `c1` while `B` is bucketed on `c2`, we may need to do the following check:
1. identify where `A.c1` and `B.c2` is used in the join keys (e.g., position 0 and 1 respectively)
2. check if the positions derived from both sides exactly match each other (this becomes more complicated if a key appears in multiple positions within the join keys.)

In order to achieve the above, this proposes the following:
```scala
trait ShuffleSpec {
  // Used as a cost indicator to shuffle children
  def numPartitions: Int

  // Used to check whether this spec is compatible with `other`
  def isCompatibleWith(other: ShuffleSpec): Boolean

  // Used to create a new partitioning for the other `distribution` in case `isCompatibleWith` failed.
  def createPartitioning(distribution: Distribution): Partitioning
}
```

A similar API is also required if we are going to support DSv2 `DataSourcePartitioning` as output partitioning in bucket join scenario, or support custom hash functions such as `HiveHash` for bucketing. With the former, even if both `A` and `B` are partitioned on columns `c1` and `c2` in the above example, they could be partitioned via different transform expressions, e.g., `A` is on `(bucket(32, c1), day(c2)` while `B` is on `(bucket(32, c1), hour(c2)`. This means we'll need to compare the partitioning from both sides of the join which makes the current approach with `Partitioning.satisfies` insufficient. The same API `isCompatibleWith` can potentially be reused for the purpose.

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

Yes, now bucket join will be enabled for more cases as mentioned above.

### How was this patch tested?

1. Added a new test suite `ShuffleSpecSuite`
2. Added additional tests in `EnsureRequirementsSuite`.

Closes #32875 from sunchao/SPARK-35703.

Authored-by: Chao Sun <sunchao@apple.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: sunchao

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.plans.physical
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types.{DataType, IntegerType}
 
@@ -87,31 +89,6 @@ case class ClusteredDistribution(
   }
 }
 
-/**
- * Represents data where tuples have been clustered according to the hash of the given
- * `expressions`. The hash function is defined as `HashPartitioning.partitionIdExpression`, so only
- * [[HashPartitioning]] can satisfy this distribution.
- *
- * This is a strictly stronger guarantee than [[ClusteredDistribution]]. Given a tuple and the
- * number of partitions, this distribution strictly requires which partition the tuple should be in.
- */
-case class HashClusteredDistribution(
-    expressions: Seq[Expression],
-    requiredNumPartitions: Option[Int] = None) extends Distribution {
-  require(
-    expressions != Nil,
-    "The expressions for hash of a HashClusteredDistribution should not be Nil. " +
-      "An AllTuples should be used to represent a distribution that only has " +
-      "a single partition.")
-
-  override def createPartitioning(numPartitions: Int): Partitioning = {
-    assert(requiredNumPartitions.isEmpty || requiredNumPartitions.get == numPartitions,
-      s"This HashClusteredDistribution requires ${requiredNumPartitions.get} partitions, but " +
-        s"the actual number of partitions is $numPartitions.")
-    HashPartitioning(expressions, numPartitions)
-  }
-}
-
 /**
  * Represents data where tuples have been ordered according to the `ordering`
  * [[Expression Expressions]]. Its requirement is defined as the following:
@@ -171,6 +148,17 @@ trait Partitioning {
     required.requiredNumPartitions.forall(_ == numPartitions) && satisfies0(required)
   }
 
+  /**
+   * Creates a shuffle spec for this partitioning and its required distribution. The
+   * spec is used in the scenario where an operator has multiple children (e.g., join), and is
+   * used to decide whether this child is co-partitioned with others, therefore whether extra
+   * shuffle shall be introduced.
+   *
+   * @param distribution the required clustered distribution for this partitioning
+   */
+  def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+    throw new IllegalStateException(s"Unexpected partitioning: ${getClass.getSimpleName}")
+
   /**
    * The actual method that defines whether this [[Partitioning]] can satisfy the given
    * [[Distribution]], after the `numPartitions` check.
@@ -202,6 +190,9 @@ case object SinglePartition extends Partitioning {
     case _: BroadcastDistribution => false
     case _ => true
   }
+
+  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+    SinglePartitionShuffleSpec
 }
 
 /**
@@ -219,17 +210,16 @@ case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int)
   override def satisfies0(required: Distribution): Boolean = {
     super.satisfies0(required) || {
       required match {
-        case h: HashClusteredDistribution =>
-          expressions.length == h.expressions.length && expressions.zip(h.expressions).forall {
-            case (l, r) => l.semanticEquals(r)
-          }
         case ClusteredDistribution(requiredClustering, _) =>
           expressions.forall(x => requiredClustering.exists(_.semanticEquals(x)))
         case _ => false
       }
     }
   }
 
+  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+    HashShuffleSpec(this, distribution)
+
   /**
    * Returns an expression that will produce a valid partition ID(i.e. non-negative and is less
    * than numPartitions) based on hashing expressions.
@@ -288,6 +278,9 @@ case class RangePartitioning(ordering: Seq[SortOrder], numPartitions: Int)
     }
   }
 
+  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec =
+    RangeShuffleSpec(this.numPartitions, distribution)
+
   override protected def withNewChildrenInternal(
       newChildren: IndexedSeq[Expression]): RangePartitioning =
     copy(ordering = newChildren.asInstanceOf[Seq[SortOrder]])
@@ -330,6 +323,11 @@ case class PartitioningCollection(partitionings: Seq[Partitioning])
   override def satisfies0(required: Distribution): Boolean =
     partitionings.exists(_.satisfies(required))
 
+  override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
+    val filtered = partitionings.filter(_.satisfies(distribution))
+    ShuffleSpecCollection(filtered.map(_.createShuffleSpec(distribution)))
+  }
+
   override def toString: String = {
     partitionings.map(_.toString).mkString("(", " or ", ")")
   }
@@ -352,3 +350,151 @@ case class BroadcastPartitioning(mode: BroadcastMode) extends Partitioning {
     case _ => false
   }
 }
+
+/**
+ * This is used in the scenario where an operator has multiple children (e.g., join) and one or more
+ * of which have their own requirement regarding whether its data can be considered as
+ * co-partitioned from others. This offers APIs for:
+ *
+ *   - Comparing with specs from other children of the operator and check if they are compatible.
+ *      When two specs are compatible, we can say their data are co-partitioned, and Spark will
+ *      potentially be able to eliminate shuffle if necessary.
+ *   - Creating a partitioning that can be used to re-partition another child, so that to make it
+ *      having a compatible partitioning as this node.
+ */
+trait ShuffleSpec {
+  /**
+   * Returns the number of partitions of this shuffle spec
+   */
+  def numPartitions: Int
+
+  /**
+   * Returns true iff this spec is compatible with the provided shuffle spec.
+   *
+   * A true return value means that the data partitioning from this spec can be seen as
+   * co-partitioned with the `other`, and therefore no shuffle is required when joining the two
+   * sides.
+   */
+  def isCompatibleWith(other: ShuffleSpec): Boolean
+
+  /**
+   * Whether this shuffle spec can be used to create partitionings for the other children.
+   */
+  def canCreatePartitioning: Boolean = false
+
+  /**
+   * Creates a partitioning that can be used to re-partition the other side with the given
+   * clustering expressions.
+   *
+   * This will only be called when:
+   *  - [[canCreatePartitioning]] returns true.
+   *  - [[isCompatibleWith]] returns false on the side where the `clustering` is from.
+   */
+  def createPartitioning(clustering: Seq[Expression]): Partitioning =
+    throw new UnsupportedOperationException("Operation unsupported for " +
+        s"${getClass.getCanonicalName}")
+}
+
+case object SinglePartitionShuffleSpec extends ShuffleSpec {
+  override def isCompatibleWith(other: ShuffleSpec): Boolean = {
+    other.numPartitions == 1
+  }
+
+  override def canCreatePartitioning: Boolean = true
+
+  override def createPartitioning(clustering: Seq[Expression]): Partitioning =
+    SinglePartition
+
+  override def numPartitions: Int = 1
+}
+
+case class RangeShuffleSpec(
+    numPartitions: Int,
+    distribution: ClusteredDistribution) extends ShuffleSpec {
+
+  override def isCompatibleWith(other: ShuffleSpec): Boolean = other match {
+    case SinglePartitionShuffleSpec => numPartitions == 1
+    case ShuffleSpecCollection(specs) => specs.exists(isCompatibleWith)
+    // `RangePartitioning` is not compatible with any other partitioning since it can't guarantee
+    // data are co-partitioned for all the children, as range boundaries are randomly sampled.
+    case _ => false
+  }
+}
+
+case class HashShuffleSpec(
+    partitioning: HashPartitioning,
+    distribution: ClusteredDistribution) extends ShuffleSpec {
+  lazy val hashKeyPositions: Seq[mutable.BitSet] =
+    createHashKeyPositions(distribution.clustering, partitioning.expressions)
+
+  override def isCompatibleWith(other: ShuffleSpec): Boolean = other match {
+    case SinglePartitionShuffleSpec =>
+      partitioning.numPartitions == 1
+    case otherHashSpec @ HashShuffleSpec(otherPartitioning, otherDistribution) =>
+      // we need to check:
+      //  1. both distributions have the same number of clustering expressions
+      //  2. both partitioning have the same number of partitions
+      //  3. both partitioning have the same number of expressions
+      //  4. each pair of expression from both has overlapping positions in their
+      //     corresponding distributions.
+      distribution.clustering.length == otherDistribution.clustering.length &&
+      partitioning.numPartitions == otherPartitioning.numPartitions &&
+      partitioning.expressions.length == otherPartitioning.expressions.length && {
+        val otherHashKeyPositions = otherHashSpec.hashKeyPositions
+        hashKeyPositions.zip(otherHashKeyPositions).forall { case (left, right) =>
+          left.intersect(right).nonEmpty
+        }
+      }
+    case ShuffleSpecCollection(specs) =>
+      specs.exists(isCompatibleWith)
+    case _ =>
+      false
+  }
+
+  override def canCreatePartitioning: Boolean = true
+
+  override def createPartitioning(clustering: Seq[Expression]): Partitioning = {
+    val exprs = hashKeyPositions.map(v => clustering(v.head))
+    HashPartitioning(exprs, partitioning.numPartitions)
+  }
+
+  override def numPartitions: Int = partitioning.numPartitions
+
+  /**
+   * Returns a sequence where each element is a set of positions of the key in `hashKeys` to its
+   * positions in `requiredClusterKeys`. For instance, if `requiredClusterKeys` is [a, b, b] and
+   * `hashKeys` is [a, b], the result will be [(0), (1, 2)].
+   */
+  private def createHashKeyPositions(
+      requiredClusterKeys: Seq[Expression],
+      hashKeys: Seq[Expression]): Seq[mutable.BitSet] = {
+    val distKeyToPos = mutable.Map.empty[Expression, mutable.BitSet]
+    requiredClusterKeys.zipWithIndex.foreach { case (distKey, distKeyPos) =>
+      distKeyToPos.getOrElseUpdate(distKey.canonicalized, mutable.BitSet.empty).add(distKeyPos)
+    }
+
+    hashKeys.map(k => distKeyToPos(k.canonicalized))
+  }
+}
+
+case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec {
+  override def isCompatibleWith(other: ShuffleSpec): Boolean = {
+    specs.exists(_.isCompatibleWith(other))
+  }
+
+  override def canCreatePartitioning: Boolean =
+    specs.forall(_.canCreatePartitioning)
+
+  override def createPartitioning(clustering: Seq[Expression]): Partitioning = {
+    // as we only consider # of partitions as the cost now, it doesn't matter which one we choose
+    // since they should all have the same # of partitions.
+    require(specs.map(_.numPartitions).toSet.size == 1, "expected all specs in the collection " +
+      "to have the same number of partitions")
+    specs.head.createPartitioning(clustering)
+  }
+
+  override def numPartitions: Int = {
+    require(specs.nonEmpty, "expected specs to be non-empty")
+    specs.head.numPartitions
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala

@@ -133,11 +133,6 @@ class DistributionSuite extends SparkFunSuite {
       ClusteredDistribution(Seq($"a", $"b", $"c")),
       true)
 
-    checkSatisfied(
-      SinglePartition,
-      HashClusteredDistribution(Seq($"a", $"b", $"c")),
-      true)
-
     checkSatisfied(
       SinglePartition,
       OrderedDistribution(Seq($"a".asc, $"b".asc, $"c".asc)),
@@ -172,23 +167,6 @@ class DistributionSuite extends SparkFunSuite {
       ClusteredDistribution(Seq($"d", $"e")),
       false)
 
-    // HashPartitioning can satisfy HashClusteredDistribution iff its hash expressions are exactly
-    // same with the required hash clustering expressions.
-    checkSatisfied(
-      HashPartitioning(Seq($"a", $"b", $"c"), 10),
-      HashClusteredDistribution(Seq($"a", $"b", $"c")),
-      true)
-
-    checkSatisfied(
-      HashPartitioning(Seq($"c", $"b", $"a"), 10),
-      HashClusteredDistribution(Seq($"a", $"b", $"c")),
-      false)
-
-    checkSatisfied(
-      HashPartitioning(Seq($"a", $"b"), 10),
-      HashClusteredDistribution(Seq($"a", $"b", $"c")),
-      false)
-
     // HashPartitioning cannot satisfy OrderedDistribution
     checkSatisfied(
       HashPartitioning(Seq($"a", $"b", $"c"), 10),
@@ -269,12 +247,6 @@ class DistributionSuite extends SparkFunSuite {
       RangePartitioning(Seq($"a".asc, $"b".asc, $"c".asc), 10),
       ClusteredDistribution(Seq($"c", $"d")),
       false)
-
-    // RangePartitioning cannot satisfy HashClusteredDistribution
-    checkSatisfied(
-      RangePartitioning(Seq($"a".asc, $"b".asc, $"c".asc), 10),
-      HashClusteredDistribution(Seq($"a", $"b", $"c")),
-      false)
   }
 
   test("Partitioning.numPartitions must match Distribution.requiredNumPartitions to satisfy it") {
@@ -283,21 +255,11 @@ class DistributionSuite extends SparkFunSuite {
       ClusteredDistribution(Seq($"a", $"b", $"c"), Some(10)),
       false)
 
-    checkSatisfied(
-      SinglePartition,
-      HashClusteredDistribution(Seq($"a", $"b", $"c"), Some(10)),
-      false)
-
     checkSatisfied(
       HashPartitioning(Seq($"a", $"b", $"c"), 10),
       ClusteredDistribution(Seq($"a", $"b", $"c"), Some(5)),
       false)
 
-    checkSatisfied(
-      HashPartitioning(Seq($"a", $"b", $"c"), 10),
-      HashClusteredDistribution(Seq($"a", $"b", $"c"), Some(5)),
-      false)
-
     checkSatisfied(
       RangePartitioning(Seq($"a".asc, $"b".asc, $"c".asc), 10),
       ClusteredDistribution(Seq($"a", $"b", $"c"), Some(5)),