[SPARK-46170][SQL] Support inject adaptive query post planner strategy rules in SparkSessionExtensions #44074

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

Jump to bottom

Closed

ulysses-you wants to merge 2 commits into apache:master from ulysses-you:post-planner

Contributor

ulysses-you commented Nov 29, 2023

What changes were proposed in this pull request?

This pr adds a new extension entrance queryPostPlannerStrategyRules in SparkSessionExtensions. It will be applied between plannerStrategy and queryStagePrepRules in AQE, so it can get the whole plan before injecting exchanges.

Why are the changes needed?

a part of #44013

Does this PR introduce any user-facing change?

no, only for develop

How was this patch tested?

add test

Was this patch authored or co-authored using generative AI tooling?

no


          Support inject adaptive query post planner strategy rules in SparkSes…

7ef5e30

…sionExtensions

github-actions bot added the SQL label

Contributor Author

ulysses-you commented Nov 29, 2023

cc @dongjoon-hyun @cloud-fan this pr is only for adding Adaptive Query Post Planner Strategy Rules part.

Contributor Author

ulysses-you commented Nov 29, 2023

CI link https://github.com/ulysses-you/spark/actions/runs/7029365126

dongjoon-hyun reviewed

View reviewed changes

Member

dongjoon-hyun left a comment

Thank you for making a PR, @ulysses-you .

dongjoon-hyun reviewed

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala Outdated

    
                type ColumnarRuleBuilder = SparkSession => ColumnarRule

                type QueryStagePrepRuleBuilder = SparkSession => Rule[SparkPlan]

                type QueryStageOptimizerRuleBuilder = SparkSession => Rule[SparkPlan]

                type QueryPostPlannerStrategyBuilder = SparkSession => Rule[SparkPlan]

Member

dongjoon-hyun Nov 29, 2023 •

edited

Loading

Shall we put this before type QueryStagePrepRuleBuilder because this rule is supposed to be used before queryStagePreparationRules?

dongjoon-hyun reviewed

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala Outdated

    
                private[this] val queryStageOptimizerRuleBuilders =

                  mutable.Buffer.empty[QueryStageOptimizerRuleBuilder]

                private[this] val queryPostPlannerStrategyRuleBuilders =

                  mutable.Buffer.empty[QueryPostPlannerStrategyBuilder]

Member

dongjoon-hyun Nov 29, 2023

Ditto. Shall we put this before private[this] val queryStagePrepRuleBuilders?

dongjoon-hyun reviewed

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala

    
                    session: SparkSession): Seq[Rule[SparkPlan]] = {

                  queryPostPlannerStrategyRuleBuilders.map(_.apply(session)).toSeq

                }

Member

dongjoon-hyun Nov 29, 2023

ditto. Shall we put this before private[sql] def buildQueryStagePrepRules?

dongjoon-hyun reviewed

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala Outdated

    
                def injectQueryPostPlannerStrategyRule(builder: QueryPostPlannerStrategyBuilder): Unit = {

                  queryPostPlannerStrategyRuleBuilders += builder

                }

Member

dongjoon-hyun Nov 29, 2023

ditto. Shall we put this before def injectQueryStagePrepRule?

dongjoon-hyun approved these changes

View reviewed changes

Member

dongjoon-hyun left a comment •

edited

Loading

+1, it looks reasonable to me. I left a few minor comments. Thank you, @ulysses-you .

Also, cc @sunchao , too

LuciferYang reviewed

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala Outdated

    
                  withSession(extensions) { session =>

                    assert(session.sessionState.adaptiveRulesHolder.queryPostPlannerStrategyRules

                      .contains(MyQueryPostPlannerStrategyRule))

                    import session.sqlContext.implicits._

Contributor

LuciferYang Nov 29, 2023

how about just use session.implicits._?

LuciferYang reviewed

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala Outdated

    
                    import session.sqlContext.implicits._

                    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",

                        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") {

                      val input = Seq((10), (20), (10)).toDF("c1")

Contributor

LuciferYang Nov 29, 2023

Why is it necessary to add parentheses to each element? for readability?

LuciferYang reviewed

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala

    
                        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") {

                      val input = Seq((10), (20), (10)).toDF("c1")

                      val df = input.groupBy("c1").count()

                      df.collect()

Contributor

LuciferYang Nov 29, 2023

I tried it, and even if df.collect() is not executed, this test case can still pass. So, is it necessary?

Contributor Author

ulysses-you Nov 30, 2023

it is used to make sure we are checking the final plan

LuciferYang reviewed

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala

    
                      df.collect()

                      assert(df.rdd.partitions.length == 1)

                      assert(find(df.queryExecution.executedPlan) {

                        case s: ShuffleExchangeExec if s.outputPartitioning == SinglePartition => true

Contributor

LuciferYang Nov 29, 2023

How about case s: ShuffleExchangeExec => s.outputPartitioning == SinglePartition

Additionally, a personal opinion unrelated to this pr: If there is an exists function in AdaptiveSparkPlanHelper, would this assertion be simpler to write?

LuciferYang reviewed

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala Outdated

    
                        case s: ShuffleExchangeExec if s.outputPartitioning == SinglePartition => true

                        case _ => false

                      }.isDefined)

                      assert(find(df.queryExecution.executedPlan) {

Contributor

LuciferYang Nov 29, 2023

Seems using collectFirst could eliminate the case _ => false branch?


          address comments

93a727a

Contributor Author

ulysses-you commented Nov 30, 2023

thank you @dongjoon-hyun @LuciferYang , addressed comments

LuciferYang approved these changes

View reviewed changes

ulysses-you closed this in

9bb358b

Contributor Author

ulysses-you commented Nov 30, 2023

thanks for review, merged to master.

ulysses-you deleted the post-planner branch

November 30, 2023 09:15

ulysses-you added a commit to ulysses-you/spark that referenced this pull request


          [SPARK-46170][SQL] Support inject adaptive query post planner strateg…

e842f97

…y rules in SparkSessionExtensions

### What changes were proposed in this pull request?

This pr adds a new extension entrance `queryPostPlannerStrategyRules` in `SparkSessionExtensions`. It will be applied between plannerStrategy and queryStagePrepRules in AQE, so it can get the whole plan before injecting exchanges.

### Why are the changes needed?

a part of apache#44013

### Does this PR introduce _any_ user-facing change?

no, only for develop

### How was this patch tested?

add test

### Was this patch authored or co-authored using generative AI tooling?

no

Closes apache#44074 from ulysses-you/post-planner.

Authored-by: ulysses-you <[email protected]>
Signed-off-by: youxiduo <[email protected]>

ulysses-you mentioned this pull request

[SPARK-46170][SQL][3.5] Support inject adaptive query post planner strategy rules in SparkSessionExtensions #45037

Closed

yaooqinn pushed a commit that referenced this pull request


          [SPARK-46170][SQL][3.5] Support inject adaptive query post planner st…

1b36e3c

…rategy rules in SparkSessionExtensions

This pr is backport #44074 for branch-3.5 since 3.5 is a lts version

### What changes were proposed in this pull request?

This pr adds a new extension entrance `queryPostPlannerStrategyRules` in `SparkSessionExtensions`. It will be applied between plannerStrategy and queryStagePrepRules in AQE, so it can get the whole plan before injecting exchanges.

### Why are the changes needed?

3.5 is a lts version

### Does this PR introduce _any_ user-facing change?

no, only for develop

### How was this patch tested?

add test

### Was this patch authored or co-authored using generative AI tooling?

no

Closes #44074 from ulysses-you/post-planner.

Authored-by: ulysses-you <ulyssesyou18gmail.com>

Closes #45037 from ulysses-you/SPARK-46170.

Authored-by: ulysses-you <[email protected]>
Signed-off-by: Kent Yao <[email protected]>

turboFei pushed a commit to turboFei/spark that referenced this pull request


          [SPARK-46170][SQL][3.5] Support inject adaptive query post planner st…

1bc00e8

…rategy rules in SparkSessionExtensions (apache#366)

This pr is backport apache#44074 for branch-3.5 since 3.5 is a lts version

### What changes were proposed in this pull request?

This pr adds a new extension entrance `queryPostPlannerStrategyRules` in `SparkSessionExtensions`. It will be applied between plannerStrategy and queryStagePrepRules in AQE, so it can get the whole plan before injecting exchanges.

### Why are the changes needed?

3.5 is a lts version

### Does this PR introduce _any_ user-facing change?

no, only for develop

### How was this patch tested?

add test

### Was this patch authored or co-authored using generative AI tooling?

no

Closes apache#44074 from ulysses-you/post-planner.

Authored-by: ulysses-you <ulyssesyou18gmail.com>

Closes apache#45037 from ulysses-you/SPARK-46170.

Authored-by: ulysses-you <[email protected]>

Signed-off-by: Kent Yao <[email protected]>
Co-authored-by: ulysses-you <[email protected]>

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

SQL