[SPARK-18079] [SQL] CollectLimitExec.executeToIterator should perform…

… per-partition limits

We have an internal product that needs this. See upstream PR [1].

[1] apache#15614

Co-authored-by: Patrick Woody <pwoody@palantir.com>
Co-authored-by: Josh Casale <jcasale@palantir.com>
Co-authored-by: Will Raschkowski <wraschkowski@palantir.com>

FORK.md

@@ -1,4 +1,5 @@
 # Difference with upstream
+* [SPARK-18079](https://issues.apache.org/jira/browse/SPARK-18079) - CollectLimitExec.executeToIterator should perform per-partition limits
 * [SPARK-20952](https://issues.apache.org/jira/browse/SPARK-20952) - ParquetFileFormat should forward TaskContext to its forkjoinpool
 * [SPARK-26626](https://issues.apache.org/jira/browse/SPARK-26626) - Limited the maximum size of repeatedly substituted aliases
 * [SPARK-25200](https://issues.apache.org/jira/browse/SPARK-25200) - Allow setting HADOOP_CONF_DIR as a spark config

sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala

@@ -45,6 +45,12 @@ case class CollectLimitExec(limit: Int, child: SparkPlan) extends LimitExec {
   override def output: Seq[Attribute] = child.output
   override def outputPartitioning: Partitioning = SinglePartition
   override def executeCollect(): Array[InternalRow] = child.executeTake(limit)
+
+  // TODO(palantir): Reopen upstream PR for SPARK-18079; we need this for an internal product (DP)
+  override def executeToIterator(): Iterator[InternalRow] = {
+    LocalLimitExec(limit, child).executeToIterator().take(limit)
+  }
+
   private val serializer: Serializer = new UnsafeRowSerializer(child.output.size)
   private lazy val writeMetrics =
     SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext)