Skip to content

[SPARK-24556][SQL] Always rewrite output partitioning in ReusedExchangeExec and InMemoryTableScanExec #21564

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
yucai wants to merge 8 commits into from
Closed

[SPARK-24556][SQL] Always rewrite output partitioning in ReusedExchangeExec and InMemoryTableScanExec #21564

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
f37139b
[SPARK-24556][SQL] ReusedExchange should rewrite output partitioning …
Jun 14, 2018
0ef99cc
improve
Jun 14, 2018
405ba94
improve import
Jun 14, 2018
9da85e0
rework unit tests
Jun 16, 2018
85dc1bc
upate tests
Jun 16, 2018
7d1c2f2
update tests
Jun 16, 2018
6744a9e
no need union in cache table
Jun 19, 2018
dcd0ce9
update tests
Jun 19, 2018
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
6 changes: 3 additions & 3 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning}
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.execution.{ColumnarBatchScan, LeafExecNode, SparkPlan, WholeStageCodegenExec}
import org.apache.spark.sql.execution.vectorized._
import org.apache.spark.sql.types._
Expand Down Expand Up @@ -169,8 +169,8 @@ case class InMemoryTableScanExec(
// But the cached version could alias output, so we need to replace output.
override def outputPartitioning: Partitioning = {
relation.cachedPlan.outputPartitioning match {
case h: HashPartitioning => updateAttribute(h).asInstanceOf[HashPartitioning]
case _ => relation.cachedPlan.outputPartitioning
case e: Expression => updateAttribute(e).asInstanceOf[Partitioning]
case other => other
}
}

Expand Down
4 changes: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/Exchange.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import org.apache.spark.broadcast
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, Expression, SortOrder}
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning}
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan, UnaryExecNode}
import org.apache.spark.sql.internal.SQLConf
Expand Down Expand Up @@ -70,7 +70,7 @@ case class ReusedExchangeExec(override val output: Seq[Attribute], child: Exchan
}

override def outputPartitioning: Partitioning = child.outputPartitioning match {
case h: HashPartitioning => h.copy(expressions = h.expressions.map(updateAttr))
case e: Expression => updateAttr(e).asInstanceOf[Partitioning]
case other => other

@viirya viirya Jun 16, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM

}

Expand Down
68 changes: 66 additions & 2 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,13 @@
package org.apache.spark.sql.execution

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{execution, Row}
import org.apache.spark.sql.{execution, DataFrame, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Range, Repartition, Sort, Union}
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.execution.columnar.InMemoryRelation
import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReusedExchangeExec, ReuseExchange, ShuffleExchangeExec}
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.functions._
Expand Down Expand Up @@ -686,6 +686,70 @@ class PlannerSuite extends SharedSQLContext {
Range(1, 2, 1, 1)))
df.queryExecution.executedPlan.execute()
}

test("SPARK-24556: always rewrite output partitioning in ReusedExchangeExec" +
"and InMemoryTableScanExec") {
def checkOutputPartitioningRewrite(
plans: Seq[SparkPlan],

@cloud-fan cloud-fan Jun 19, 2018

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.

now we can take a single spark plan

@yucai yucai Jun 19, 2018 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

How do you think if we merge check*OutputPartitioningRewrite together?

    def checkPlanAndOutputPartitioningRewrite(
        df: DataFrame,
        expectedPlanClass: Class[_],
        expectedPartitioningClass: Class[_]): Unit = {
      val plans = df.queryExecution.executedPlan.collect {
        case r: ReusedExchangeExec => r
        case m: InMemoryTableScanExec => m
      }
      assert(plans.size == 1)
      val plan = plans.head
      assert(plan.getClass == expectedPlanClass)
      val partitioning = plan.outputPartitioning
      assert(partitioning.getClass == expectedPartitioningClass)
      val partitionedAttrs = partitioning.asInstanceOf[Expression].references
      assert(partitionedAttrs.subsetOf(plan.outputSet))
    }

@yucai yucai Jun 19, 2018

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

@cloud-fan I still use Seq, so I can make checkReusedExchangeOutputPartitioningRewrite and checkInMemoryTableScanOutputPartitioningRewrite simpler. Kindly let me know if you have better idea.

expectedPartitioningClass: Class[_]): Unit = {
plans.foreach { plan =>
val partitioning = plan.outputPartitioning
assert(partitioning.getClass == expectedPartitioningClass)
val partitionedAttrs = partitioning.asInstanceOf[Expression].references
assert(partitionedAttrs.subsetOf(plan.outputSet))
}
}

def checkReusedExchangeOutputPartitioningRewrite(
df: DataFrame,
expectedPartitioningClass: Class[_]): Unit = {
val reusedExchange = df.queryExecution.executedPlan.collect {
case r: ReusedExchangeExec => r
}
assert(reusedExchange.size == 1)
checkOutputPartitioningRewrite(reusedExchange, expectedPartitioningClass)
}

def checkInMemoryTableScanOutputPartitioningRewrite(
df: DataFrame,
expectedPartitioningClass: Class[_]): Unit = {
val inMemoryScans = df.queryExecution.executedPlan.collect {
case m: InMemoryTableScanExec => m
}
assert(inMemoryScans.size == 2)
checkOutputPartitioningRewrite(inMemoryScans, expectedPartitioningClass)
}

// ReusedExchange is HashPartitioning
val df1 = Seq(1 -> "a").toDF("i", "j").repartition($"i")
val df2 = Seq(1 -> "a").toDF("i", "j").repartition($"i")
checkReusedExchangeOutputPartitioningRewrite(df1.union(df2), classOf[HashPartitioning])

// ReusedExchange is RangePartitioning
val df3 = Seq(1 -> "a").toDF("i", "j").orderBy($"i")
val df4 = Seq(1 -> "a").toDF("i", "j").orderBy($"i")
checkReusedExchangeOutputPartitioningRewrite(df3.union(df4), classOf[RangePartitioning])

// InMemoryTableScan is HashPartitioning
val df5 = df1.persist()
val df6 = df2.persist()
checkInMemoryTableScanOutputPartitioningRewrite(df5.union(df6), classOf[HashPartitioning])

@cloud-fan cloud-fan Jun 18, 2018

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.

why do we need to test table cache with union?

@yucai yucai Jun 19, 2018

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

I want to make sure both cache have the right output partitioning, so test the second cache table only?

@cloud-fan cloud-fan Jun 19, 2018

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.

union is used to trigger exchange reuse, but it's unnecessary to test cache.


// InMemoryTableScan is RangePartitioning
val df7 = spark.range(1, 100, 1, 10).toDF().persist()
val df8 = spark.range(1, 100, 1, 10).toDF().persist()
checkInMemoryTableScanOutputPartitioningRewrite(df7.union(df8), classOf[RangePartitioning])

// InMemoryTableScan is PartitioningCollection
withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
val df1 =
Seq(1 -> "a").toDF("i", "j").join(Seq(1 -> "a").toDF("m", "n"), $"i" === $"m").persist()
val df2 =
Seq(1 -> "a").toDF("i", "j").join(Seq(1 -> "a").toDF("m", "n"), $"i" === $"m").persist()
checkInMemoryTableScanOutputPartitioningRewrite(
df1.union(df2), classOf[PartitioningCollection])
}
}
}

// Used for unit-testing EnsureRequirements
Expand Down