[SPARK-38768][SQL] Remove Limit from plan if complete push down limit…

… to data source (#474) * [SPARK-38768][SQL] Remove `Limit` from plan if complete push down limit to data source Currently, Spark supports push down limit to data source. If limit could pushed down and Data source only have one partition, DS V2 still do limit again. This PR want remove `Limit` from plan if complete push down limit to data source. Improve performance. 'No'. New feature. Tests updated. Closes apache#36043 from beliefer/SPARK-38768. Authored-by: Jiaan Geng <beliefer@163.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com> * [SPARK-38391][SPARK-38768][SQL][FOLLOWUP] Add comments for `pushLimit` and `pushTopN` of `PushDownUtils` `pushLimit` and `pushTopN` of `PushDownUtils` returns tuple of boolean. It will be good to explain what the boolean value represents. Make DS V2 API more friendly to developers. 'No'. Just update comments. N/A Closes apache#36092 from beliefer/SPARK-38391_SPARK-38768_followup. Authored-by: Jiaan Geng <beliefer@163.com> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org> * [SPARK-37960][SQL][FOLLOWUP] Make the testing CASE WHEN query more reasonable Some testing CASE WHEN queries are not carefully written and do not make sense. In the future, the optimizer may get smarter and get rid of the CASE WHEN completely, and then we loose test coverage. This PR updates some CASE WHEN queries to make them more reasonable. future-proof test coverage. 'No'. N/A Closes apache#36125 from beliefer/SPARK-37960_followup3. Authored-by: Jiaan Geng <beliefer@163.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com> * update spark version Co-authored-by: Jiaan Geng <beliefer@163.com>
Kyligence · Aug 28, 2023 · 30a1edd · 30a1edd
1 parent ddadcf4
commit 30a1edd
Show file tree

Hide file tree

Showing 40 changed files with 90 additions and 55 deletions.
diff --git a/assembly/pom.xml b/assembly/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/common/kvstore/pom.xml b/common/kvstore/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/network-common/pom.xml b/common/network-common/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/network-shuffle/pom.xml b/common/network-shuffle/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/network-yarn/pom.xml b/common/network-yarn/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/sketch/pom.xml b/common/sketch/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/tags/pom.xml b/common/tags/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/common/unsafe/pom.xml b/common/unsafe/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/core/pom.xml b/core/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/examples/pom.xml b/examples/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/external/avro/pom.xml b/external/avro/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/docker-integration-tests/pom.xml b/external/docker-integration-tests/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kafka-0-10-assembly/pom.xml b/external/kafka-0-10-assembly/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kafka-0-10-sql/pom.xml b/external/kafka-0-10-sql/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kafka-0-10-token-provider/pom.xml b/external/kafka-0-10-token-provider/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kafka-0-10/pom.xml b/external/kafka-0-10/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kinesis-asl-assembly/pom.xml b/external/kinesis-asl-assembly/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/kinesis-asl/pom.xml b/external/kinesis-asl/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/external/spark-ganglia-lgpl/pom.xml b/external/spark-ganglia-lgpl/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/graphx/pom.xml b/graphx/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/hadoop-cloud/pom.xml b/hadoop-cloud/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/launcher/pom.xml b/launcher/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/mllib-local/pom.xml b/mllib-local/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/mllib/pom.xml b/mllib/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/pom.xml b/pom.xml
@@ -26,7 +26,7 @@
   </parent>
   <groupId>org.apache.spark</groupId>
   <artifactId>spark-parent_2.12</artifactId>
-  <version>3.2.0-kylin-4.x-r66</version>
+  <version>3.2.0-kylin-4.x-r69</version>
   <packaging>pom</packaging>
   <name>Spark Project Parent POM</name>
   <url>https://spark.apache.org/</url>

diff --git a/repl/pom.xml b/repl/pom.xml
@@ -21,7 +21,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../pom.xml</relativePath>
   </parent>
 

diff --git a/resource-managers/kubernetes/core/pom.xml b/resource-managers/kubernetes/core/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../../pom.xml</relativePath>
   </parent>
 

diff --git a/resource-managers/kubernetes/integration-tests/pom.xml b/resource-managers/kubernetes/integration-tests/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../../pom.xml</relativePath>
   </parent>
 

diff --git a/resource-managers/mesos/pom.xml b/resource-managers/mesos/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/resource-managers/yarn/pom.xml b/resource-managers/yarn/pom.xml
@@ -20,7 +20,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/sql/catalyst/pom.xml b/sql/catalyst/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownLimit.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownLimit.java
@@ -33,4 +33,10 @@ public interface SupportsPushDownLimit extends ScanBuilder {
    * Pushes down LIMIT to the data source.
    */
   boolean pushLimit(int limit);
+
+  /**
+   * Whether the LIMIT is partially pushed or not. If it returns true, then Spark will do LIMIT
+   * again. This method will only be called when {@link #pushLimit} returns true.
+   */
+  default boolean isPartiallyPushed() { return true; }
 }
diff --git a/sql/core/pom.xml b/sql/core/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.apache.spark</groupId>
     <artifactId>spark-parent_2.12</artifactId>
-    <version>3.2.0-kylin-4.x-r66</version>
+    <version>3.2.0-kylin-4.x-r69</version>
     <relativePath>../../pom.xml</relativePath>
   </parent>
 

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
@@ -122,11 +122,11 @@ object PushDownUtils extends PredicateHelper {
    *         the second Boolean value represents whether to push down partially, which means
    *         Spark will keep the Limit and do it again.
    */
-  def pushLimit(scanBuilder: ScanBuilder, limit: Int): Boolean = {
+  def pushLimit(scanBuilder: ScanBuilder, limit: Int): (Boolean, Boolean) = {
     scanBuilder match {
-      case s: SupportsPushDownLimit =>
-        s.pushLimit(limit)
-      case _ => false
+      case s: SupportsPushDownLimit if s.pushLimit(limit) =>
+        (true, s.isPartiallyPushed)
+      case _ => (false, false)
     }
   }
 

diff --git a/...src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/...src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
@@ -369,13 +369,13 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper wit
     }
   }
 
-  private def pushDownLimit(plan: LogicalPlan, limit: Int): LogicalPlan = plan match {
+  private def pushDownLimit(plan: LogicalPlan, limit: Int): (LogicalPlan, Boolean) = plan match {
     case operation @ ScanOperation(_, filter, sHolder: ScanBuilderHolder) if filter.isEmpty =>
-      val limitPushed = PushDownUtils.pushLimit(sHolder.builder, limit)
-      if (limitPushed) {
+      val (isPushed, isPartiallyPushed) = PushDownUtils.pushLimit(sHolder.builder, limit)
+      if (isPushed) {
         sHolder.pushedLimit = Some(limit)
       }
-      operation
+      (operation, isPushed && !isPartiallyPushed)
     case s @ Sort(order, _, operation @ ScanOperation(project, filter, sHolder: ScanBuilderHolder))
         if filter.isEmpty && CollapseProject.canCollapseExpressions(
           order, project, alwaysInline = true) =>
@@ -389,27 +389,32 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper wit
           sHolder.pushedLimit = Some(limit)
           sHolder.sortOrders = orders
           if (isPartiallyPushed) {
-            s
+            (s, false)
           } else {
-            operation
+            (operation, true)
           }
         } else {
-          s
+          (s, false)
         }
       } else {
-        s
+        (s, false)
       }
     case p: Project =>
-      val newChild = pushDownLimit(p.child, limit)
-      p.withNewChildren(Seq(newChild))
-    case other => other
+      val (newChild, isPartiallyPushed) = pushDownLimit(p.child, limit)
+      (p.withNewChildren(Seq(newChild)), isPartiallyPushed)
+    case other => (other, false)
   }
 
   def pushDownLimits(plan: LogicalPlan): LogicalPlan = plan.transform {
     case globalLimit @ Limit(IntegerLiteral(limitValue), child) =>
-      val newChild = pushDownLimit(child, limitValue)
-      val newLocalLimit = globalLimit.child.asInstanceOf[LocalLimit].withNewChildren(Seq(newChild))
-      globalLimit.withNewChildren(Seq(newLocalLimit))
+      val (newChild, canRemoveLimit) = pushDownLimit(child, limitValue)
+      if (canRemoveLimit) {
+        newChild
+      } else {
+        val newLocalLimit =
+          globalLimit.child.asInstanceOf[LocalLimit].withNewChildren(Seq(newChild))
+        globalLimit.withNewChildren(Seq(newLocalLimit))
+      }
   }
 
   private def getWrappedScan(