[GitHub] [incubator-kyuubi] ulysses-you commented on a diff in pull request #3962: Add two conditions to decide if add shuffle before writing

[GitBox <gi...@apache.org>]

ulysses-you commented on code in PR #3962:
URL: https://github.com/apache/incubator-kyuubi/pull/3962#discussion_r1045583449


##########
extensions/spark/kyuubi-extension-spark-common/src/main/scala/org/apache/kyuubi/sql/RepartitionBeforeWritingBase.scala:
##########
@@ -108,14 +108,37 @@ abstract class RepartitionBeforeWritingHiveBase extends RepartitionBuilder {
   }
 }
 
-trait RepartitionBeforeWriteHelper {
-  def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = plan match {
-    case Project(_, child) => canInsertRepartitionByExpression(child)
-    case SubqueryAlias(_, child) => canInsertRepartitionByExpression(child)
-    case Limit(_, _) => false
-    case _: Sort => false
-    case _: RepartitionByExpression => false
-    case _: Repartition => false
-    case _ => true
+trait RepartitionBeforeWriteHelper extends Rule[LogicalPlan] {
+  private def hasBenefit(plan: LogicalPlan): Boolean = {
+    def probablyHasShuffle: Boolean = plan.find {
+      case _: Join => true
+      case _: Aggregate => true
+      case _: Distinct => true
+      case _: Deduplicate => true
+      case _: Window => true
+      case s: Sort if s.global => true
+      case _: RepartitionOperation => true
+      case _: GlobalLimit => true
+      case _ => false
+    }.isDefined
+
+    conf.getConf(KyuubiSQLConf.INSERT_REPARTITION_BEFORE_WRITE_IF_NO_SHUFFLE) || probablyHasShuffle
+  }
+
+  def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = {
+    def canInsert(p: LogicalPlan): Boolean = p match {
+      case Project(_, child) => canInsert(child)
+      case SubqueryAlias(_, child) => canInsert(child)
+      case Limit(_, _) => false
+      case _: Sort => false
+      case _: RepartitionByExpression => false
+      case _: Repartition => false
+      case _ => true
+    }
+
+    // 1. make sure AQE is enabled, otherwise it is no meaning to add a shuffle
+    // 2. make sure it does not break the semantics of original plan
+    // 3. try to reduce the performance regression if add a shuffle

Review Comment:
   addressed



-- 
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: notifications-unsubscribe@kyuubi.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@kyuubi.apache.org
For additional commands, e-mail: notifications-help@kyuubi.apache.org