You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2022/05/30 06:39:08 UTC
[spark] branch master updated: [SPARK-37523][SQL] Re-optimize partitions in Distribution and Ordering if numPartitions is not specified
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 098523c8c08 [SPARK-37523][SQL] Re-optimize partitions in Distribution and Ordering if numPartitions is not specified
098523c8c08 is described below
commit 098523c8c08d03c0b8d9b970314e538c1c7573fc
Author: huaxingao <hu...@apple.com>
AuthorDate: Mon May 30 14:38:45 2022 +0800
[SPARK-37523][SQL] Re-optimize partitions in Distribution and Ordering if numPartitions is not specified
### What changes were proposed in this pull request?
Re-optimize partitions in Distribution and Ordering if numPartitions is not specified.
Support both strictly required distribution and best effort distribution. For best effort distribution, if user doesn't request a specific number of partitions, Spark will split skewed partitions and coalesce small partitions (`RebalancePartitions` will be used and both `OptimizeSkewInRebalancePartitions` and `CoalesceShufflePartitions` will be applied). For strictly required distribution, if user doesn't request a specific number of partitions, Spark will coalesce small partitions b [...]
### Why are the changes needed?
optimization
### Does this PR introduce _any_ user-facing change?
In interface `RequiresDistributionAndOrdering`, a new API `distributionStrictlyRequired()` will be added to tell if the distribution required by this write is strictly required or best effort only
```
default boolean distributionStrictlyRequired() { return true; }
```
### How was this patch tested?
Existing and new tests
Closes #36644 from huaxingao/skew2.
Authored-by: huaxingao <hu...@apple.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../write/RequiresDistributionAndOrdering.java | 12 +
.../plans/logical/basicLogicalOperators.scala | 64 ++--
.../sql/connector/catalog/InMemoryTable.scala | 5 +-
.../connector/catalog/InMemoryTableCatalog.scala | 5 +-
.../v2/DistributionAndOrderingUtils.scala | 20 +-
.../sql/execution/datasources/v2/V2Writes.scala | 10 +-
.../WriteDistributionAndOrderingSuite.scala | 355 ++++++++++++++++-----
7 files changed, 348 insertions(+), 123 deletions(-)
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RequiresDistributionAndOrdering.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RequiresDistributionAndOrdering.java
index 983e6b0fffb..a3d08338c7a 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RequiresDistributionAndOrdering.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RequiresDistributionAndOrdering.java
@@ -47,6 +47,18 @@ public interface RequiresDistributionAndOrdering extends Write {
*/
Distribution requiredDistribution();
+ /**
+ * Returns if the distribution required by this write is strictly required or best effort only.
+ * <p>
+ * If true, Spark will strictly distribute incoming records across partitions to satisfy
+ * the required distribution before passing the records to the data source table on write.
+ * Otherwise, Spark may apply certain optimizations to speed up the query but break
+ * the distribution requirement.
+ *
+ * @return true if the distribution required by this write is strictly required; false otherwise.
+ */
+ default boolean distributionStrictlyRequired() { return true; }
+
/**
* Returns the number of partitions required by this write.
* <p>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
index 29a4142eb12..b4c6e19d0bc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.catalyst.plans.logical
-import org.apache.spark.sql.catalyst.AliasIdentifier
+import org.apache.spark.sql.catalyst.{AliasIdentifier, SQLConfHelper}
import org.apache.spark.sql.catalyst.analysis.{AnsiTypeCoercion, MultiInstanceRelation, Resolver, TypeCoercion, TypeCoercionBase}
import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable}
import org.apache.spark.sql.catalyst.catalog.CatalogTable.VIEW_STORING_ANALYZED_PLAN
@@ -1486,6 +1486,35 @@ case class Repartition(numPartitions: Int, shuffle: Boolean, child: LogicalPlan)
copy(child = newChild)
}
+trait HasPartitionExpressions extends SQLConfHelper {
+
+ val numPartitions = optNumPartitions.getOrElse(conf.numShufflePartitions)
+ require(numPartitions > 0, s"Number of partitions ($numPartitions) must be positive.")
+
+ def partitionExpressions: Seq[Expression]
+
+ def optNumPartitions: Option[Int]
+
+ protected def partitioning: Partitioning = if (partitionExpressions.isEmpty) {
+ RoundRobinPartitioning(numPartitions)
+ } else {
+ val (sortOrder, nonSortOrder) = partitionExpressions.partition(_.isInstanceOf[SortOrder])
+ require(sortOrder.isEmpty || nonSortOrder.isEmpty,
+ s"${getClass.getSimpleName} expects that either all its `partitionExpressions` are of type " +
+ "`SortOrder`, which means `RangePartitioning`, or none of them are `SortOrder`, which " +
+ "means `HashPartitioning`. In this case we have:" +
+ s"""
+ |SortOrder: $sortOrder
+ |NonSortOrder: $nonSortOrder
+ """.stripMargin)
+ if (sortOrder.nonEmpty) {
+ RangePartitioning(sortOrder.map(_.asInstanceOf[SortOrder]), numPartitions)
+ } else {
+ HashPartitioning(partitionExpressions, numPartitions)
+ }
+ }
+}
+
/**
* This method repartitions data using [[Expression]]s into `optNumPartitions`, and receives
* information about the number of partitions during execution. Used when a specific ordering or
@@ -1496,31 +1525,13 @@ case class Repartition(numPartitions: Int, shuffle: Boolean, child: LogicalPlan)
case class RepartitionByExpression(
partitionExpressions: Seq[Expression],
child: LogicalPlan,
- optNumPartitions: Option[Int]) extends RepartitionOperation {
-
- val numPartitions = optNumPartitions.getOrElse(conf.numShufflePartitions)
- require(numPartitions > 0, s"Number of partitions ($numPartitions) must be positive.")
+ optNumPartitions: Option[Int]) extends RepartitionOperation with HasPartitionExpressions {
override val partitioning: Partitioning = {
- val (sortOrder, nonSortOrder) = partitionExpressions.partition(_.isInstanceOf[SortOrder])
-
- require(sortOrder.isEmpty || nonSortOrder.isEmpty,
- s"${getClass.getSimpleName} expects that either all its `partitionExpressions` are of type " +
- "`SortOrder`, which means `RangePartitioning`, or none of them are `SortOrder`, which " +
- "means `HashPartitioning`. In this case we have:" +
- s"""
- |SortOrder: $sortOrder
- |NonSortOrder: $nonSortOrder
- """.stripMargin)
-
if (numPartitions == 1) {
SinglePartition
- } else if (sortOrder.nonEmpty) {
- RangePartitioning(sortOrder.map(_.asInstanceOf[SortOrder]), numPartitions)
- } else if (nonSortOrder.nonEmpty) {
- HashPartitioning(nonSortOrder, numPartitions)
} else {
- RoundRobinPartitioning(numPartitions)
+ super.partitioning
}
}
@@ -1551,19 +1562,12 @@ object RepartitionByExpression {
case class RebalancePartitions(
partitionExpressions: Seq[Expression],
child: LogicalPlan,
- initialNumPartitionOpt: Option[Int] = None) extends UnaryNode {
+ optNumPartitions: Option[Int] = None) extends UnaryNode with HasPartitionExpressions {
override def maxRows: Option[Long] = child.maxRows
override def output: Seq[Attribute] = child.output
override val nodePatterns: Seq[TreePattern] = Seq(REBALANCE_PARTITIONS)
- def partitioning: Partitioning = {
- val initialNumPartitions = initialNumPartitionOpt.getOrElse(conf.numShufflePartitions)
- if (partitionExpressions.isEmpty) {
- RoundRobinPartitioning(initialNumPartitions)
- } else {
- HashPartitioning(partitionExpressions, initialNumPartitions)
- }
- }
+ override val partitioning: Partitioning = super.partitioning
override protected def withNewChildInternal(newChild: LogicalPlan): RebalancePartitions =
copy(child = newChild)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTable.scala
index 61b0dfa3b9b..ec2b5f7ec7a 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTable.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTable.scala
@@ -52,7 +52,8 @@ class InMemoryTable(
override val properties: util.Map[String, String],
val distribution: Distribution = Distributions.unspecified(),
val ordering: Array[SortOrder] = Array.empty,
- val numPartitions: Option[Int] = None)
+ val numPartitions: Option[Int] = None,
+ val isDistributionStrictlyRequired: Boolean = true)
extends Table with SupportsRead with SupportsWrite with SupportsDelete
with SupportsMetadataColumns {
@@ -363,6 +364,8 @@ class InMemoryTable(
override def build(): Write = new Write with RequiresDistributionAndOrdering {
override def requiredDistribution: Distribution = distribution
+ override def distributionStrictlyRequired = isDistributionStrictlyRequired
+
override def requiredOrdering: Array[SortOrder] = ordering
override def requiredNumPartitions(): Int = {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
index 428aec70367..80a1a61d78a 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
@@ -100,7 +100,8 @@ class BasicInMemoryTableCatalog extends TableCatalog {
properties: util.Map[String, String],
distribution: Distribution,
ordering: Array[SortOrder],
- requiredNumPartitions: Option[Int]): Table = {
+ requiredNumPartitions: Option[Int],
+ distributionStrictlyRequired: Boolean = true): Table = {
if (tables.containsKey(ident)) {
throw new TableAlreadyExistsException(ident)
}
@@ -109,7 +110,7 @@ class BasicInMemoryTableCatalog extends TableCatalog {
val tableName = s"$name.${ident.quoted}"
val table = new InMemoryTable(tableName, schema, partitions, properties, distribution,
- ordering, requiredNumPartitions)
+ ordering, requiredNumPartitions, distributionStrictlyRequired)
tables.put(ident, table)
namespaces.putIfAbsent(ident.namespace.toList, Map())
table
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala
index 275255c9a3d..dbe02198567 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala
@@ -19,16 +19,15 @@ package org.apache.spark.sql.execution.datasources.v2
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.expressions.V2ExpressionUtils._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, RepartitionByExpression, Sort}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, RebalancePartitions, RepartitionByExpression, Sort}
import org.apache.spark.sql.connector.distributions._
import org.apache.spark.sql.connector.write.{RequiresDistributionAndOrdering, Write}
import org.apache.spark.sql.errors.QueryCompilationErrors
-import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.util.collection.Utils.sequenceToOption
object DistributionAndOrderingUtils {
- def prepareQuery(write: Write, query: LogicalPlan, conf: SQLConf): LogicalPlan = write match {
+ def prepareQuery(write: Write, query: LogicalPlan): LogicalPlan = write match {
case write: RequiresDistributionAndOrdering =>
val numPartitions = write.requiredNumPartitions()
@@ -41,15 +40,16 @@ object DistributionAndOrderingUtils {
}
val queryWithDistribution = if (distribution.nonEmpty) {
- val finalNumPartitions = if (numPartitions > 0) {
- numPartitions
- } else {
- conf.numShufflePartitions
- }
+ val optNumPartitions = if (numPartitions > 0) Some(numPartitions) else None
// the conversion to catalyst expressions above produces SortOrder expressions
// for OrderedDistribution and generic expressions for ClusteredDistribution
- // this allows RepartitionByExpression to pick either range or hash partitioning
- RepartitionByExpression(distribution, query, finalNumPartitions)
+ // this allows RebalancePartitions/RepartitionByExpression to pick either
+ // range or hash partitioning
+ if (write.distributionStrictlyRequired()) {
+ RepartitionByExpression(distribution, query, optNumPartitions)
+ } else {
+ RebalancePartitions(distribution, query, optNumPartitions)
+ }
} else if (numPartitions > 0) {
throw QueryCompilationErrors.numberOfPartitionsNotAllowedWithUnspecifiedDistributionError()
} else {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala
index 2fd1d52fd98..4422743c5ac 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala
@@ -44,7 +44,7 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
case a @ AppendData(r: DataSourceV2Relation, query, options, _, None) =>
val writeBuilder = newWriteBuilder(r.table, options, query.schema)
val write = writeBuilder.build()
- val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, conf)
+ val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query)
a.copy(write = Some(write), query = newQuery)
case o @ OverwriteByExpression(r: DataSourceV2Relation, deleteExpr, query, options, _, None) =>
@@ -68,7 +68,7 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
throw QueryExecutionErrors.overwriteTableByUnsupportedExpressionError(table)
}
- val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, conf)
+ val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query)
o.copy(write = Some(write), query = newQuery)
case o @ OverwritePartitionsDynamic(r: DataSourceV2Relation, query, options, _, None) =>
@@ -80,7 +80,7 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
case _ =>
throw QueryExecutionErrors.dynamicPartitionOverwriteUnsupportedByTableError(table)
}
- val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, conf)
+ val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query)
o.copy(write = Some(write), query = newQuery)
case WriteToMicroBatchDataSource(
@@ -90,14 +90,14 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
val write = buildWriteForMicroBatch(table, writeBuilder, outputMode)
val microBatchWrite = new MicroBatchWrite(batchId, write.toStreaming)
val customMetrics = write.supportedCustomMetrics.toSeq
- val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, conf)
+ val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query)
WriteToDataSourceV2(relation, microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, None) =>
val rowSchema = StructType.fromAttributes(rd.dataInput)
val writeBuilder = newWriteBuilder(r.table, Map.empty, rowSchema)
val write = writeBuilder.build()
- val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, conf)
+ val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query)
// project away any metadata columns that could be used for distribution and ordering
rd.copy(write = Some(write), query = Project(rd.dataInput, newQuery))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
index 36efe5ec1d2..335d7b410f6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
@@ -27,12 +27,15 @@ import org.apache.spark.sql.connector.distributions.{Distribution, Distributions
import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, NullOrdering, SortDirection, SortOrder}
import org.apache.spark.sql.connector.expressions.LogicalExpressions._
import org.apache.spark.sql.execution.{QueryExecution, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec
import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.streaming.MemoryStream
import org.apache.spark.sql.execution.streaming.sources.ContinuousMemoryStream
import org.apache.spark.sql.functions.lit
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.{StreamingQueryException, Trigger}
+import org.apache.spark.sql.test.SQLTestData.TestData
import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
import org.apache.spark.sql.util.QueryExecutionListener
@@ -53,15 +56,15 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
.add("data", StringType)
test("ordered distribution and sort with same exprs: append") {
- checkOrderedDistributionAndSortWithSameExprs("append")
+ checkOrderedDistributionAndSortWithSameExprsInVariousCases("append")
}
test("ordered distribution and sort with same exprs: overwrite") {
- checkOrderedDistributionAndSortWithSameExprs("overwrite")
+ checkOrderedDistributionAndSortWithSameExprsInVariousCases("overwrite")
}
test("ordered distribution and sort with same exprs: overwriteDynamic") {
- checkOrderedDistributionAndSortWithSameExprs("overwriteDynamic")
+ checkOrderedDistributionAndSortWithSameExprsInVariousCases("overwriteDynamic")
}
test("ordered distribution and sort with same exprs: micro-batch append") {
@@ -100,13 +103,23 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkOrderedDistributionAndSortWithSameExprs(microBatchPrefix + "complete", Some(10))
}
- private def checkOrderedDistributionAndSortWithSameExprs(command: String): Unit = {
- checkOrderedDistributionAndSortWithSameExprs(command, None)
+ private def checkOrderedDistributionAndSortWithSameExprsInVariousCases(cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkOrderedDistributionAndSortWithSameExprs(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkOrderedDistributionAndSortWithSameExprs(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
)
@@ -120,7 +133,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
Seq.empty
)
)
- val writePartitioning = orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ } else {
+ orderedWritePartitioning(writeOrdering, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -128,19 +145,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
targetNumPartitions,
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("clustered distribution and sort with same exprs: append") {
- checkClusteredDistributionAndSortWithSameExprs("append")
+ checkClusteredDistributionAndSortWithSameExprsInVariousCases("append")
}
test("clustered distribution and sort with same exprs: overwrite") {
- checkClusteredDistributionAndSortWithSameExprs("overwrite")
+ checkClusteredDistributionAndSortWithSameExprsInVariousCases("overwrite")
}
test("clustered distribution and sort with same exprs: overwriteDynamic") {
- checkClusteredDistributionAndSortWithSameExprs("overwriteDynamic")
+ checkClusteredDistributionAndSortWithSameExprsInVariousCases("overwriteDynamic")
}
test("clustered distribution and sort with same exprs: micro-batch append") {
@@ -179,13 +199,23 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkClusteredDistributionAndSortWithSameExprs(microBatchPrefix + "complete", Some(10))
}
- private def checkClusteredDistributionAndSortWithSameExprs(command: String): Unit = {
- checkClusteredDistributionAndSortWithSameExprs(command, None)
+ private def checkClusteredDistributionAndSortWithSameExprsInVariousCases(cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkClusteredDistributionAndSortWithSameExprs(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkClusteredDistributionAndSortWithSameExprs(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.DESCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -208,7 +238,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
)
)
val writePartitioningExprs = Seq(attr("data"), attr("id"))
- val writePartitioning = clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ } else {
+ clusteredWritePartitioning(writePartitioningExprs, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -216,19 +250,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
targetNumPartitions,
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("clustered distribution and sort with extended exprs: append") {
- checkClusteredDistributionAndSortWithExtendedExprs("append")
+ checkClusteredDistributionAndSortWithExtendedExprsInVariousCases("append")
}
test("clustered distribution and sort with extended exprs: overwrite") {
- checkClusteredDistributionAndSortWithExtendedExprs("overwrite")
+ checkClusteredDistributionAndSortWithExtendedExprsInVariousCases("overwrite")
}
test("clustered distribution and sort with extended exprs: overwriteDynamic") {
- checkClusteredDistributionAndSortWithExtendedExprs("overwriteDynamic")
+ checkClusteredDistributionAndSortWithExtendedExprsInVariousCases("overwriteDynamic")
}
test("clustered distribution and sort with extended exprs: micro-batch append") {
@@ -271,13 +308,23 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkClusteredDistributionAndSortWithExtendedExprs(microBatchPrefix + "complete", Some(10))
}
- private def checkClusteredDistributionAndSortWithExtendedExprs(command: String): Unit = {
- checkClusteredDistributionAndSortWithExtendedExprs(command, None)
+ private def checkClusteredDistributionAndSortWithExtendedExprsInVariousCases(cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkClusteredDistributionAndSortWithExtendedExprs(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkClusteredDistributionAndSortWithExtendedExprs(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.DESCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -300,7 +347,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
)
)
val writePartitioningExprs = Seq(attr("data"))
- val writePartitioning = clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ } else {
+ clusteredWritePartitioning(writePartitioningExprs, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -308,7 +359,10 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
targetNumPartitions,
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("unspecified distribution and local sort: append") {
@@ -466,15 +520,15 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
}
test("ordered distribution and sort with manual global sort: append") {
- checkOrderedDistributionAndSortWithManualGlobalSort("append")
+ checkOrderedDistributionAndSortWithManualGlobalSortInVariousCases("append")
}
test("ordered distribution and sort with manual global sort: overwrite") {
- checkOrderedDistributionAndSortWithManualGlobalSort("overwrite")
+ checkOrderedDistributionAndSortWithManualGlobalSortInVariousCases("overwrite")
}
test("ordered distribution and sort with manual global sort: overwriteDynamic") {
- checkOrderedDistributionAndSortWithManualGlobalSort("overwriteDynamic")
+ checkOrderedDistributionAndSortWithManualGlobalSortInVariousCases("overwriteDynamic")
}
test("ordered distribution and sort with manual global sort with numPartitions: append") {
@@ -490,13 +544,23 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkOrderedDistributionAndSortWithManualGlobalSort("overwriteDynamic", Some(10))
}
- private def checkOrderedDistributionAndSortWithManualGlobalSort(command: String): Unit = {
- checkOrderedDistributionAndSortWithManualGlobalSort(command, None)
+ private def checkOrderedDistributionAndSortWithManualGlobalSortInVariousCases(cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkOrderedDistributionAndSortWithManualGlobalSort(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkOrderedDistributionAndSortWithManualGlobalSort(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -517,7 +581,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
Seq.empty
)
)
- val writePartitioning = orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ } else {
+ orderedWritePartitioning(writeOrdering, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -526,19 +594,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
writeTransform = df => df.orderBy("data", "id"),
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("ordered distribution and sort with incompatible global sort: append") {
- checkOrderedDistributionAndSortWithIncompatibleGlobalSort("append")
+ checkOrderedDistributionAndSortWithIncompatibleGlobalSortInVariousCases("append")
}
test("ordered distribution and sort with incompatible global sort: overwrite") {
- checkOrderedDistributionAndSortWithIncompatibleGlobalSort("overwrite")
+ checkOrderedDistributionAndSortWithIncompatibleGlobalSortInVariousCases("overwrite")
}
test("ordered distribution and sort with incompatible global sort: overwriteDynamic") {
- checkOrderedDistributionAndSortWithIncompatibleGlobalSort("overwriteDynamic")
+ checkOrderedDistributionAndSortWithIncompatibleGlobalSortInVariousCases("overwriteDynamic")
}
test("ordered distribution and sort with incompatible global sort with numPartitions: append") {
@@ -555,13 +626,24 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkOrderedDistributionAndSortWithIncompatibleGlobalSort("overwriteDynamic", Some(10))
}
- private def checkOrderedDistributionAndSortWithIncompatibleGlobalSort(command: String): Unit = {
- checkOrderedDistributionAndSortWithIncompatibleGlobalSort(command, None)
+ private def checkOrderedDistributionAndSortWithIncompatibleGlobalSortInVariousCases(
+ cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkOrderedDistributionAndSortWithIncompatibleGlobalSort(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkOrderedDistributionAndSortWithIncompatibleGlobalSort(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -582,7 +664,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
Seq.empty
)
)
- val writePartitioning = orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ } else {
+ orderedWritePartitioning(writeOrdering, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -591,19 +677,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
writeTransform = df => df.orderBy(df("data").desc, df("id").asc),
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("ordered distribution and sort with manual local sort: append") {
- checkOrderedDistributionAndSortWithManualLocalSort("append")
+ checkOrderedDistributionAndSortWithManualLocalSortInVariousCases("append")
}
test("ordered distribution and sort with manual local sort: overwrite") {
- checkOrderedDistributionAndSortWithManualLocalSort("overwrite")
+ checkOrderedDistributionAndSortWithManualLocalSortInVariousCases("overwrite")
}
test("ordered distribution and sort with manual local sort: overwriteDynamic") {
- checkOrderedDistributionAndSortWithManualLocalSort("overwriteDynamic")
+ checkOrderedDistributionAndSortWithManualLocalSortInVariousCases("overwriteDynamic")
}
test("ordered distribution and sort with manual local sort with numPartitions: append") {
@@ -619,13 +708,23 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkOrderedDistributionAndSortWithManualLocalSort("overwriteDynamic", Some(10))
}
- private def checkOrderedDistributionAndSortWithManualLocalSort(command: String): Unit = {
- checkOrderedDistributionAndSortWithManualLocalSort(command, None)
+ private def checkOrderedDistributionAndSortWithManualLocalSortInVariousCases(cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkOrderedDistributionAndSortWithManualLocalSort(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkOrderedDistributionAndSortWithManualLocalSort(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -646,7 +745,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
Seq.empty
)
)
- val writePartitioning = orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ orderedWritePartitioning(writeOrdering, targetNumPartitions)
+ } else {
+ orderedWritePartitioning(writeOrdering, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -655,19 +758,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
writeTransform = df => df.sortWithinPartitions("data", "id"),
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("clustered distribution and local sort with manual global sort: append") {
- checkClusteredDistributionAndLocalSortWithManualGlobalSort("append")
+ checkClusteredDistributionAndLocalSortWithManualGlobalSortInVariousCases("append")
}
test("clustered distribution and local sort with manual global sort: overwrite") {
- checkClusteredDistributionAndLocalSortWithManualGlobalSort("overwrite")
+ checkClusteredDistributionAndLocalSortWithManualGlobalSortInVariousCases("overwrite")
}
test("clustered distribution and local sort with manual global sort: overwriteDynamic") {
- checkClusteredDistributionAndLocalSortWithManualGlobalSort("overwriteDynamic")
+ checkClusteredDistributionAndLocalSortWithManualGlobalSortInVariousCases("overwriteDynamic")
}
test("clustered distribution and local sort with manual global sort with numPartitions: append") {
@@ -684,13 +790,24 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkClusteredDistributionAndLocalSortWithManualGlobalSort("overwriteDynamic", Some(10))
}
- private def checkClusteredDistributionAndLocalSortWithManualGlobalSort(command: String): Unit = {
- checkClusteredDistributionAndLocalSortWithManualGlobalSort(command, None)
+ private def checkClusteredDistributionAndLocalSortWithManualGlobalSortInVariousCases(
+ cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkClusteredDistributionAndLocalSortWithManualGlobalSort(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkClusteredDistributionAndLocalSortWithManualGlobalSort(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.DESCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -712,7 +829,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
)
)
val writePartitioningExprs = Seq(attr("data"))
- val writePartitioning = clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ } else {
+ clusteredWritePartitioning(writePartitioningExprs, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -721,19 +842,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
writeTransform = df => df.orderBy("data", "id"),
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("clustered distribution and local sort with manual local sort: append") {
- checkClusteredDistributionAndLocalSortWithManualLocalSort("append")
+ checkClusteredDistributionAndLocalSortWithManualLocalSortInVariousCases("append")
}
test("clustered distribution and local sort with manual local sort: overwrite") {
- checkClusteredDistributionAndLocalSortWithManualLocalSort("overwrite")
+ checkClusteredDistributionAndLocalSortWithManualLocalSortInVariousCases("overwrite")
}
test("clustered distribution and local sort with manual local sort: overwriteDynamic") {
- checkClusteredDistributionAndLocalSortWithManualLocalSort("overwriteDynamic")
+ checkClusteredDistributionAndLocalSortWithManualLocalSortInVariousCases("overwriteDynamic")
}
test("clustered distribution and local sort with manual local sort with numPartitions: append") {
@@ -750,13 +874,24 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
checkClusteredDistributionAndLocalSortWithManualLocalSort("overwriteDynamic", Some(10))
}
- private def checkClusteredDistributionAndLocalSortWithManualLocalSort(command: String): Unit = {
- checkClusteredDistributionAndLocalSortWithManualLocalSort(command, None)
+ private def checkClusteredDistributionAndLocalSortWithManualLocalSortInVariousCases(
+ cmd: String) = {
+ Seq(true, false).foreach { distributionStrictlyRequired =>
+ Seq(true, false).foreach { dataSkewed =>
+ Seq(true, false).foreach { coalesce =>
+ checkClusteredDistributionAndLocalSortWithManualLocalSort(
+ cmd, None, distributionStrictlyRequired, dataSkewed, coalesce)
+ }
+ }
+ }
}
private def checkClusteredDistributionAndLocalSortWithManualLocalSort(
command: String,
- targetNumPartitions: Option[Int]): Unit = {
+ targetNumPartitions: Option[Int] = None,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
val tableOrdering = Array[SortOrder](
sort(FieldReference("data"), SortDirection.DESCENDING, NullOrdering.NULLS_FIRST),
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)
@@ -778,7 +913,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
)
)
val writePartitioningExprs = Seq(attr("data"))
- val writePartitioning = clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ val writePartitioning = if (!coalesce) {
+ clusteredWritePartitioning(writePartitioningExprs, targetNumPartitions)
+ } else {
+ clusteredWritePartitioning(writePartitioningExprs, Some(1))
+ }
checkWriteRequirements(
tableDistribution,
@@ -787,7 +926,10 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWritePartitioning = writePartitioning,
expectedWriteOrdering = writeOrdering,
writeTransform = df => df.orderBy("data", "id"),
- writeCommand = command)
+ writeCommand = command,
+ distributionStrictlyRequired = distributionStrictlyRequired,
+ dataSkewed = dataSkewed,
+ coalesce = coalesce)
}
test("continuous mode does not support write distribution and ordering") {
@@ -845,6 +987,7 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
}
}
+ // scalastyle:off argcount
private def checkWriteRequirements(
tableDistribution: Distribution,
tableOrdering: Array[SortOrder],
@@ -853,8 +996,11 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWriteOrdering: Seq[catalyst.expressions.SortOrder],
writeTransform: DataFrame => DataFrame = df => df,
writeCommand: String,
- expectAnalysisException: Boolean = false): Unit = {
-
+ expectAnalysisException: Boolean = false,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
+ // scalastyle:on argcount
if (writeCommand.startsWith(microBatchPrefix)) {
checkMicroBatchWriteRequirements(
tableDistribution,
@@ -874,10 +1020,14 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWriteOrdering,
writeTransform,
writeCommand,
- expectAnalysisException)
+ expectAnalysisException,
+ distributionStrictlyRequired,
+ dataSkewed,
+ coalesce)
}
}
+ // scalastyle:off argcount
private def checkBatchWriteRequirements(
tableDistribution: Distribution,
tableOrdering: Array[SortOrder],
@@ -886,12 +1036,22 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
expectedWriteOrdering: Seq[catalyst.expressions.SortOrder],
writeTransform: DataFrame => DataFrame = df => df,
writeCommand: String = "append",
- expectAnalysisException: Boolean = false): Unit = {
+ expectAnalysisException: Boolean = false,
+ distributionStrictlyRequired: Boolean = true,
+ dataSkewed: Boolean = false,
+ coalesce: Boolean = false): Unit = {
+ // scalastyle:on argcount
catalog.createTable(ident, schema, Array.empty, emptyProps, tableDistribution,
- tableOrdering, tableNumPartitions)
+ tableOrdering, tableNumPartitions, distributionStrictlyRequired)
- val df = spark.createDataFrame(Seq((1, "a"), (2, "b"), (3, "c"))).toDF("id", "data")
+ val df = if (!dataSkewed) {
+ spark.createDataFrame(Seq((1, "a"), (2, "b"), (3, "c"))).toDF("id", "data")
+ } else {
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(if (i > 4) 5 else i, i.toString)), 3)
+ .toDF("id", "data")
+ }
val writer = writeTransform(df).writeTo(tableNameAsString)
def executeCommand(): SparkPlan = writeCommand match {
@@ -905,12 +1065,51 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
executeCommand()
}
} else {
- val executedPlan = executeCommand()
-
- checkPartitioningAndOrdering(executedPlan, expectedWritePartitioning, expectedWriteOrdering)
+ if (coalesce) {
+ val executedPlan = executeCommand()
+ val read = collect(executedPlan) {
+ case r: AQEShuffleReadExec => r
+ }
+ assert(read.size == 1)
+ assert(read.head.partitionSpecs.size == 1)
+ checkPartitioningAndOrdering(
+ // num of partition in expectedWritePartitioning is 1
+ executedPlan, expectedWritePartitioning, expectedWriteOrdering, 1)
+ } else {
+ if (dataSkewed && !distributionStrictlyRequired) {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "5",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") {
+ val executedPlan = executeCommand()
+ val read = collect(executedPlan) {
+ case r: AQEShuffleReadExec => r
+ }
+ assert(read.size == 1)
+ // skew data: 144, 88, 88, 144, 80
+ // after repartition: 72, 72, 88, 88, 72, 72, 80
+ assert(read.head.partitionSpecs.size >= 7)
+
+ checkPartitioningAndOrdering(
+ executedPlan, expectedWritePartitioning, expectedWriteOrdering, 1, true)
+ }
+ } else {
+ withSQLConf(
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") {
+ val executedPlan = executeCommand()
+ checkPartitioningAndOrdering(
+ executedPlan, expectedWritePartitioning, expectedWriteOrdering, 1)
+ }
+ }
+ }
checkAnswer(spark.table(tableNameAsString), df)
}
+ catalog.dropTable(ident)
}
private def checkMicroBatchWriteRequirements(
@@ -985,7 +1184,8 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
plan: SparkPlan,
partitioning: physical.Partitioning,
ordering: Seq[catalyst.expressions.SortOrder],
- maxNumShuffles: Int = 1): Unit = {
+ maxNumShuffles: Int = 1,
+ skewSplit: Boolean = false): Unit = {
val sorts = collect(plan) { case s: SortExec => s }
assert(sorts.size <= 1, "must be at most one sort")
@@ -1005,7 +1205,12 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
actualPartitioning
case other => other
}
- assert(actualPartitioning == expectedPartitioning, "partitioning must match")
+
+ if (skewSplit) {
+ assert(actualPartitioning.numPartitions > conf.numShufflePartitions)
+ } else {
+ assert(actualPartitioning == expectedPartitioning, "partitioning must match")
+ }
val actualOrdering = plan.outputOrdering
val expectedOrdering = ordering.map(resolveAttrs(_, plan))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org