Skip to content

[SPARK-32932][SQL] Do not use local shuffle reader at final stage on write command #29797

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

Closed
manuzhang wants to merge 7 commits into from
Closed

[SPARK-32932][SQL] Do not use local shuffle reader at final stage on write command #29797

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
84134b0
[SPARK-32932][SQL] Do not use local shuffle reader at final stage on …
manuzhang Sep 18, 2020
b112133
Support V2TableWriteExec
manuzhang Oct 9, 2020
f05b458
Update tests
manuzhang Oct 12, 2020
4fdee62
Remove unused imports
manuzhang Oct 12, 2020
39dc4ca
Address comments
manuzhang Oct 14, 2020
364ee28
Use NoopDataSource to test DataSource v2
manuzhang Oct 14, 2020
d391269
Fix style
manuzhang Oct 14, 2020
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
11 changes: 10 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule}
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec._
import org.apache.spark.sql.execution.command.DataWritingCommandExec
import org.apache.spark.sql.execution.exchange._
import org.apache.spark.sql.execution.ui.{SparkListenerSQLAdaptiveExecutionUpdate, SparkListenerSQLAdaptiveSQLMetricUpdates, SQLPlanMetric}
import org.apache.spark.sql.internal.SQLConf
Expand Down Expand Up @@ -102,6 +103,14 @@ case class AdaptiveSparkPlanExec(
OptimizeLocalShuffleReader(conf)
)

@transient private val finalStageOptimizerRules: Seq[Rule[SparkPlan]] =

@cloud-fan cloud-fan Oct 12, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

it's only called once, can be a def

context.qe.sparkPlan match {
case _: DataWritingCommandExec =>

@cloud-fan cloud-fan Sep 29, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

we need to match all writing commands, including DS v1, v2 and file source. Maybe we can create a tagging trait like UserDefinedExpression, to tag these writing commands.

@manuzhang manuzhang Oct 1, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

It seems DSv2 is not ready for write as per https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala#L988-L994.

Meanwhile, will it too big a change for those interfaces to extend the tagging trait ?

@cloud-fan cloud-fan Oct 6, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

File source v2 is not ready yet, but it doesn't mean DS v2 is not ready for writing. Please follow InsertAdaptiveSparkPlan.applyInternal

@manuzhang manuzhang Oct 6, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Sure, is there a UT for DS v2 write ? I find only V1 is used for write no matter format I specify.

@cloud-fan cloud-fan Oct 7, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

See DataSourceV2Suite, we have testing v2 source that support writing.

@manuzhang manuzhang Oct 9, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I've added match for V2TableWriteExec and UT. Not sure I've got it right as v2 API is quite new to me. Please help review. Thanks.

queryStageOptimizerRules.filterNot(_.isInstanceOf[OptimizeLocalShuffleReader])
Comment thread
dongjoon-hyun marked this conversation as resolved.
case _ =>
queryStageOptimizerRules
}

// A list of physical optimizer rules to be applied right after a new stage is created. The input
// plan to these rules has exchange as its root node.
@transient private val postStageCreationRules = Seq(
Expand Down Expand Up @@ -235,7 +244,7 @@ case class AdaptiveSparkPlanExec(
// Run the final plan when there's no more unfinished stages.
currentPhysicalPlan = applyPhysicalRules(
result.newPlan,
queryStageOptimizerRules ++ postStageCreationRules,
finalStageOptimizerRules ++ postStageCreationRules,
Some((planChangeLogger, "AQE Final Query Stage Optimization")))
isFinalPlan = true
executionId.foreach(onUpdatePlan(_, Seq(currentPhysicalPlan)))
Expand Down
22 changes: 22 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec
import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{IntegerType, StructType}
import org.apache.spark.util.Utils
Expand Down Expand Up @@ -1258,4 +1259,25 @@ class AdaptiveQueryExecSuite
}
}
}

test("SPARK-32932: Do not use local shuffle reader at final stage on DataWritingCommand") {
withSQLConf(SQLConf.PARTITION_OVERWRITE_MODE.key -> PartitionOverwriteMode.DYNAMIC.toString,
SQLConf.SHUFFLE_PARTITIONS.key -> "5",
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true") {
withTable("t") {
val data = for (
i <- 1 to 10;
j <- 1 to 3
) yield (i, j)
data.toDF("a", "b")
.repartition($"b")
.write
.partitionBy("b")
.mode("overwrite")
.saveAsTable("t")

@cloud-fan cloud-fan Oct 12, 2020

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

ditto: put in one line

assert(spark.read.table("t").inputFiles.length == 3)

@cloud-fan cloud-fan Oct 12, 2020 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

it's a bit tricky to check the number of files. 3 distinct values don't always mean 3 files.

Can we use QueryExecutionListener to catch the query plan, and check there is no local shuffle reader?

@manuzhang manuzhang Oct 12, 2020

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Tests have been updated with QueryExecutionListener to check local shuffle reader. Please help review again.

}
}
}
}