Refactor apply method in InsertMapSortInRepartitionExpressions

apache · Jan 7, 2025 · 20c76e4 · 20c76e4
1 parent 75e323b
commit 20c76e4
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/...lyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala b/...lyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala
@@ -89,21 +89,18 @@ object InsertMapSortInRepartitionExpressions extends Rule[LogicalPlan] {
 
   override def apply(plan: LogicalPlan): LogicalPlan = {
     plan.transformUpWithPruning(_.containsPattern(REPARTITION_OPERATION)) {
-      case rep @
-        RepartitionByExpression(partitionExprs, child, optNumPartitions, optAdvisoryPartitionSize)
+      case rep: RepartitionByExpression
         if rep.partitionExpressions.exists(mapTypeExistsRecursively) =>
         val exprToMapSort = new mutable.HashMap[Expression, Expression]
-        val newPartitionExprs = partitionExprs.map { expr =>
+        val newPartitionExprs = rep.partitionExpressions.map { expr =>
           val inserted = insertMapSortRecursively(expr)
           if (expr.ne(inserted)) {
             exprToMapSort.getOrElseUpdate(expr.canonicalized, inserted)
           } else {
             expr
           }
         }
-        RepartitionByExpression(
-          newPartitionExprs, child, optNumPartitions, optAdvisoryPartitionSize
-        )
+        rep.copy(partitionExpressions = newPartitionExprs)
     }
   }
 }