From d67071e33a5cfca81a700a58eed55d16a348e93f Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 25 Apr 2024 12:59:57 -0700 Subject: [PATCH] Implement equals and hashCode for CometShuffleExchangeExec --- .../shuffle/CometShuffleExchangeExec.scala | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index faa52cd06..e05ac7a23 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -50,6 +50,8 @@ import org.apache.spark.util.MutablePair import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordComparator} import org.apache.spark.util.random.XORShiftRandom +import com.google.common.base.Objects + import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde} import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.QueryPlanSerde.serializeDataType @@ -193,6 +195,21 @@ case class CometShuffleExchangeExec( override protected def withNewChildInternal(newChild: SparkPlan): CometShuffleExchangeExec = copy(child = newChild) + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometShuffleExchangeExec => + this.outputPartitioning == other.outputPartitioning && + this.shuffleOrigin == other.shuffleOrigin && this.child == other.child && + this.shuffleType == other.shuffleType && + this.advisoryPartitionSize == other.advisoryPartitionSize + case _ => + false + } + } + + override def hashCode(): Int = + Objects.hashCode(outputPartitioning, shuffleOrigin, shuffleType, advisoryPartitionSize, child) } object CometShuffleExchangeExec extends ShimCometShuffleExchangeExec {