diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeLocalShuffleReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeLocalShuffleReader.scala
index 44fd9fe21f..b17af00053 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeLocalShuffleReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeLocalShuffleReader.scala
@@ -141,15 +141,8 @@ object OptimizeLocalShuffleReader extends CustomShuffleReaderRule {
     case s: ShuffleQueryStageExec =>
       s.mapStats.isDefined && supportLocalReader(s.shuffle)
     case CustomShuffleReaderExec(s: ShuffleQueryStageExec, partitionSpecs) =>
-      s.shuffle.shuffleOrigin match {
-        case ENSURE_REQUIREMENTS =>
-          s.mapStats.isDefined && partitionSpecs.nonEmpty && supportLocalReader(s.shuffle)
-        case REBALANCE_PARTITIONS_BY_NONE =>
-          // Use LocalShuffleReader only when we can't CoalesceShufflePartitions
-          s.mapStats.exists(_.bytesByPartitionId.length == partitionSpecs.size) &&
-            partitionSpecs.nonEmpty && supportLocalReader(s.shuffle)
-        case _ => false
-      }
+      s.mapStats.isDefined && partitionSpecs.nonEmpty && supportLocalReader(s.shuffle) &&
+        s.shuffle.shuffleOrigin == ENSURE_REQUIREMENTS
     case _ => false
   }