Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SPARK-18089][SQL] Remove shuffle codes in CollectLimitExec #15596

Closed
wants to merge 8 commits into from
Closed

[SPARK-18089][SQL] Remove shuffle codes in CollectLimitExec #15596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3d24f79
Remove shuffle codes in CollectLimitExec.
viirya Oct 21, 2016
76a3eaf
Fix test.
viirya Oct 24, 2016
58e8383
Remove CollectLimitExec.
viirya Oct 25, 2016
fbf4fd6
Merge remote-tracking branch 'upstream/master' into refactor-collectl…
viirya Oct 25, 2016
82ebff4
Polishing comment.
viirya Oct 25, 2016
44c64e0
fix test.
viirya Oct 25, 2016
86b4e42
Minimize the necessary changes.
viirya Oct 26, 2016
89c0c62
Revert CollectLimitExec back.
viirya Oct 27, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
execution.TakeOrderedAndProjectExec(
limit, order, projectList, planLater(child)) :: Nil
case logical.Limit(IntegerLiteral(limit), child) =>
execution.CollectLimitExec(limit, planLater(child)) :: Nil
execution.CollectLimitExec(
limit,
execution.LocalLimitExec(limit, planLater(child))) :: Nil
case other => planLater(other) :: Nil
}
case logical.Limit(IntegerLiteral(limit), logical.Sort(order, true, child)) =>
Expand Down
24 changes: 17 additions & 7 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,14 +35,24 @@ import org.apache.spark.util.Utils
case class CollectLimitExec(limit: Int, child: SparkPlan) extends UnaryExecNode {
override def output: Seq[Attribute] = child.output
override def outputPartitioning: Partitioning = SinglePartition
override def executeCollect(): Array[InternalRow] = child.executeTake(limit)
private val serializer: Serializer = new UnsafeRowSerializer(child.output.size)
override def requiredChildDistribution: List[Distribution] = AllTuples :: Nil
override def executeCollect(): Array[InternalRow] = child match {
// Shuffling injected. WholeStageCodegenExec enabled.
case ShuffleExchange(_, WholeStageCodegenExec(l: LocalLimitExec), _) =>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are removing an optimization here right? We can greatly reduce the number of shuffled records by applying the limit before anything gets shuffled.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I use a LocalLimit to do this optimization. I think we should use existing physical plans as much as possible, instead of rdd manipulation.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One advantage is LocalLimit supports whole stage codegen. We can also easily get the idea of this optimization from physical plans.

l.child.executeTake(limit)

// Shuffling injected. WholeStageCodegenExec disabled.
case ShuffleExchange(_, l: LocalLimitExec, _) => l.child.executeTake(limit)

// No shuffled injected. WholeStageCodegenExec enabled.
case WholeStageCodegenExec(l: LocalLimitExec) => l.child.executeTake(limit)

// No shuffling injected. WholeStageCodegenExec disabled.
case l: LocalLimitExec => l.child.executeTake(limit)
}

protected override def doExecute(): RDD[InternalRow] = {
val locallyLimited = child.execute().mapPartitionsInternal(_.take(limit))
val shuffled = new ShuffledRowRDD(
ShuffleExchange.prepareShuffleDependency(
locallyLimited, child.output, SinglePartition, serializer))
shuffled.mapPartitionsInternal(_.take(limit))
child.execute().mapPartitionsInternal(_.take(limit))
}
}

Expand Down
16 changes: 12 additions & 4 deletions sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.SortOrder
import org.apache.spark.sql.catalyst.plans.logical.Aggregate
import org.apache.spark.sql.catalyst.util.StringUtils
import org.apache.spark.sql.execution.aggregate
import org.apache.spark.sql.execution.LocalLimitExec
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, CartesianProductExec, SortMergeJoinExec}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
Expand Down Expand Up @@ -2684,11 +2685,18 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
}

test("SPARK-17515: CollectLimit.execute() should perform per-partition limits") {
val df = spark.range(1, 100, 1, numPartitions = 10).limit(1)
val localLimit = df.queryExecution.executedPlan.collect {
case l: LocalLimitExec => l
}
assert(localLimit.nonEmpty)
val numRecordsRead = spark.sparkContext.longAccumulator
spark.range(1, 100, 1, numPartitions = 10).map { x =>
numRecordsRead.add(1)
x
}.limit(1).queryExecution.toRdd.count()
localLimit.head.execute().mapPartitionsInternal { iter =>
iter.map { x =>
numRecordsRead.add(1)
x
}
}.count
assert(numRecordsRead.value === 10)
}

Expand Down
14 changes: 9 additions & 5 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,11 +95,15 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
assert(metrics1.contains("numOutputRows"))
assert(metrics1("numOutputRows").value === 3)

val df2 = spark.createDataset(Seq(1, 2, 3)).limit(2)
df2.collect()
val metrics2 = df2.queryExecution.executedPlan.collectLeaves().head.metrics
assert(metrics2.contains("numOutputRows"))
assert(metrics2("numOutputRows").value === 2)
Seq("true", "false").map { codeGen =>
withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codeGen) {
val df2 = spark.createDataset(Seq(1, 2, 3)).coalesce(1).limit(2)
assert(df2.collect().length === 2)
val metrics2 = df2.queryExecution.executedPlan.collectLeaves().head.metrics
assert(metrics2.contains("numOutputRows"))
assert(metrics2("numOutputRows").value === 2)
}
}
}

test("Filter metrics") {
Expand Down