cloud-fan commented on code in PR #37612:
URL: https://github.com/apache/spark/pull/37612#discussion_r955761128
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PropagateEmptyRelation.scala:
##########
@@ -162,14 +177,53 @@ abstract class PropagateEmptyRelationBase extends
Rule[LogicalPlan] with CastSup
case _ => p
}
}
+
+ protected val repartitionTreePattern: TreePattern
+ protected def userSpecifiedRepartition(p: LogicalPlan): Boolean
+
+ /**
+ * Add a [[ROOT_REPARTITION]] tag for the root user-specified repartition so
this rule can
+ * skip optimize it.
+ */
+ protected def addTagForRootRepartition(plan: LogicalPlan): LogicalPlan = {
+ var isRootRepartition = true
+ plan.transformDownWithPruning(_.containsPattern(repartitionTreePattern)) {
+ case repartition if repartition.getTagValue(ROOT_REPARTITION).isEmpty &&
+ isRootRepartition && userSpecifiedRepartition(repartition) =>
+ repartition.setTagValue(ROOT_REPARTITION, ())
+ repartition
+ case project: Project => project
+ case filter: Filter => filter
+ case other if isRootRepartition =>
+ isRootRepartition = false
+ other
+ case other => other
+ }
+ }
}
/**
* This rule runs in the normal optimizer
*/
object PropagateEmptyRelation extends PropagateEmptyRelationBase {
- override def apply(plan: LogicalPlan): LogicalPlan =
plan.transformUpWithPruning(
- _.containsAnyPattern(LOCAL_RELATION, TRUE_OR_FALSE_LITERAL), ruleId) {
- commonApplyFunc
+ override protected val repartitionTreePattern: TreePattern =
REPARTITION_OPERATION
+
+ protected def userSpecifiedRepartition(p: LogicalPlan): Boolean = p match {
+ case _: Repartition => true
+ case r: RepartitionByExpression
+ if r.optNumPartitions.isDefined || r.partitionExpressions.nonEmpty =>
true
+ case _ => false
+ }
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
Review Comment:
is it possible to move the `apply` method to the base class, so that we can
share more code?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
