ulysses-you commented on code in PR #59630:
URL: https://github.com/apache/spark/pull/59630#discussion_r4422047096
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala:
##########
@@ -707,8 +708,11 @@ class CacheManager extends Logging with
AdaptiveSparkPlanHelper {
// extra shuffle
disableConfigs =
disableConfigs :+
SQLConf.ADAPTIVE_EXECUTION_APPLY_FINAL_STAGE_SHUFFLE_OPTIMIZATIONS
+ SparkSession.getOrCloneSessionWithConfigsOff(session, disableConfigs,
+ Seq(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED))
Review Comment:
Updated the canChangeCachedPlanOutputPartitioning documentation in
0a98c1e6a70. It now names union shuffle push-down and clarifies that false
disables both optimizations for cache materialization, while true retains their
configured enablement.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala:
##########
@@ -692,7 +692,8 @@ class CacheManager extends Logging with
AdaptiveSparkPlanHelper {
/**
* If `CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING` is enabled, return the
session with disabled
- * `ADAPTIVE_EXECUTION_APPLY_FINAL_STAGE_SHUFFLE_OPTIMIZATIONS`.
+ * `ADAPTIVE_EXECUTION_APPLY_FINAL_STAGE_SHUFFLE_OPTIMIZATIONS` and
+ * `PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED`.
Review Comment:
Corrected "is enabled" to "is disabled" in 0a98c1e6a70, matching the
condition below. The inversion predates this PR; the correction also covers the
newly listed push-down setting.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/PushDownShuffleExchangeThroughUnionSuite.scala:
##########
@@ -0,0 +1,2496 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.adaptive
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{
+ Alias, Ascending, Attribute, AttributeReference, ExprId, GreaterThan,
IsNotNull, Literal, Rand,
+ RowNumber, SortOrder}
+import org.apache.spark.sql.catalyst.optimizer.BuildLeft
+import org.apache.spark.sql.catalyst.plans.physical.{
+ ClusteredDistribution, CoalescedBoundary, CoalescedHashPartitioning,
Distribution,
+ HashPartitioning, NullAwareHashPartitioning, Partitioning,
PartitioningCollection,
+ RoundRobinPartitioning,
+ UnknownPartitioning}
+import org.apache.spark.sql.execution.{
+ FilterExec, InputAdapter, LocalLimitExec, ProjectExec, QueryExecution,
ScalarSubquery, SortExec,
+ SparkPlan, StampUnionDecisions, SubqueryExec, UnaryExecNode,
UnionConfSnapshot, UnionExec,
+ WholeStageCodegenExec}
+import org.apache.spark.sql.execution.exchange.{
+ ENSURE_REQUIREMENTS, EnsureRequirements, REPARTITION_BY_COL,
REPARTITION_BY_NUM,
+ ShuffleExchangeExec, ValidateRequirements}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.window.{Final => FinalWindowGroupLimit,
Partial => PartialWindowGroupLimit, WindowGroupLimitExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.IntegerType
+
+class PushDownShuffleExchangeThroughUnionSuite extends QueryTest with
SharedSparkSession {
+ import testImplicits._
+
+ // The suite exercises the rule itself, so it is on for every test whatever
the default is.
+ override protected def sparkConf =
+ super.sparkConf
+ .set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key, "true")
+
.set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_JOIN_INPUTS_ENABLED.key,
"true")
+
.set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_JOIN_BRANCHES_ENABLED.key,
"true")
+
.set(SQLConf.PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_WITH_SUBQUERIES_ENABLED.key,
"true")
Review Comment:
Added a subquery-free arm under allowWithSubqueries=false in 0a98c1e6a70. It
checks that the context flag remains false and the union exchange is removed.
The scalar-subquery arm still checks the shared context and retained exchange.
Unconditionally marking queries as containing subqueries makes the new arm fail.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -8526,6 +8526,134 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED =
+ buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.enabled")
+ .internal()
+ .doc("When set to true, during AQE query stage preparation a
ShuffleExchangeExec with " +
+ "hash partitioning added by EnsureRequirements on top of a UnionExec
is pushed down " +
+ "below the union, shuffling only the children that do not already
report the picked " +
+ "layout. The target count must be at least the smaller of the replaced
exchange's " +
+ "count and the effective initial shuffle partition count, including " +
+ "spark.sql.adaptive.coalescePartitions.initialPartitionNum when set. "
+
+ "Removing the union's direct shuffle stage " +
+ "can reduce AQE skew handling and join conversion opportunities.
Keeping a shuffled " +
+ "join child's layout can make skew join optimization or local shuffle
reads invalidate " +
+ "an ancestor's distribution. AQE may discard that optimization or add
another shuffle " +
+ "when skew optimization is forced. Join inputs and join branches are
opt-in. Requires " +
+ "spark.sql.unionOutputPartitioning to be enabled. Adaptive cost
compensation applies " +
+ "only to SimpleCostEvaluator. If
spark.sql.adaptive.customCostEvaluatorClass is set, " +
+ "this rule is disabled unless explicitly enabled. With a custom
evaluator, adapt its " +
+ "cost model to the replacement of one union exchange by multiple
branch " +
+ "exchanges; otherwise the increased exchange count may cause AQE to
reject a re-plan " +
+ "and its other runtime optimizations.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createOptional
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_LOOK_THROUGH_OPERATORS_ENABLED =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.lookThroughOperators.enabled")
+ .internal()
+ .doc("When set to true, the shuffle push-down through a union may look
through the " +
+ "operators between the shuffle and the union: a project, a local sort,
a filter, a " +
+ "partial aggregate and a partial window group limit. A filter, a
partial aggregate and " +
+ "a partial window group limit reduce the data on their way, so the
child is then " +
+ "shuffled before the reduction, which can send more data than the
single shuffle it " +
+ "replaces. Requires " +
+ s"${PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key} to be enabled.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_NON_COALESCIBLE_BRANCHES_ENABLED =
+ buildConf(
+
"spark.sql.optimizer.pushDownShuffleThroughUnion.keepNonCoalescibleBranches.enabled")
+ .internal()
+ .doc("When set to true, fixed-count shuffles, bucketed or cached scans,
and branches " +
+ "containing exploding joins may offer push-down targets. This can
prevent coalescing " +
+ "of the consumer's compatible shuffle group or reduce its target
partition size. " +
+ "Rebalance shuffles remain ineligible. This only affects target " +
+ "selection when spark.sql.adaptive.coalescePartitions.enabled is on.
Requires " +
+ s"${PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key} to be enabled.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_MAX_MOVED_BRANCHES =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.maxMovedBranches")
+ .internal()
+ .doc("The maximum number of union children to move onto a push-down
target. Existing " +
+ "unshared requirement or column-repartition exchanges may be
retargeted; other moved " +
+ "children get a new exchange. Child count and order are preserved.
Targets requiring " +
+ "more moved children are skipped. Zero only allows targets " +
+ "already offered by every child. Requires " +
+ s"${PUSH_DOWN_SHUFFLE_THROUGH_UNION_ENABLED.key} to be enabled.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .intConf
+ .checkValue(_ >= 0, "The maximum number of moved branches must be
non-negative.")
+ .createWithDefault(3)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_MAX_UNION_CHILDREN =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.maxUnionChildren")
+ .internal()
+ .doc("The maximum total number of union children eligible for shuffle
push-down, " +
+ "including kept and moved children. Zero disables push-down. This
bounds the number " +
+ "of branch tails executed sequentially within each union output task.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .intConf
+ .checkValue(_ >= 0, "The maximum number of union children must be
non-negative.")
+ .createWithDefault(10)
+
+ val
PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_PARTITION_KEYS_SUBSET_OF_CLUSTER_KEYS_ENABLED
=
+ buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion." +
+ "allowPartitionKeysSubsetOfClusterKeys.enabled")
+ .internal()
+ .doc("When true, shuffle push-down may use partition keys that are a
subset of the " +
+ "consumer's clustering keys if
spark.sql.requireAllClusterKeysForCoPartition is false " +
+ "and the consumer accepts that distribution. Such targets can
concentrate rows in " +
+ "fewer partitions. When false, targets must match all clustering
keys.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_WITH_SUBQUERIES_ENABLED =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.allowWithSubqueries.enabled")
+ .internal()
+ .doc("When true, allows shuffle push-down in queries containing scalar,
IN or dynamic " +
+ "pruning subqueries. Independent adaptive plans can make different
rewrite decisions, " +
+ "reducing exchange reuse or disabling broadcast-only dynamic partition
pruning. " +
+ "When false, push-down is disabled for the main query and all its
subqueries.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_ALLOW_JOIN_INPUTS_ENABLED =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.allowJoinInputs.enabled")
+ .internal()
+ .doc("When true, allows shuffle push-down for inputs of shuffled joins.
Removing a " +
+ "direct shuffle stage can prevent adaptive skew join optimization or
join conversion.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_DOWN_SHUFFLE_THROUGH_UNION_KEEP_JOIN_BRANCHES_ENABLED =
+
buildConf("spark.sql.optimizer.pushDownShuffleThroughUnion.keepJoinBranches.enabled")
+ .internal()
+ .doc("When true, union children may offer layouts produced by shuffled
joins above their " +
Review Comment:
Updated the configuration documentation in 0a98c1e6a70 to include
broadcast-hash-join probes backed by requirement shuffles, including stages and
adaptive reads. User column-repartition-backed probes retain their existing
eligibility.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
