From d17effbe8c657a94ddd04769e294ec8962a104a3 Mon Sep 17 00:00:00 2001
From: ulysses-you <ulyssesyou18@gmail.com>
Date: Fri, 25 Jun 2021 20:48:51 +0800
Subject: [PATCH] [KYUUBI #715] [SQL] Check if need add repartition through
 project

<!--
Thanks for sending a pull request!

Here are some tips for you:
  1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html
  2. If the PR is related to an issue in https://github.com/NetEase/kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'.
  3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'.
-->

### _Why are the changes needed?_
<!--
Please clarify why the changes are needed. For instance,
  1. If you add a feature, you can talk about the use case of it.
  2. If you fix a bug, you can clarify why it is a bug.
-->
Avoid unnecessary shuffle

### _How was this patch tested?_
Add test

Closes #715 from ulysses-you/through-project.

Closes #715

5942320a [ulysses-you] through-project

Authored-by: ulysses-you <ulyssesyou18@gmail.com>
Signed-off-by: Cheng Pan <379377944@qq.com>
---
 .../main/scala/org/apache/kyuubi/sql/KyuubiAnalysis.scala    | 1 +
 .../scala/org/apache/spark/sql/KyuubiExtensionSuite.scala    | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/dev/kyuubi-extension-spark_3.1/src/main/scala/org/apache/kyuubi/sql/KyuubiAnalysis.scala b/dev/kyuubi-extension-spark_3.1/src/main/scala/org/apache/kyuubi/sql/KyuubiAnalysis.scala
index c9a1d0bb2..3d9b408d2 100644
--- a/dev/kyuubi-extension-spark_3.1/src/main/scala/org/apache/kyuubi/sql/KyuubiAnalysis.scala
+++ b/dev/kyuubi-extension-spark_3.1/src/main/scala/org/apache/kyuubi/sql/KyuubiAnalysis.scala
@@ -153,6 +153,7 @@ case class RepartitionBeforeWriteHive(session: SparkSession) extends Rule[Logica
 
 object RepartitionBeforeWriteHelper {
   def canInsertRepartitionByExpression(plan: LogicalPlan): Boolean = plan match {
+    case Project(_, child) => canInsertRepartitionByExpression(child)
     case Limit(_, _) => false
     case _: Sort => false
     case _: RepartitionByExpression => false
diff --git a/dev/kyuubi-extension-spark_3.1/src/test/scala/org/apache/spark/sql/KyuubiExtensionSuite.scala b/dev/kyuubi-extension-spark_3.1/src/test/scala/org/apache/spark/sql/KyuubiExtensionSuite.scala
index d65455fa6..9afe6c43c 100644
--- a/dev/kyuubi-extension-spark_3.1/src/test/scala/org/apache/spark/sql/KyuubiExtensionSuite.scala
+++ b/dev/kyuubi-extension-spark_3.1/src/test/scala/org/apache/spark/sql/KyuubiExtensionSuite.scala
@@ -110,6 +110,11 @@ class KyuubiExtensionSuite extends QueryTest with SQLTestUtils with AdaptiveSpar
           sql(s"CREATE TABLE tmp1 $storage PARTITIONED BY(c2) AS " +
             s"SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2)")
         }
+
+        withTable("tmp1") {
+          sql(s"CREATE TABLE tmp1 $storage PARTITIONED BY(c2) AS " +
+            s"SELECT * FROM VALUES(1, 'a'),(2, 'b') AS t(c1, c2) DISTRIBUTE BY rand()")
+        }
       }
     }
   }