You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/09/08 01:23:50 UTC
(spark) branch branch-4.x updated: [SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 813f474e9008 [SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output
813f474e9008 is described below
commit 813f474e900818342bfdea2840f838dcb8a81975
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Tue Sep 8 09:23:16 2026 +0800
[SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output
### What changes were proposed in this pull request?
Keep the reported key-grouped partitioning when some of its keys are column-pruned out of the scan
output, instead of dropping the whole partitioning, so a storage-partitioned join still applies.
- `V2ScanPartitioningAndOrdering` keeps the reported keys as long as at least one of them is still in
the scan output. The full key list is kept, not the surviving subset, to stay positionally aligned
with the raw partition-key rows.
- `DataSourceV2ScanExecBase` splits the two views of that report. `reportedKeyedPartitioning` is the
source's own, at full key width, and is what reads and orders the raw `HasPartitionKey` rows
(`BatchScanExec.filteredPartitions`). `outputPartitioning` projects the pruned key positions away,
so the partitioning Spark plans against only references output columns. `replanWithRuntimeFilters`
takes `Option[KeyedPartitioning]`, so the projected view cannot reach the raw key rows.
- `DataSourceV2ScanRelation.doCanonicalize` and `BatchScanExec` (`doCanonicalize` / `equals` /
`hashCode`, through a shared `prunedKeyGroupedPartitioning`) ignore partitioning keys that
reference pruned columns, and the canonicalized ordering keeps only its leading run of sort orders
over output columns. Otherwise a dangling attribute's exprId keeps otherwise-equivalent scans
unequal and defeats subplan merging and exchange reuse.
Whether the surviving projection is usable is decided one layer down, by `KeyedPartitioning`:
pruning a key can map keys that were distinct onto the same key, and such a projection is marked
`isCollapsed`, so grouping it stays gated on
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` in `mayGroupToSatisfy`. A
projection that collapses nothing leaves the keys unique and needs no opt-in. No config doc changes:
the config keeps the roles it already has.
### Why are the changes needed?
An extra partition key that is column-pruned out of the scan output should not disable a
storage-partitioned join. Dropping the whole report also lost it for the case where the surviving
keys are still unique, which is sound with no opt-in at all.
Background: SPARK-40429 (#37886) introduced the drop-to-`None` guard to avoid a "missing inputs"
plan-validation error when a partitioning key referenced a pruned column. That failure mode no
longer applies: SPARK-40259 made `DataSourceV2ScanRelation.references` return empty (partitioning,
ordering and pushed filters are scan metadata, not references to resolve), so a dangling
partitioning key no longer trips plan validation. This change relies on that. SPARK-40259 is on
`master` only, so this change is master-only too and is not backportable to `branch-4.x` as it
stands.
### Does this PR introduce _any_ user-facing change?
Yes. When a partition key is pruned out of the scan output, Spark now reports the projected
partitioning instead of no partitioning, so a query that previously shuffled can plan a
storage-partitioned join:
- with `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` off (default), when
the surviving keys are still unique;
- with it on, also when pruning collapses distinct keys onto the same key.
Results are unchanged either way.
### How was this patch tested?
New tests in `KeyGroupedPartitioningSuite`:
- a join whose key is a subset of the partition keys, with the extra key pruned, triggers SPJ with
`allowKeysSubsetOfPartitionKeys` on and shuffles with it off;
- a pruned key that collapses nothing keeps SPJ with the config off;
- a pruned *leading* key keeps the raw key rows aligned: one test on row loss, one on the
`ClassCastException` from reading a leading `String` key at the projected `Integer` type;
- runtime filtering with a pruned leading key re-plans the scan's partitions against the full-width
key rows. This one lives in `KeyGroupedPartitioningRuntimeFilterTests`, so it runs against both
runtime-filtering fixtures; only the V2 one prunes a column, since the Catalyst fixture's scan
keeps the full table schema;
- all partition keys pruned reports no partitioning, even with the config on;
- a self-join with a pruned key stays canonicalizable, and a pruned key does not defeat plan reuse.
Pointing `filteredPartitions` at the projected partitioning fails exactly the three tests that pin
it (row loss, `ClassCastException`, runtime filtering) and nothing else.
`MergeSubplansSuite`: identical scans whose reported ordering is on a pruned column are
deduplicated, not fused. Verified this test fails without the ordering half of the canonicalization
change.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (qwen3.8-max, Claude Opus 5)
Closes #58522 from ulysses-you/worktree-spj-columnprune.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
(cherry picked from commit cfbec6d3d3ee545e1f58a9f17b5eb33f296d830e)
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../datasources/v2/DataSourceV2Relation.scala | 14 +-
.../execution/datasources/v2/BatchScanExec.scala | 20 +-
.../datasources/v2/DataSourceV2ScanExecBase.scala | 38 ++-
.../execution/datasources/v2/PushDownUtils.scala | 34 ++-
.../v2/V2ScanPartitioningAndOrdering.scala | 20 +-
.../sql/execution/planmerging/PlanMerger.scala | 6 +-
.../DataSourceV2CatalystRuntimeFilterSuite.scala | 2 +-
.../connector/KeyGroupedPartitioningSuite.scala | 329 ++++++++++++++++++++-
.../execution/planmerging/MergeSubplansSuite.scala | 27 ++
9 files changed, 447 insertions(+), 43 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
index d74cffb14f41..0cce60d42f3a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
@@ -343,11 +343,21 @@ case class DataSourceV2ScanRelation(
output = this.relation.output.map(QueryPlan.normalizeExpressions(_, this.relation.output))
),
output = this.output.map(QueryPlan.normalizeExpressions(_, this.output)),
+ // keyGroupedPartitioning may reference columns pruned out of `output`, which is kept as long
+ // as any key survives. A pruned key carries no information for plan comparison, since the
+ // physical outputPartitioning projects it away, so drop it before normalizing; otherwise the
+ // dangling attribute's exprId would keep otherwise-equivalent scans unequal and defeat
+ // subplan merging.
keyGroupedPartitioning = keyGroupedPartitioning.map(
- _.map(QueryPlan.normalizeExpressions(_, output))
+ _.filter(_.references.subsetOf(outputSet))
+ .map(QueryPlan.normalizeExpressions(_, output))
),
+ // ordering may likewise reference columns pruned out of `output`. Ordering is prefix-based,
+ // so keep only the leading run of sort orders that reference output columns and drop the rest
+ // before normalizing, for the same reason as keyGroupedPartitioning above.
ordering = ordering.map(
- _.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
+ _.takeWhile(_.references.subsetOf(outputSet))
+ .map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
),
// pushedFilters may reference columns pruned out of `output` (see the field doc), so they are
// normalized against the relation's full output rather than `output`.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
index 2e3394ac2a08..4b22f898f41e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
@@ -66,12 +66,22 @@ case class BatchScanExec(
case other: BatchScanExec =>
this.batch != null && this.batch == other.batch &&
this.runtimeFilters == other.runtimeFilters &&
- this.keyGroupedPartitioning == other.keyGroupedPartitioning
+ this.prunedKeyGroupedPartitioning == other.prunedKeyGroupedPartitioning
case _ =>
false
}
- override def hashCode(): Int = Objects.hash(batch, runtimeFilters, keyGroupedPartitioning)
+ override def hashCode(): Int = Objects.hash(batch, runtimeFilters, prunedKeyGroupedPartitioning)
+
+ /**
+ * The reported partitioning keys restricted to those still present in `output`. A key may
+ * reference a column pruned out of the scan, which is kept as long as any key survives (see
+ * V2ScanPartitioningAndOrdering). Such a dangling key carries no information, since the physical
+ * outputPartitioning projects it away, so `doCanonicalize`, `equals` and `hashCode` all use
+ * this view to stay consistent about ignoring it.
+ */
+ @transient lazy val prunedKeyGroupedPartitioning: Option[Seq[Expression]] =
+ keyGroupedPartitioning.map(_.filter(_.references.subsetOf(outputSet)))
@transient override lazy val inputPartitions: Seq[InputPartition] =
batch.planInputPartitions().toImmutableArraySeq
@@ -83,7 +93,7 @@ case class BatchScanExec(
runtimeFilters,
table,
output,
- outputPartitioning,
+ reportedKeyedPartitioning,
inputPartitions)
override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory()
@@ -106,8 +116,8 @@ case class BatchScanExec(
runtimeFilters = QueryPlan.normalizePredicates(
runtimeFilters.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral)),
output),
- keyGroupedPartitioning = keyGroupedPartitioning.map(_.map(
- QueryPlan.normalizeExpressions(_, output))))
+ keyGroupedPartitioning = prunedKeyGroupedPartitioning.map(
+ _.map(QueryPlan.normalizeExpressions(_, output))))
}
override def simpleString(maxFields: Int): String = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala
index cde8b1041a95..c897caa020e3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala
@@ -89,9 +89,18 @@ trait DataSourceV2ScanExecBase
|""".stripMargin
}
- // A `lazy val` because the planner asks a node for its partitioning many times, and the
- // key-grouped arm sorts every partition key, wraps each one and runs a `distinct` over them.
- @transient override lazy val outputPartitioning: physical.Partitioning = {
+ /**
+ * The partitioning as the source reported it: one key per input partition, holding every
+ * reported key position, with the partitions in the order those full keys sort into. It is built
+ * from the raw, full-width `HasPartitionKey.partitionKey()` rows, so a consumer of those rows
+ * (`filteredPartitions`) must take the keys, the key types and the order from here, not from the
+ * possibly-projected `outputPartitioning`.
+ *
+ * A `lazy val` because this is the expensive half: it sorts every partition key, wraps each one
+ * and runs a `distinct` over them, and both `outputPartitioning` and `filteredPartitions` ask
+ * for it.
+ */
+ @transient protected lazy val reportedKeyedPartitioning: Option[KeyedPartitioning] = {
keyGroupedPartitioning match {
case Some(exprs) if conf.v2BucketingEnabled && KeyedPartitioning.supportsExpressions(exprs) &&
inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) =>
@@ -99,12 +108,29 @@ trait DataSourceV2ScanExecBase
val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes)
val partitionKeys =
inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
- KeyedPartitioning(exprs, partitionKeys)
- case _ =>
- super.outputPartitioning
+ Some(KeyedPartitioning(exprs, partitionKeys))
+ case _ => None
}
}
+ // A `lazy val` because the planner asks a node for its partitioning many times, and each ask
+ // would otherwise re-project the reported keys.
+ @transient override lazy val outputPartitioning: physical.Partitioning =
+ reportedKeyedPartitioning match {
+ case Some(partitioning) =>
+ // A partition key may reference a column that was pruned out of the scan output (see
+ // V2ScanPartitioningAndOrdering). Project such unresolvable key positions away so the
+ // reported partitioning only references output columns.
+ val exprs = partitioning.expressions
+ val resolvablePositions = exprs.indices.filter(i => exprs(i).references.subsetOf(outputSet))
+ if (resolvablePositions.isEmpty) {
+ super.outputPartitioning
+ } else {
+ partitioning.project(resolvablePositions)
+ }
+ case _ => super.outputPartitioning
+ }
+
/**
* Returns the output ordering for this scan. When the source reports ordering via
* `SupportsReportOrdering`, that ordering is returned as-is. Otherwise, when the output
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
index 3a5796d3448f..d2b6c0bfa13c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
@@ -24,7 +24,7 @@ import org.apache.spark.internal.{Logging, LogKeys}
import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils}
import org.apache.spark.sql.catalyst.plans.logical.SampleMethod
-import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning}
+import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.catalyst.util.{CharVarcharUtils, InternalRowComparableWrapper}
@@ -273,23 +273,25 @@ object PushDownUtils extends Logging {
}
/**
- * Pushes runtime filters into `scan` and re-plans its input partitions. For scans whose
- * `outputPartitioning` is a [[KeyedPartitioning]] (SPJ-active), validates that the data source
- * preserved the original partitioning and pads with `None` to preserve key alignment with the
- * pre-filter partition set.
+ * Pushes runtime filters into `scan` and re-plans its input partitions. When the scan reported a
+ * [[KeyedPartitioning]] (SPJ-active), validates that the data source preserved the original
+ * partitioning and pads with `None` to preserve key alignment with the pre-filter partition set.
*
* Notes:
* - `filter` is mutating, and Spark may call this more than once for the same `scan` instance
* (see [[pushRuntimeFilters]]); successive calls are additive.
- * - When `outputPartitioning` is a [[KeyedPartitioning]], every split from
- * `planInputPartitions()` used on this path must implement [[HasPartitionKey]].
+ * - With a [[KeyedPartitioning]], every split from `planInputPartitions()` used on this path
+ * must implement [[HasPartitionKey]].
*
* @param scan the V2 scan to push filters into
* @param runtimeFilters runtime filters to translate and push
* @param table the table backing the scan, used to derive the partition-predicate
* schema for iterative [[PartitionPredicate]] pushdown
* @param output scan output attributes
- * @param outputPartitioning Spark-side output partitioning (used for SPJ validation)
+ * @param keyedPartitioning the partitioning as the source reported it, at its full key width.
+ * The raw [[HasPartitionKey]] rows are read and ordered against it, so
+ * a projected partitioning must not be passed here: it would read each
+ * key row at the wrong positions and types
* @param originalPartitions unfiltered partitions, consulted only when no runtime filters fire
* @return one entry per original input partition: `Some(part)` for surviving partitions and
* `None` for partition keys whose splits were entirely pruned (SPJ alignment)
@@ -299,15 +301,15 @@ object PushDownUtils extends Logging {
runtimeFilters: Seq[Expression],
table: Table,
output: Seq[AttributeReference],
- outputPartitioning: Partitioning,
+ keyedPartitioning: Option[KeyedPartitioning],
originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
if (filtered) {
// call toBatch again to get filtered partitions
val newPartitions = scan.toBatch.planInputPartitions()
- outputPartitioning match {
- case k: KeyedPartitioning =>
+ keyedPartitioning match {
+ case Some(k) =>
if (newPartitions.exists(!_.isInstanceOf[HasPartitionKey])) {
throw new SparkException("Data source must have preserved the original partitioning " +
"during runtime filtering: not all partitions implement HasPartitionKey after " +
@@ -344,22 +346,22 @@ object PushDownUtils extends Logging {
fps.map(Some).padTo(size, None)
}
- case _ =>
+ case None =>
// no validation is needed as the data source did not report any specific partitioning
newPartitions.toSeq.map(Some)
}
} else {
val parts = originalPartitions
- (outputPartitioning match {
- case k: KeyedPartitioning =>
+ (keyedPartitioning match {
+ case Some(k) =>
if (parts.exists(!_.isInstanceOf[HasPartitionKey])) {
throw new SparkException("Original partitions must implement HasPartitionKey when " +
- "outputPartitioning is KeyedPartitioning.")
+ "the scan reported a KeyedPartitioning.")
}
parts.sortBy(_.asInstanceOf[HasPartitionKey].partitionKey())(k.keyRowOrdering)
- case _ => parts
+ case None => parts
}).map(Some)
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala
index 165425aea19d..936da28f638e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala
@@ -50,15 +50,17 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging {
val partitioning = sequenceToOption(
kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation, relation.funCatalog))
.toImmutableArraySeq)
- if (partitioning.isEmpty) {
- None
- } else {
- if (partitioning.get.forall(p => p.references.subsetOf(d.outputSet))) {
- partitioning
- } else {
- None
- }
- }
+ // Keep the partitioning when at least one of its keys is still in the scan output: the
+ // scan projects the pruned key positions away when reporting its physical output
+ // partitioning (see DataSourceV2ScanExecBase.outputPartitioning). When no key survives,
+ // and likewise when the source reported no key at all, there is nothing to report.
+ // Grouping a projection that collapsed distinct keys onto the same key stays gated on
+ // allowKeysSubsetOfPartitionKeys one layer down, in KeyedPartitioning.mayGroupToSatisfy.
+ //
+ // A kept pruned key leaves a dangling attribute on the relation. What keeps that off
+ // `missingInput`, and so past the optimizer's plan-change validation, is the
+ // `DataSourceV2ScanRelation.references` override; see the comment there.
+ partitioning.filter(_.exists(_.references.subsetOf(d.outputSet)))
case _: UnknownPartitioning => None
case p =>
logWarning(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
index d84979b0b300..0cf12e8b5e1d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
@@ -981,9 +981,9 @@ class PlanMerger(
// the merged scan's split count and partition values can still differ from an input's, since it
// may push a different best-effort filter and so prune differently. And a report the merged scan
// GAINS is not a degradation either. For partitioning that is because an input dropped its own
- // only where a pruned column left the expressions inexpressible over that scan's output
- // (V2ScanPartitioningAndOrdering's partitioning pass is reference-subset guarded), not because
- // the source stopped reporting; the ordering pass has no such guard, so an ordering report is
+ // only where none of its keys survived in that scan's output (V2ScanPartitioningAndOrdering's
+ // partitioning pass keeps the report whenever any key survives), not because the source stopped
+ // reporting; the ordering pass has no such guard, so an ordering report is
// never dropped by pruning and a gained one can only come from the source. Either way, keeping it
// is exactly the win this merge is after.
private def mergeDegradesReporting(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
index 45cca0c40a68..04223078fe23 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
@@ -670,7 +670,7 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = {
val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter)
PushDownUtils.replanWithRuntimeFilters(scan, Seq(EqualTo(partAttr, Literal(1))), table,
- Seq(partAttr), partitioning, originalPartitions = Seq.empty)
+ Seq(partAttr), Some(partitioning), originalPartitions = Seq.empty)
}
val keyDropped = intercept[SparkException](replanAfterFiltering(Seq(new InputPartition {})))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 9120c01231f7..8d20faf1a93b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.execution.{
SparkPlan,
UnionExec}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec}
-import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements}
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements}
import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec}
import org.apache.spark.sql.functions.{col, max}
import org.apache.spark.sql.internal.SQLConf
@@ -325,6 +325,47 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui
}
}
}
+
+ test("SPARK-59248: runtime filtering with a pruned partition key keeps the full key rows") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ 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)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+ val pCols = Array(
+ Column.create("store_id", IntegerType),
+ Column.create("item_id", LongType),
+ Column.create("price", FloatType))
+ createTable(purchases, pCols, Array(identity("store_id"), identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(10, 1, 42.0), (20, 1, 44.0), (30, 2, 11.0), (40, 3, 19.5)")
+
+ // `store_id` is the leading partition key and is not selected, so it is pruned and the
+ // reported partitioning is projected onto item_id. The runtime filter on item_id re-plans the
+ // scan's partitions, which must stay keyed and ordered by the full (store_id, item_id) rows,
+ // not the projected item_id-only key types.
+ //
+ // `item_id` must match `items.id` in type. A narrower type makes the join cast it, and
+ // `translateRuntimeFilterV2` cannot translate a cast, so no filter reaches the scan and the
+ // re-planning path never runs.
+ //
+ // Only the V2 instance of this trait reaches that path with a pruned key. The Catalyst
+ // fixture's scan keeps the full table schema, so `store_id` stays in the output and nothing
+ // is projected away; that instance runs a no-pruning variant of the same query.
+ val df = sql(
+ s"SELECT p.item_id, p.price from testcat.ns.$items i, testcat.ns.$purchases p " +
+ "WHERE i.id = p.item_id AND i.price > 20.0 ORDER BY p.item_id, p.price")
+ checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f)))
+ }
+ }
}
@ExtendedSQLTest
@@ -2598,6 +2639,292 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59248: join key subset of partition keys, extra partition key pruned from the " +
+ "output") {
+ // Both tables are partitioned by (id, data). The join is only on `id`, and `data` is not
+ // selected, so it is column-pruned out of both scan outputs. The partitioning is kept and
+ // projected onto `id`; since `data` has several values per `id` the projection collapses the
+ // keys, so grouping them needs allowKeysSubsetOfPartitionKeys. With it SPJ triggers; without it
+ // both sides shuffle.
+ val table1 = "prune_t1"
+ val table2 = "prune_t2"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-01' as timestamp))")
+
+ createTable(table2, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'ee', cast('2020-01-01' as timestamp)), " +
+ "(4, 'ff', cast('2020-01-01' as timestamp))")
+
+ // Selecting only `id` prunes the other partition key `data` (and `ts`) from both scans. The
+ // expected result is the within-`id` cross product (id=2 matches 2 x 2 rows, id=3 matches
+ // 1 x 1).
+ val expected = Seq(Row(2), Row(2), Row(2), Row(2), Row(3))
+ val query =
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.id AS id
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.id = t2.id ORDER BY id
+ |""".stripMargin
+
+ Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowKeysSubsetOfPartitionKeys.toString) {
+ val df = sql(query)
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan)
+ if (allowKeysSubsetOfPartitionKeys) {
+ assert(shuffles.isEmpty, "SPJ should be triggered even though `data` is pruned")
+ assert(groupPartitions.nonEmpty, "GroupPartitionsExec should coalesce on the join key")
+ // The reported partitioning is kept on the scan even though `data` is pruned ...
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.nonEmpty)
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isDefined,
+ "partitioning should be kept despite the pruned key")
+ // ... but the physical output partitioning must only reference output columns, so the
+ // pruned column reaches no consumer (shuffle spec, ordering, plan equality).
+ scan.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.expressions.forall(_.references.subsetOf(scan.outputSet)),
+ s"partitioning ${kp.expressions} references a column outside ${scan.output}")
+ case other =>
+ fail(s"expected KeyedPartitioning but got $other")
+ }
+ }
+ } else {
+ assert(shuffles.nonEmpty, "SPJ should not be triggered without the config")
+ assert(groupPartitions.isEmpty)
+ }
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ test("SPARK-59248: a pruned leading partition key must not lose join rows") {
+ // The table is partitioned by (store_id, dept_id) and the leading key `store_id` is pruned, so
+ // the surviving partitioning is on dept_id while the raw partition-key rows still lead with
+ // store_id. The other side is unkeyed and is shuffled into the table's layout. Sorting the raw
+ // partitions by the projected (dept_id-only) key ordering reordered the partitions inside a
+ // store_id group against the advertised dept_id order and silently dropped the rows whose label
+ // swapped; the full-width input ordering keeps the two aligned.
+ val cols = Array(Column.create("store_id", IntegerType), Column.create("dept_id", IntegerType))
+ createTable("prune_lead_t", cols, Array(identity("store_id"), identity("dept_id")))
+ sql("INSERT INTO testcat.ns.prune_lead_t VALUES " +
+ "(1, 20), (1, 10), (1, 30), (2, 5), (2, 40), (3, 7), (3, 1)")
+
+ withTempView("other") {
+ spark.range(1, 41).selectExpr("cast(id as int) as dept_id").createOrReplaceTempView("other")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t", "o")}
+ |t.dept_id AS dept_id
+ |FROM testcat.ns.prune_lead_t t JOIN other o ON t.dept_id = o.dept_id
+ |ORDER BY dept_id
+ |""".stripMargin)
+ checkAnswer(df, Seq(1, 5, 7, 10, 20, 30, 40).map(Row(_)))
+ // Only the unkeyed side shuffles; assert the plan shape too, so a silent fall-back to
+ // shuffling both sides does not pass on rows alone.
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).size == 1,
+ s"only the unkeyed side of the join should shuffle:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59248: a pruned leading partition key of another type must not fail the scan") {
+ // Partitioned by (data, id) with the String `data` leading. The join is on `id` and `data` is
+ // pruned, so the projected partitioning is on the Integer `id` while the raw partition-key rows
+ // still lead with the String `data`. Sorting those rows with the projected (Integer-only) key
+ // ordering used to cast the leading String to an Integer and throw ClassCastException; the
+ // full-width input ordering reads each column at its own type.
+ val cols = Array(Column.create("data", StringType), Column.create("id", IntegerType))
+ val partition = Array(identity("data"), identity("id"))
+ createTable("prune_cce_t1", cols, partition)
+ sql("INSERT INTO testcat.ns.prune_cce_t1 VALUES ('a', 1), ('b', 2), ('c', 3)")
+ createTable("prune_cce_t2", cols, partition)
+ sql("INSERT INTO testcat.ns.prune_cce_t2 VALUES ('a', 1), ('b', 2)")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.id AS id
+ |FROM testcat.ns.prune_cce_t1 t1 JOIN testcat.ns.prune_cce_t2 t2
+ |ON t1.id = t2.id ORDER BY id
+ |""".stripMargin)
+ // Each id is unique per table, so the projection collapses nothing and SPJ applies with no
+ // exchange; assert the plan shape, not just the rows, so a regression to shuffle is caught.
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "SPJ should be triggered for the pruned leading key of another type")
+ checkAnswer(df, Seq(Row(1), Row(2)))
+ }
+ }
+
+ test("SPARK-59248: a pruned key that collapses nothing keeps SPJ without the config") {
+ // Partitioned by (dept_id, store_id) with one store_id per dept_id, so pruning store_id leaves
+ // the projected dept_id keys unique. No grouping is needed, so
+ // allowKeysSubsetOfPartitionKeys is not required and SPJ applies with the config off.
+ val cols = Array(Column.create("dept_id", IntegerType), Column.create("store_id", IntegerType))
+ val partition = Array(identity("dept_id"), identity("store_id"))
+ createTable("prune_nocollapse_t1", cols, partition)
+ sql("INSERT INTO testcat.ns.prune_nocollapse_t1 VALUES (10, 100), (20, 200), (30, 300)")
+ createTable("prune_nocollapse_t2", cols, partition)
+ sql("INSERT INTO testcat.ns.prune_nocollapse_t2 VALUES (10, 100), (20, 200)")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.dept_id AS dept_id
+ |FROM testcat.ns.prune_nocollapse_t1 t1 JOIN testcat.ns.prune_nocollapse_t2 t2
+ |ON t1.dept_id = t2.dept_id ORDER BY dept_id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "a non-collapsing pruned key should keep SPJ without the config")
+ checkAnswer(df, Seq(Row(10), Row(20)))
+ }
+ }
+
+ test("SPARK-59248: scan reports no partitioning when all partition keys are pruned") {
+ // The table is partitioned by (id, data), but the query selects only `ts`, so both partition
+ // keys are column-pruned out of the scan output. Even with allowKeysSubsetOfPartitionKeys on,
+ // no partition key survives in the output, so the scan must not keep a dangling
+ // KeyedPartitioning and reports no (unknown) partitioning.
+ val table1 = "prune_all_keys"
+ createTable(table1, columns, Array(identity("id"), identity("data")))
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-02' as timestamp))")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(s"SELECT ts FROM testcat.ns.$table1")
+ checkAnswer(df, Seq(
+ Row(Timestamp.valueOf("2020-01-01 00:00:00")),
+ Row(Timestamp.valueOf("2020-01-02 00:00:00"))))
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.length == 1)
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isEmpty,
+ s"no partition key survives in the output, got ${scan.keyGroupedPartitioning}")
+ scan.outputPartitioning match {
+ case _: physical.UnknownPartitioning => // expected: nothing left to partition by
+ case other => fail(s"expected UnknownPartitioning but got $other")
+ }
+ }
+ }
+ }
+
+ test("SPARK-59248: self-join with a pruned partition key keeps plans canonicalizable") {
+ // Same-table join where the extra partition key `data` is pruned from both scan instances. This
+ // exercises canonicalization/plan-equality over scans whose reported partitioning carries a key
+ // that is not in the scan output; results must stay correct and planning must not fail.
+ val table1 = "prune_self"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-01' as timestamp))")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("a", "b")}
+ |a.id AS id
+ |FROM testcat.ns.$table1 a JOIN testcat.ns.$table1 b
+ |ON a.id = b.id ORDER BY id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty, "SPJ should be triggered")
+ // id=1 yields 1 row, id=2 yields 2 x 2 = 4 rows, id=3 yields 1 row.
+ checkAnswer(df, Seq(Row(1), Row(2), Row(2), Row(2), Row(2), Row(3)))
+
+ // Both scan instances must survive (no incorrect dedup) and each must report a partitioning
+ // that only references its own output, even though `data` is pruned: this is what keeps the
+ // dangling key from reaching any consumer (shuffle spec, ordering, canonicalized comparison).
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.length == 2, s"expected the two self-join scans, got:\n" +
+ s"${df.queryExecution.executedPlan}")
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isDefined,
+ "partitioning should be kept despite the pruned key")
+ scan.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.expressions.forall(_.references.subsetOf(scan.outputSet)),
+ s"partitioning ${kp.expressions} references a column outside ${scan.output}")
+ case other =>
+ fail(s"expected KeyedPartitioning but got $other")
+ }
+ // Canonicalization must be stable and must not throw with a dangling key present.
+ assert(scan.canonicalized.sameResult(scan.canonicalized))
+ }
+ }
+ }
+
+ test("SPARK-59248: a pruned partition key must not defeat plan reuse") {
+ val table1 = "prune_reuse"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-02' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-03' as timestamp))")
+
+ // Self-join on the non-partition column `ts`; the other partition key `data` is pruned from
+ // both scan instances. The two legs are identical subtrees, so Spark reuses one leg's exchange
+ // for the other. With allowKeysSubsetOfPartitionKeys the scan keeps its reported partitioning,
+ // which then references the pruned `data`; that dangling key must not leak into canonicalized
+ // plan comparison and break the reuse.
+ val query =
+ s"""
+ |SELECT a.id AS id1, b.id AS id2
+ |FROM testcat.ns.$table1 a JOIN testcat.ns.$table1 b
+ |ON a.ts = b.ts ORDER BY id1, id2
+ |""".stripMargin
+
+ Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowKeysSubsetOfPartitionKeys.toString) {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 2), Row(3, 3)))
+ val plan = df.queryExecution.executedPlan
+ val reused = collect(plan) { case r: ReusedExchangeExec => r }
+ val scans = collectScans(plan)
+ assert(scans.length == 1,
+ s"the two identical legs should reuse a single scan " +
+ s"(allowKeysSubsetOfPartitionKeys=$allowKeysSubsetOfPartitionKeys):\n$plan")
+ assert(reused.length == 1,
+ s"expected one reused exchange " +
+ s"(allowKeysSubsetOfPartitionKeys=$allowKeysSubsetOfPartitionKeys):\n$plan")
+ }
+ }
+ }
+
test("SPARK-47094: SPJ: Support compatible buckets") {
val table1 = "tab1e1"
val table2 = "table2"
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
index cd16ac4ea665..849b95e0a927 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
@@ -2818,6 +2818,33 @@ class MergeSubplansSuite extends PlanTest {
comparePlans(Optimize.execute(q.analyze), q.analyze)
}
+ test("SPARK-59248: identical DSv2 scans whose reported ordering is on a pruned column are " +
+ "deduplicated, not fused") {
+ // The two subqueries read the same column and compute the same aggregate, so they are
+ // identical and the identical-plan check should deduplicate them (plan left unchanged). The
+ // table reports an ordering on `b`, which is pruned out of each scan's output (only `a` is
+ // read). Canonicalize must drop that dangling ordering: otherwise the two scans' ordering
+ // attributes carry different exprIds, the identical check fails, and the plans are wrongly
+ // fused into a CTE.
+ val table = new TestV2Table(
+ StructType(Seq(
+ StructField("a", IntegerType),
+ StructField("b", IntegerType),
+ StructField("c", StringType))),
+ reportedOrderingCols = Seq("b"))
+ val q = testRelation.select(
+ ScalarSubquery(v2ScanReportingOn(table, Seq("a")).groupBy()(sum($"a").as("sum_a"))),
+ ScalarSubquery(v2ScanReportingOn(table, Seq("a")).groupBy()(sum($"a").as("sum_a"))))
+
+ val optimized = Optimize.execute(q.analyze)
+ // Deduplicated, not fused: both subqueries survive and no merged CTE scan is introduced.
+ assert(v2Scans(optimized).length == 2,
+ s"the two identical scans must be deduplicated, not fused into one:\n$optimized")
+ assert(!optimized.isInstanceOf[WithCTE],
+ s"identical subqueries must not be extracted to a merged CTE:\n$optimized")
+ comparePlans(optimized, q.analyze)
+ }
+
test("SPARK-58549: enforce the required report on the deferred under-Filter scan build") {
// The above tests fuse scans directly under an Aggregate (no Filter), so they exercise the
// scan build at the leaf. When the scans sit under an (identical) Filter the build is instead
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org