dongjoon-hyun commented on code in PR #58591:
URL: https://github.com/apache/spark/pull/58591#discussion_r3955624462
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -274,6 +307,34 @@ case class GroupPartitionsExec(
@transient private lazy val hasCoalescing: Boolean =
groupedPartitions.exists(_._2.size > 1)
+ // All values are computed on the driver by `grouping`, so they are reported
through
+ // `sendDriverMetrics` rather than task-side accumulators. Registration
reads constructor
+ // parameters only: pruning and replication are facts of the alignment path
and are registered
+ // only in the modes that produce them, while coalescing can happen in any
mode and reports 0
+ // when nothing merged.
+ @transient override lazy val metrics: Map[String, SQLMetric] = Map(
+ "numInputPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
input partitions"),
+ "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
partitions"),
+ "numEmptyPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
empty partitions"),
+ "numCoalescedPartitions" ->
+ SQLMetrics.createMetric(sparkContext, "number of coalesced partitions"),
+ "maxPartitionsPerGroup" ->
+ SQLMetrics.createMetric(sparkContext, "max partitions per group")) ++ {
+ if (expectedPartitionKeys.isDefined && !distributePartitions) {
+ Map("numReplicatedPartitions" -> SQLMetrics.createMetric(sparkContext,
+ "number of replicated input partition reads"))
+ } else {
+ Map.empty[String, SQLMetric]
+ }
+ } ++ {
+ if (expectedPartitionKeys.isDefined) {
Review Comment:
Same theme as the block above, but this one I don't think can be narrowed:
on the `OrderedDistribution` path the expected keys are the child's own keys
(`EnsureRequirements.scala:110-115`), so pruning is 0 by construction, yet
nothing in the constructor distinguishes that producer from a join producer.
Probably enough to soften the comment at the top of this block rather than
change the condition.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -274,6 +307,34 @@ case class GroupPartitionsExec(
@transient private lazy val hasCoalescing: Boolean =
groupedPartitions.exists(_._2.size > 1)
+ // All values are computed on the driver by `grouping`, so they are reported
through
+ // `sendDriverMetrics` rather than task-side accumulators. Registration
reads constructor
+ // parameters only: pruning and replication are facts of the alignment path
and are registered
+ // only in the modes that produce them, while coalescing can happen in any
mode and reports 0
+ // when nothing merged.
+ @transient override lazy val metrics: Map[String, SQLMetric] = Map(
+ "numInputPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
input partitions"),
+ "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
partitions"),
+ "numEmptyPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
empty partitions"),
+ "numCoalescedPartitions" ->
+ SQLMetrics.createMetric(sparkContext, "number of coalesced partitions"),
+ "maxPartitionsPerGroup" ->
+ SQLMetrics.createMetric(sparkContext, "max partitions per group")) ++ {
+ if (expectedPartitionKeys.isDefined && !distributePartitions) {
+ Map("numReplicatedPartitions" -> SQLMetrics.createMetric(sparkContext,
+ "number of replicated input partition reads"))
Review Comment:
This condition is broader than the "align-replicate path" row in the PR
description.
In a plain SPJ (partial clustering off), `distributePartitions =
applyPartialClustering && !replicateXSide` is `false`, and every split count in
`mergedPartitionKeys` is 1 because it comes from
`mergeAndDedupPartitions(...).map((_, 1))` (`EnsureRequirements.scala:617`,
`:750-752`). So this metric is registered — always reporting 0 — on the most
common SPJ path, which also contradicts the comment just above ("registered
only in the modes that produce them").
Whether replication is possible is decidable from the constructor parameters
alone:
```suggestion
if (expectedPartitionKeys.exists(_.exists(_._2 > 1)) &&
!distributePartitions) {
Map("numReplicatedPartitions" -> SQLMetrics.createMetric(sparkContext,
"number of replicated input partition reads"))
```
Minor, while here: the key says `Partitions` but the description says
"reads". `numReplicatedPartitionReads` would keep the two in sync.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -359,7 +420,47 @@ case class GroupPartitionsExec(
private[v2] def kWayMergeOrdering: Seq[SortOrder] =
child.outputOrdering.map(_.copy(sameOrderExpressions = Seq.empty))
+ private def sendDriverMetrics(): Unit = {
+ val executionId =
sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+ val driverAccumUpdates = ArrayBuffer.empty[(Long, Long)]
+ def set(name: String, value: Long): Unit = {
+ val metric = metrics(name)
+ metric.set(value)
+ driverAccumUpdates += (metric.id -> value)
+ }
+ // A single pass for the three per-group counts; an empty group and a
coalesced one are
+ // mutually exclusive.
+ var numEmptyPartitions = 0
+ var numCoalescedPartitions = 0
+ var maxPartitionsPerGroup = 0
+ groupedPartitions.foreach { case (_, group) =>
+ val size = group.size
+ if (size == 0) {
+ numEmptyPartitions += 1
+ } else if (size > 1) {
+ numCoalescedPartitions += 1
+ }
+ if (size > maxPartitionsPerGroup) {
+ maxPartitionsPerGroup = size
+ }
+ }
+ set("numInputPartitions", grouping.numInputPartitions)
+ set("numPartitions", groupedPartitions.size)
+ set("numEmptyPartitions", numEmptyPartitions)
+ set("numCoalescedPartitions", numCoalescedPartitions)
+ set("maxPartitionsPerGroup", maxPartitionsPerGroup)
+ if (expectedPartitionKeys.isDefined && !distributePartitions) {
+ set("numReplicatedPartitions", grouping.numReplicatedPartitions)
+ }
+ if (expectedPartitionKeys.isDefined) {
+ set("numPrunedPartitions", grouping.numPrunedPartitions)
+ }
+ SQLMetrics.postDriverMetricsUpdatedByValue(
+ sparkContext, executionId, driverAccumUpdates.toSeq)
+ }
+
override protected def doExecute(): RDD[InternalRow] = {
+ sendDriverMetrics()
Review Comment:
Could this be guarded by a lazy val? `AQEShuffleReadExec`, which this
follows, calls it from `private lazy val shuffleRDD`
(`AQEShuffleReadExec.scala:268-276`), and `FileSourceScanExec` calls it from
its `inputRDD` lazy val. That guard does real work:
- `SQLAppStatusListener.onDriverAccumUpdates` does `exec.driverAccumUpdates
++ accumUpdates`, and that field is a `Seq[(Long, Long)]`
(`SQLAppStatusListener.scala:509`, `:433`) — a repeated post of the same
accumulator id **appends** rather than overwrites.
- `SQLAppStatusListener.scala:278-296` then appends each value into the
array, and `SUM_METRIC` is rendered as `values.sum` in
`MetricUtils.stringValue`.
So a second `doExecute()` within one execution would silently double every
value here. I could not find a path today where this node is executed twice
(exchange and subquery reuse both cache at the `Exchange` / `SubqueryExec`
level), so this is not a live bug as far as I can tell — but the failure mode
is silent and the convention is already established:
```scala
@transient private lazy val metricsSent: Unit = sendDriverMetrics()
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -274,6 +307,34 @@ case class GroupPartitionsExec(
@transient private lazy val hasCoalescing: Boolean =
groupedPartitions.exists(_._2.size > 1)
+ // All values are computed on the driver by `grouping`, so they are reported
through
+ // `sendDriverMetrics` rather than task-side accumulators. Registration
reads constructor
+ // parameters only: pruning and replication are facts of the alignment path
and are registered
+ // only in the modes that produce them, while coalescing can happen in any
mode and reports 0
+ // when nothing merged.
+ @transient override lazy val metrics: Map[String, SQLMetric] = Map(
+ "numInputPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
input partitions"),
+ "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
partitions"),
+ "numEmptyPartitions" -> SQLMetrics.createMetric(sparkContext, "number of
empty partitions"),
+ "numCoalescedPartitions" ->
+ SQLMetrics.createMetric(sparkContext, "number of coalesced partitions"),
+ "maxPartitionsPerGroup" ->
+ SQLMetrics.createMetric(sparkContext, "max partitions per group")) ++ {
Review Comment:
`maxPartitionsPerGroup` is a `SUM_METRIC`. That renders correctly for a
single driver value, but it means a max would be summed if this node ever
reported twice within one execution — see my comment on the
`sendDriverMetrics()` call in `doExecute`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -450,12 +552,20 @@ case class GroupPartitionsExec(
}
}
-/** What a [[GroupPartitionsExec]] computes once and reports from several
members. */
+/**
+ * What a [[GroupPartitionsExec]] computes once and reports from several
members.
+ * `numInputPartitions` is the child's split count; the last two fields count
the alignment's
+ * effect on the reads of those splits (see `alignToExpectedKeys`), and are 0
outside the
+ * alignment path.
+ */
private case class PartitionGrouping(
partitions: Seq[(InternalRowComparableWrapper, Seq[Int])],
isGrouped: Boolean,
isCollapsed: Boolean,
- keysRewritten: Boolean)
+ keysRewritten: Boolean,
+ numInputPartitions: Int,
Review Comment:
This field looks redundant. `childKp.numPartitions` always equals
`child.outputPartitioning.numPartitions` — `PartitioningCollection` enforces a
uniform partition count across its members, and `identityGrouping` in this same
file (`:350`) already reads it that way.
`alignToExpectedKeys` genuinely needs the argument, but `sendDriverMetrics`
could read it straight off the child instead of carrying it through the case
class.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -380,6 +481,7 @@ case class GroupPartitionsExec(
child.supportsColumnar && !(hasCoalescing && enableSortedMerge &&
canUseSortedMerge)
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ sendDriverMetrics()
Review Comment:
Same as the call in `doExecute` above — a lazy val guard would cover both
entry points at once.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -508,6 +523,281 @@ class GroupPartitionsExecSuite extends SharedSparkSession
{
assert(gpe.tryEnableSortedMerge().isEmpty)
}
}
+
+ test("SPARK-59310: basic counts without alignment") {
+ // Keys [1, 2, 1]: 3 input splits, key 1 coalesces partitions 0 and 2, 2
output partitions.
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(1))))
+ val gpe = GroupPartitionsExec(child)
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 0)
+ assert(gpe.metrics("numCoalescedPartitions").value === 1)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 2)
+ // Without expectedPartitionKeys there is no alignment, so its metrics
stay unregistered.
+ assert(!gpe.metrics.contains("numPrunedPartitions"))
+ assert(!gpe.metrics.contains("numReplicatedPartitions"))
+ }
+
+ test("SPARK-59310: zero coalesced without duplicate keys") {
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(3))))
+ val gpe = GroupPartitionsExec(child)
+ gpe.execute()
+
+ assert(gpe.metrics("numCoalescedPartitions").value === 0)
+ assert(gpe.metrics("numPartitions").value === 3)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 1)
+ }
+
+ test("SPARK-59310: distribute alignment pads and never replicates") {
+ // Child splits: key 1 -> [0], key 2 -> [1, 2]. Expected: key 1 x1, key 2
x3, key 3 x1.
+ // Key 2 spreads its 2 splits over 3 expected partitions (one empty pad)
and key 3 has no
+ // split (one more empty), so 5 output partitions with 2 empty and nothing
coalesced.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(2))))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 3, keyOf(3)
-> 1)),
+ distributePartitions = true)
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 5)
+ assert(gpe.metrics("numEmptyPartitions").value === 2)
+ assert(gpe.metrics("numPrunedPartitions").value === 0)
+ assert(gpe.metrics("numCoalescedPartitions").value === 0, "distribute
never coalesces")
+ assert(!gpe.metrics.contains("numReplicatedPartitions"), "distribute never
replicates")
+ }
+
+ test("SPARK-59310: alignment prunes unmatched keys, pads missing ones") {
+ // The expected keys carry key 1 and a key-3 slot the child does not hold,
as an inner
+ // join's intersection combined with the other side's layout would. The 2
splits of key 2
+ // cannot produce join output and never enter the alignment; the missing
key 3 pads two
+ // empty output partitions, and being empty, replicates nothing despite
its 2 slots.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(2))))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(3) -> 2)))
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 3)
+ assert(gpe.metrics("numPrunedPartitions").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 2)
+ assert(gpe.metrics("numCoalescedPartitions").value === 0)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 1)
+ assert(gpe.metrics("numReplicatedPartitions").value === 0,
+ "empty groups replicate nothing, and the single-split key 1 has no copy")
+ }
+
+ test("SPARK-59310: replicate alignment counts the reads beyond the first") {
+ // The other join side expects 2 partitions for key 1, so this side's
splits for the key are
+ // replicated to both: the slot beyond the first re-reads both splits, 2
extra input
+ // partition reads. Each output partition also coalesces the 2 splits of
the key.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(1))))
+ val gpe = GroupPartitionsExec(child, expectedPartitionKeys =
Some(Seq(keyOf(1) -> 2)))
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 2)
+ assert(gpe.metrics("numPartitions").value === 2)
+ assert(gpe.metrics("numReplicatedPartitions").value === 2)
+ assert(gpe.metrics("numCoalescedPartitions").value === 2, "both copies
merge the 2 splits")
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 0)
+ assert(gpe.metrics("numPrunedPartitions").value === 0)
+ }
+
+ test("SPARK-59310: an inner join intersection prunes both sides") {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withTable(s"testcat.ns.$items", s"testcat.ns.$purchases") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 44.0, cast('2020-01-15' as timestamp)), " +
+ s"(1, 45.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 11.0, cast('2020-01-01' as timestamp)), " +
+ s"(4, 19.5, cast('2020-02-01' as timestamp))")
+
+ // The sides hold different keys ({1,2,3} vs {1,2,4}), so grouping
alone cannot align
+ // them and the join pushes the inner intersection {1,2} down as the
expected keys.
+ // Each side then coalesces its duplicate-key splits (items merges two
groups of two,
+ // purchases one group of three) and prunes the one split of its
unmatched key
+ // (items key 3, purchases key 4); nothing is empty or replicated.
+ // Both AQE arms: the query has no shuffle, so the executed nodes and
their accumulator
+ // ids are the same with and without AQE, and the reporting chain must
work in both.
Review Comment:
This comment isn't accurate: each `sql(...)` call inside the loop builds a
fresh physical plan, and `metrics` is a lazy val on each new node instance, so
the accumulator ids differ between the two arms.
The test itself is fine — it reads `planGraph` / `executionMetrics` per
execution id — so only the comment needs fixing. Covering both AQE arms is
still worth doing for the reporting chain, which is what the second half of the
sentence says.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -149,22 +150,52 @@ case class GroupPartitionsExec(
})
}
- /** Aligns partitions based on `expectedPartitionKeys` and clustering mode.
*/
- private def alignToExpectedKeys(keyMap: Map[InternalRowComparableWrapper,
Seq[Int]]) = {
+ /**
+ * Aligns partitions based on `expectedPartitionKeys` and clustering mode.
+ *
+ * Returns the aligned groups, whether their keys ended up unique, and two
counts, both in
+ * input partition reads relative to the baseline of every input being read
exactly once.
+ * The pruned count is the reads that never happen: inputs of a key the
alignment never
+ * references. The join proved such keys cannot produce output, e.g. an
inner join keeps
+ * only the intersection. Which keys are expected is orthogonal to the mode:
an
+ * `OrderedDistribution` producer expects the child's own keys and so never
prunes, while a
+ * join producer can prune in either mode. The replicated count is the reads
that happen
+ * again: in the replicate mode of partial clustering every expected
partition of a held key
+ * re-reads all its splits, counted beyond the first pass (3 splits over 2
slots count 3).
+ * Total input reads = `numInputPartitions` - pruned + replicated.
+ */
+ private def alignToExpectedKeys(
+ keyMap: Map[InternalRowComparableWrapper, Seq[Int]],
+ numInputPartitions: Int) = {
var isGrouped = true
+ var numReplicatedPartitions = 0
+ // Splits the alignment references, accumulated over the expected keys.
Every producer of
+ // `expectedPartitionKeys` deduplicates the keys, so no split is counted
twice here.
+ var numMatchedPartitions = 0
val alignedPartitions = expectedPartitionKeys.get.flatMap { case (key,
numSplits) =>
+ // Every producer derives the split counts from a `groupBy` size or the
literal 1; a
+ // non-positive count would emit a different partition count for the key
than the other
+ // side expects, breaking the pairing this alignment exists for.
+ assert(numSplits > 0, s"expected partition key split count must be
positive: $numSplits")
Review Comment:
This adds a new runtime failure mode that isn't really about metrics: today
a non-positive count is silently absorbed by `Seq.fill(0)` / `padTo(0)`, and
after this it throws.
I did verify the contract holds for both producers
(`mergeAndDedupPartitions(...).map((_, 1))` at `EnsureRequirements.scala:617`
and `keys.groupBy(identity).view.mapValues(_.size)` at `:110` and `:726`), so
the assertion is correct — it just seems like it belongs in its own change
rather than riding along with the metrics. The `s"..."` message is by-name, so
no overhead concern.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -18,17 +18,22 @@
package org.apache.spark.sql.execution.datasources.v2
import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute,
AttributeReference, SortOrder, TransformExpression}
import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution,
KeyedPartitioning, KeyedShuffleSpec, KeyReducer, Partitioning,
PartitioningCollection, UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
+import org.apache.spark.sql.connector.KeyGroupedPartitioningSuiteBase
import org.apache.spark.sql.connector.catalog.functions.{BucketFunction,
BucketReducer, Reducer}
+import org.apache.spark.sql.connector.expressions.Expressions.identity
import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode,
SafeForKWayMerge}
+import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
+import org.apache.spark.sql.execution.ui.SparkPlanGraphNode
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{DataType, IntegerType}
-class GroupPartitionsExecSuite extends SharedSparkSession {
+class GroupPartitionsExecSuite
+ extends KeyGroupedPartitioningSuiteBase with SQLMetricsTestUtils {
Review Comment:
Widening the base class is safe as far as the existing tests go —
`DistributionAndOrderingSuiteBase` does not override `sparkConf`, so the only
real change is the `testcat` registration. Two things still give me pause:
1. Three catalog-backed e2e tests in an otherwise pure unit-test suite
change its character. `KeyGroupedPartitioningSuite` already extends this same
base and looks like a more natural home for them.
2. `KeyGroupedPartitioningSuite` carries `@ExtendedSQLTest`; this suite does
not, so the three e2e tests land in the default test bucket.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -508,6 +523,281 @@ class GroupPartitionsExecSuite extends SharedSparkSession
{
assert(gpe.tryEnableSortedMerge().isEmpty)
}
}
+
+ test("SPARK-59310: basic counts without alignment") {
+ // Keys [1, 2, 1]: 3 input splits, key 1 coalesces partitions 0 and 2, 2
output partitions.
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(1))))
+ val gpe = GroupPartitionsExec(child)
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 0)
+ assert(gpe.metrics("numCoalescedPartitions").value === 1)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 2)
+ // Without expectedPartitionKeys there is no alignment, so its metrics
stay unregistered.
+ assert(!gpe.metrics.contains("numPrunedPartitions"))
+ assert(!gpe.metrics.contains("numReplicatedPartitions"))
+ }
+
+ test("SPARK-59310: zero coalesced without duplicate keys") {
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(3))))
+ val gpe = GroupPartitionsExec(child)
+ gpe.execute()
+
+ assert(gpe.metrics("numCoalescedPartitions").value === 0)
+ assert(gpe.metrics("numPartitions").value === 3)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 1)
+ }
+
+ test("SPARK-59310: distribute alignment pads and never replicates") {
+ // Child splits: key 1 -> [0], key 2 -> [1, 2]. Expected: key 1 x1, key 2
x3, key 3 x1.
+ // Key 2 spreads its 2 splits over 3 expected partitions (one empty pad)
and key 3 has no
+ // split (one more empty), so 5 output partitions with 2 empty and nothing
coalesced.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(2))))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 3, keyOf(3)
-> 1)),
+ distributePartitions = true)
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 5)
+ assert(gpe.metrics("numEmptyPartitions").value === 2)
+ assert(gpe.metrics("numPrunedPartitions").value === 0)
+ assert(gpe.metrics("numCoalescedPartitions").value === 0, "distribute
never coalesces")
+ assert(!gpe.metrics.contains("numReplicatedPartitions"), "distribute never
replicates")
+ }
+
+ test("SPARK-59310: alignment prunes unmatched keys, pads missing ones") {
+ // The expected keys carry key 1 and a key-3 slot the child does not hold,
as an inner
+ // join's intersection combined with the other side's layout would. The 2
splits of key 2
+ // cannot produce join output and never enter the alignment; the missing
key 3 pads two
+ // empty output partitions, and being empty, replicates nothing despite
its 2 slots.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(2))))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(3) -> 2)))
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 3)
+ assert(gpe.metrics("numPartitions").value === 3)
+ assert(gpe.metrics("numPrunedPartitions").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 2)
+ assert(gpe.metrics("numCoalescedPartitions").value === 0)
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 1)
+ assert(gpe.metrics("numReplicatedPartitions").value === 0,
+ "empty groups replicate nothing, and the single-split key 1 has no copy")
+ }
+
+ test("SPARK-59310: replicate alignment counts the reads beyond the first") {
+ // The other join side expects 2 partitions for key 1, so this side's
splits for the key are
+ // replicated to both: the slot beyond the first re-reads both splits, 2
extra input
+ // partition reads. Each output partition also coalesces the 2 splits of
the key.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(1))))
+ val gpe = GroupPartitionsExec(child, expectedPartitionKeys =
Some(Seq(keyOf(1) -> 2)))
+ gpe.execute()
+
+ assert(gpe.metrics("numInputPartitions").value === 2)
+ assert(gpe.metrics("numPartitions").value === 2)
+ assert(gpe.metrics("numReplicatedPartitions").value === 2)
+ assert(gpe.metrics("numCoalescedPartitions").value === 2, "both copies
merge the 2 splits")
+ assert(gpe.metrics("maxPartitionsPerGroup").value === 2)
+ assert(gpe.metrics("numEmptyPartitions").value === 0)
+ assert(gpe.metrics("numPrunedPartitions").value === 0)
+ }
+
+ test("SPARK-59310: an inner join intersection prunes both sides") {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withTable(s"testcat.ns.$items", s"testcat.ns.$purchases") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 44.0, cast('2020-01-15' as timestamp)), " +
+ s"(1, 45.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 11.0, cast('2020-01-01' as timestamp)), " +
+ s"(4, 19.5, cast('2020-02-01' as timestamp))")
+
+ // The sides hold different keys ({1,2,3} vs {1,2,4}), so grouping
alone cannot align
+ // them and the join pushes the inner intersection {1,2} down as the
expected keys.
+ // Each side then coalesces its duplicate-key splits (items merges two
groups of two,
+ // purchases one group of three) and prunes the one split of its
unmatched key
+ // (items key 3, purchases key 4); nothing is empty or replicated.
+ // Both AQE arms: the query has no shuffle, so the executed nodes and
their accumulator
+ // ids are the same with and without AQE, and the reporting chain must
work in both.
+ Seq(false, true).foreach { aqeEnabled =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->
aqeEnabled.toString) {
+ val df = sql(s"SELECT i.id FROM testcat.ns.$items i JOIN
testcat.ns.$purchases p " +
+ "ON i.id = p.item_id")
+ val previousExecutionIds = currentExecutionIds()
+ checkAnswer(df, Seq.fill(6)(Row(1L)) ++ Seq.fill(2)(Row(2L)))
+ val executionIds = currentExecutionIds().diff(previousExecutionIds)
+ assert(executionIds.size === 1)
+ val executionId = executionIds.head
+
+ // The metrics must survive the full reporting path: set on the
driver during
+ // doExecute, posted to the listener bus, and readable back from
the status store
+ // the SQL UI renders.
+ val metricValues = statusStore.executionMetrics(executionId)
+ val groupNodes =
+ statusStore.planGraph(executionId).nodes.filter(_.name ==
"GroupPartitions")
+ assert(groupNodes.size === 2, "one GroupPartitionsExec per join
side")
+ groupNodes.foreach { node =>
+ assert(metricValue(metricValues, node, "number of input
partitions") === "5")
+ assert(metricValue(metricValues, node, "number of partitions")
=== "2")
+ assert(metricValue(metricValues, node, "number of empty
partitions") === "0")
+ assert(metricValue(metricValues, node, "number of pruned input
partitions") === "1")
+ assert(metricValue(metricValues, node,
+ "number of replicated input partition reads") === "0",
+ "no expected key carries multiple splits, so nothing is
replicated")
+ }
+ assert(groupNodes.map(metricValue(metricValues, _, "number of
coalesced partitions"))
+ .sorted === Seq("1", "2"))
+ assert(groupNodes.map(metricValue(metricValues, _, "max partitions
per group"))
+ .sorted === Seq("2", "3"))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-59310: a disjoint inner join prunes both sides to empty end to
end") {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTable(s"testcat.ns.$items", s"testcat.ns.$purchases") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp))")
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(3, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(4, 44.0, cast('2020-01-15' as timestamp))")
+
+ // The key sets {1, 2} and {3, 4} are disjoint, so the inner join's
intersection is
+ // empty: the alignment emits no output partition on either side, each
side prunes both
+ // of its inputs, and doExecute takes the empty-RDD branch. The
metrics are sent before
+ // that branch, so the total pruning still reaches the store.
+ val df = sql(s"SELECT i.id FROM testcat.ns.$items i JOIN
testcat.ns.$purchases p " +
+ "ON i.id = p.item_id")
+ val previousExecutionIds = currentExecutionIds()
+ checkAnswer(df, Nil)
+ val executionIds = currentExecutionIds().diff(previousExecutionIds)
+ assert(executionIds.size === 1)
+ val executionId = executionIds.head
+
+ val metricValues = statusStore.executionMetrics(executionId)
+ val groupNodes =
+ statusStore.planGraph(executionId).nodes.filter(_.name ==
"GroupPartitions")
+ assert(groupNodes.size === 2, "one GroupPartitionsExec per join side")
+ groupNodes.foreach { node =>
+ assert(metricValue(metricValues, node, "number of input partitions")
=== "2")
+ assert(metricValue(metricValues, node, "number of partitions") ===
"0")
+ assert(metricValue(metricValues, node, "number of pruned input
partitions") === "2")
+ assert(metricValue(metricValues, node, "number of empty partitions")
=== "0")
+ assert(metricValue(metricValues, node, "number of coalesced
partitions") === "0")
+ assert(metricValue(metricValues, node, "max partitions per group")
=== "0")
+ assert(metricValue(metricValues, node,
+ "number of replicated input partition reads") === "0")
+ }
+ }
+ }
+ }
+
+ test("SPARK-59310: partial clustering replicates the smaller side") {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTable(s"testcat.ns.$items", s"testcat.ns.$purchases") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 41.0, cast('2020-01-02' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 45.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 50.0, cast('2020-01-02' as timestamp)), " +
+ s"(1, 55.0, cast('2020-01-02' as timestamp)), " +
+ s"(2, 15.0, cast('2020-01-02' as timestamp)), " +
+ s"(2, 20.0, cast('2020-01-03' as timestamp)), " +
+ s"(2, 22.0, cast('2020-01-03' as timestamp)), " +
+ s"(3, 20.0, cast('2020-02-01' as timestamp))")
+
+ // Partial clustering picks the side with fewer splits to replicate:
items (4 splits,
+ // key 1 x2, key 2 x1, key 3 x1) groups per key and copies each group
into every
+ // expected slot, while purchases (7 splits) keeps them, one per slot.
The slots come
+ // from purchases: key 1 x3, key 2 x3, key 3 x1. Items' extra reads:
(3-1) x 2 splits
+ // for key 1, (3-1) x 1 for key 2, none for key 3's single slot.
+ val df = sql(s"SELECT i.id FROM testcat.ns.$items i JOIN
testcat.ns.$purchases p " +
+ "ON i.id = p.item_id")
+ val previousExecutionIds = currentExecutionIds()
+ checkAnswer(df, Seq.fill(6)(Row(1L)) ++ Seq.fill(3)(Row(2L)) ++
Seq(Row(3L)))
+ val executionIds = currentExecutionIds().diff(previousExecutionIds)
+ assert(executionIds.size === 1)
+ val executionId = executionIds.head
+
+ val metricValues = statusStore.executionMetrics(executionId)
+ val groupNodes =
+ statusStore.planGraph(executionId).nodes.filter(_.name ==
"GroupPartitions")
+ assert(groupNodes.size === 2, "one GroupPartitionsExec per join side")
+ val byInputPartitions = groupNodes.map { node =>
+ metricValue(metricValues, node, "number of input partitions") -> node
+ }.toMap
+ assert(byInputPartitions.keySet === Set("4", "7"))
+ val replicatedSide = byInputPartitions("4")
+ val distributeSide = byInputPartitions("7")
+ Seq(replicatedSide, distributeSide).foreach { node =>
+ assert(metricValue(metricValues, node, "number of partitions") ===
"7")
+ assert(metricValue(metricValues, node, "number of empty partitions")
=== "0")
+ assert(metricValue(metricValues, node, "number of pruned input
partitions") === "0")
+ }
+ assert(metricValue(metricValues, replicatedSide,
+ "number of replicated input partition reads") === "6")
+ assert(!distributeSide.metrics.exists(
+ _.name == "number of replicated input partition reads"),
+ "the distribute side holds one split per slot and registers no
replicated metric")
+ assert(metricValue(metricValues, replicatedSide, "number of coalesced
partitions") === "3",
+ "the three copies of key 1's two-split group each merge their
splits")
+ assert(metricValue(metricValues, replicatedSide, "max partitions per
group") === "2")
+ assert(metricValue(metricValues, distributeSide, "number of coalesced
partitions") === "0")
+ assert(metricValue(metricValues, distributeSide, "max partitions per
group") === "1")
+ }
+ }
+ }
+
+ private case class ExecutableKeyedLeaf(kp: KeyedPartitioning)
Review Comment:
Nit: this is nested inside the suite class, whereas `DummyLeafSparkPlan` and
`DummySparkPlan` in this same file are top level. An inner-class `SparkPlan`
carries an outer reference, which can break `TreeNode.makeCopy` reflection;
nothing fails here because the tests only call `execute()`, but moving it out
would match the file's convention.
--
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]
