Skip to content

[SPARK-46609][SQL] Avoid exponential explosion in PartitioningPreservingUnaryExecNode #44614

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
wants to merge 1 commit into from
Closed

[SPARK-46609][SQL] Avoid exponential explosion in PartitioningPreservingUnaryExecNode #44614

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
with AliasAwareOutputExpression {
final override def outputPartitioning: Partitioning = {
val partitionings: Seq[Partitioning] = if (hasAlias) {
flattenPartitioning(child.outputPartitioning).flatMap {
flattenPartitioning(child.outputPartitioning).iterator.flatMap {
Copy link
Contributor Author

@cloud-fan cloud-fan Jan 5, 2024

Choose a reason for hiding this comment

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

We apply the threshold when expanding each partitioning inside PartitioningCollection. But if there are many partitionings, we don't limit anything and may produce a lot of partitionings.

case e: Expression =>
// We need unique partitionings but if the input partitioning is
// `HashPartitioning(Seq(id + id))` and we have `id -> a` and `id -> b` aliases then after
Expand All @@ -44,7 +44,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
.take(aliasCandidateLimit)
.asInstanceOf[LazyList[Partitioning]]
case o => Seq(o)
}
}.take(aliasCandidateLimit).toSeq
} else {
// Filter valid partitiongs (only reference output attributes of the current plan node)
val outputSet = AttributeSet(outputExpressions.map(_.toAttribute))
Expand Down
30 changes: 28 additions & 2 deletions ...src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,14 @@

package org.apache.spark.sql.execution

import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection, UnknownPartitioning}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference}
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.StringType

class ProjectedOrderingAndPartitioningSuite
extends SharedSparkSession with AdaptiveSparkPlanHelper {
Expand Down Expand Up @@ -101,6 +104,22 @@ class ProjectedOrderingAndPartitioningSuite
}
}

test("SPARK-46609: Avoid exponential explosion in PartitioningPreservingUnaryExecNode") {
withSQLConf(SQLConf.EXPRESSION_PROJECTION_CANDIDATE_LIMIT.key -> "2") {
val output = Seq(AttributeReference("a", StringType)(), AttributeReference("b", StringType)())
val plan = ProjectExec(
Seq(
Alias(output(0), "a1")(),
Alias(output(0), "a2")(),
Alias(output(1), "b1")(),
Alias(output(1), "b2")()
),
DummyLeafPlanExec(output)
)
assert(plan.outputPartitioning.asInstanceOf[PartitioningCollection].partitionings.length == 2)
}
}

test("SPARK-42049: Improve AliasAwareOutputExpression - multi-references to complex " +
"expressions") {
val df2 = spark.range(2).repartition($"id" + $"id").selectExpr("id + id as a", "id + id as b")
Expand Down Expand Up @@ -192,3 +211,10 @@ class ProjectedOrderingAndPartitioningSuite
assert(outputOrdering.head.sameOrderExpressions.size == 0)
}
}

private case class DummyLeafPlanExec(output: Seq[Attribute]) extends LeafExecNode {
override protected def doExecute(): RDD[InternalRow] = null
override def outputPartitioning: Partitioning = {
PartitioningCollection(output.map(attr => HashPartitioning(Seq(attr), 4)))
}
}