From 560b5b0ad1bc6eeb37f89ff2e2b48e53098adfe5 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Fri, 4 Sep 2026 18:36:01 +0800 Subject: [PATCH 1/3] [SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output When a key-grouped partitioning key is column-pruned out of the scan output, the reported partitioning is kept (with allowKeysSubsetOfPartitionKeys enabled) and projected onto the surviving keys, so storage-partitioned join still applies. Canonicalization and plan equality ignore such pruned keys so subplan merging and exchange reuse are unaffected. Assisted-by: Claude Fable 5 --- .../datasources/v2/DataSourceV2Relation.scala | 14 +- .../datasources/v2/BatchScanExec.scala | 18 +- .../v2/DataSourceV2ScanExecBase.scala | 12 +- .../v2/V2ScanPartitioningAndOrdering.scala | 59 +++-- .../KeyGroupedPartitioningSuite.scala | 229 +++++++++++++++++- .../planmerging/MergeSubplansSuite.scala | 27 +++ 6 files changed, 328 insertions(+), 31 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 d74cffb14f41b..349da8bf9358b 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` (kept when operation + // keys may be a subset of the partition keys). 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 2e3394ac2a082..7bebaa217c465 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 (kept when operation keys may be a subset of the + * partition keys). 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 @@ -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 cde8b1041a952..8835ec9b14ee4 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 @@ -99,7 +99,17 @@ trait DataSourceV2ScanExecBase val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes) val partitionKeys = inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering) - KeyedPartitioning(exprs, partitionKeys) + val partitioning = KeyedPartitioning(exprs, partitionKeys) + // A partition key may reference a column that was pruned out of the scan output (kept only + // when operation keys may be a subset of the partition keys, see + // V2ScanPartitioningAndOrdering). Project such unresolvable key positions away so the + // reported partitioning only references output columns. + val resolvablePositions = exprs.indices.filter(i => exprs(i).references.subsetOf(outputSet)) + if (resolvablePositions.isEmpty) { + super.outputPartitioning + } else { + partitioning.project(resolvablePositions) + } case _ => super.outputPartitioning } 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 165425aea19dc..6f20d91191f77 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 @@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_SCAN_RELATION import org.apache.spark.sql.connector.read.{SupportsReportOrdering, SupportsReportPartitioning} import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, UnknownPartitioning} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.collection.Utils.sequenceToOption @@ -41,33 +42,45 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { } } - private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning( + private def partitioning(plan: LogicalPlan) = { + val allowKeysSubsetOfPartitionKeys = SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys + plan.transformDownWithPruning( _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { - case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) - if d.keyGroupedPartitioning.isEmpty => - val catalystPartitioning = scan.outputPartitioning() match { - case kgp: KeyGroupedPartitioning => - 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 { + case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) + if d.keyGroupedPartitioning.isEmpty => + val catalystPartitioning = scan.outputPartitioning() match { + case kgp: KeyGroupedPartitioning => + val partitioning = sequenceToOption( + kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation, relation.funCatalog)) + .toImmutableArraySeq) + if (partitioning.isEmpty) { None + } else { + val inOutput = partitioning.get.map(p => p.references.subsetOf(d.outputSet)) + if (inOutput.forall(identity)) { + partitioning + } else if (inOutput.exists(identity) && allowKeysSubsetOfPartitionKeys) { + // Some partition keys were pruned out of the scan output. Keep the full + // partitioning when operation keys may be a subset of the partition keys: the scan + // projects the unresolvable key positions away when reporting its physical output + // partitioning (see DataSourceV2ScanExecBase.outputPartitioning). Keeping the full + // list, rather than the resolvable subset, preserves positional alignment with the + // partition keys. + partitioning + } else { + None + } } - } - case _: UnknownPartitioning => None - case p => - logWarning( - log"Spark ignores the partitioning ${MDC(CLASS_NAME, p.getClass.getSimpleName)}. " + - log"Please use KeyGroupedPartitioning for better performance") - None - } + case _: UnknownPartitioning => None + case p => + logWarning( + log"Spark ignores the partitioning ${MDC(CLASS_NAME, p.getClass.getSimpleName)}. " + + log"Please use KeyGroupedPartitioning for better performance") + None + } - d.copy(keyGroupedPartitioning = catalystPartitioning) + d.copy(keyGroupedPartitioning = catalystPartitioning) + } } private def ordering(plan: LogicalPlan) = plan.transformDownWithPruning( 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 02fc8afb7d9de..df66cf6750879 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 @@ -42,7 +42,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 @@ -2597,6 +2597,233 @@ 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. Without + // allowKeysSubsetOfPartitionKeys the pruned `data` key drops the reported partitioning and both + // sides shuffle; with it, the partitioning is kept and projected onto `id`, so SPJ triggers. + 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: 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-59248: a pruned source-reported ordering must not defeat exchange reuse") { + val table1 = "ord_reuse" + // Partitioned by `id` (kept in the output) and the source also reports an ordering on `data`. + // The self-join prunes `data`, leaving only the reported ordering dangling; that dangling + // ordering must not keep the two identical legs from being reused. + createTable(table1, columns, Array(identity("id")), + Array(sort(column("data"), SortDirection.ASCENDING))) + 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))") + + 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 + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql(query) + checkAnswer(df, Seq(Row(1, 1), Row(2, 2), Row(3, 3))) + val plan = df.queryExecution.executedPlan + // The scan still reports the (dangling) ordering; the point is that it does not block reuse. + collectScans(plan).foreach { scan => + assert(scan.ordering.exists(_.exists(_.child.references.exists(_.name == "data"))), + s"expected the scan to report an ordering on the pruned `data`:\n$plan") + } + val scans = collectScans(plan) + val reused = collect(plan) { case r: ReusedExchangeExec => r } + assert(scans.length == 1, s"the two identical legs should reuse a single scan:\n$plan") + assert(reused.length == 1, s"expected one reused exchange:\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 cd16ac4ea6657..849b95e0a9279 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 From 17964d54b2d1ee69431a9d0f54cc7b9ea25e4d3c Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Mon, 7 Sep 2026 10:21:44 +0800 Subject: [PATCH 2/3] [SPARK-59248][SQL] Separate the reported and projected scan partitioning DataSourceV2ScanExecBase now reports two views of a key-grouped partitioning: reportedKeyedPartitioning at the source's full key width, which reads and orders the raw HasPartitionKey rows, and outputPartitioning with the pruned key positions projected away, which is what Spark plans against. Mixing them reads each key row at the wrong positions and types, losing join rows or throwing ClassCastException on a pruned leading key of another type, so replanWithRuntimeFilters takes Option[KeyedPartitioning]. V2ScanPartitioningAndOrdering no longer gates the report on allowKeysSubsetOfPartitionKeys. It has no partition key values, so it cannot tell a projection that collapses distinct keys onto one from a projection that leaves them unique; KeyedPartitioning.mayGroupToSatisfy makes that distinction and keeps the config over the collapsing case, while a non-collapsing projection is sound with no opt-in. Assisted-by: Claude Opus 5 --- .../datasources/v2/DataSourceV2Relation.scala | 10 +- .../datasources/v2/BatchScanExec.scala | 6 +- .../v2/DataSourceV2ScanExecBase.scala | 34 +++- .../datasources/v2/PushDownUtils.scala | 34 ++-- .../v2/V2ScanPartitioningAndOrdering.scala | 65 +++---- .../execution/planmerging/PlanMerger.scala | 6 +- ...taSourceV2CatalystRuntimeFilterSuite.scala | 2 +- .../KeyGroupedPartitioningSuite.scala | 168 ++++++++++++++---- 8 files changed, 212 insertions(+), 113 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 349da8bf9358b..0cce60d42f3aa 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,11 @@ 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` (kept when operation - // keys may be a subset of the partition keys). 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 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( _.filter(_.references.subsetOf(outputSet)) .map(QueryPlan.normalizeExpressions(_, 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 7bebaa217c465..4b22f898f41e3 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 @@ -75,8 +75,8 @@ case class BatchScanExec( /** * The reported partitioning keys restricted to those still present in `output`. A key may - * reference a column pruned out of the scan (kept when operation keys may be a subset of the - * partition keys). Such a dangling key carries no information, since the physical + * 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. */ @@ -93,7 +93,7 @@ case class BatchScanExec( runtimeFilters, table, output, - outputPartitioning, + reportedKeyedPartitioning, inputPartitions) override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory() 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 8835ec9b14ee4..c897caa020e3b 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,21 +108,28 @@ trait DataSourceV2ScanExecBase val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes) val partitionKeys = inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering) - val partitioning = KeyedPartitioning(exprs, partitionKeys) - // A partition key may reference a column that was pruned out of the scan output (kept only - // when operation keys may be a subset of the partition keys, see + 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 + case _ => super.outputPartitioning } - } /** * Returns the output ordering for this scan. When the source reports ordering via 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 3a5796d3448fc..d2b6c0bfa13c4 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 6f20d91191f77..936da28f638ed 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 @@ -24,7 +24,6 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_SCAN_RELATION import org.apache.spark.sql.connector.read.{SupportsReportOrdering, SupportsReportPartitioning} import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, UnknownPartitioning} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.collection.Utils.sequenceToOption @@ -42,45 +41,35 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { } } - private def partitioning(plan: LogicalPlan) = { - val allowKeysSubsetOfPartitionKeys = SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys - plan.transformDownWithPruning( + private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning( _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { - case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) - if d.keyGroupedPartitioning.isEmpty => - val catalystPartitioning = scan.outputPartitioning() match { - case kgp: KeyGroupedPartitioning => - val partitioning = sequenceToOption( - kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation, relation.funCatalog)) - .toImmutableArraySeq) - if (partitioning.isEmpty) { - None - } else { - val inOutput = partitioning.get.map(p => p.references.subsetOf(d.outputSet)) - if (inOutput.forall(identity)) { - partitioning - } else if (inOutput.exists(identity) && allowKeysSubsetOfPartitionKeys) { - // Some partition keys were pruned out of the scan output. Keep the full - // partitioning when operation keys may be a subset of the partition keys: the scan - // projects the unresolvable key positions away when reporting its physical output - // partitioning (see DataSourceV2ScanExecBase.outputPartitioning). Keeping the full - // list, rather than the resolvable subset, preserves positional alignment with the - // partition keys. - partitioning - } else { - None - } - } - case _: UnknownPartitioning => None - case p => - logWarning( - log"Spark ignores the partitioning ${MDC(CLASS_NAME, p.getClass.getSimpleName)}. " + - log"Please use KeyGroupedPartitioning for better performance") - None - } + case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) + if d.keyGroupedPartitioning.isEmpty => + val catalystPartitioning = scan.outputPartitioning() match { + case kgp: KeyGroupedPartitioning => + val partitioning = sequenceToOption( + kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation, relation.funCatalog)) + .toImmutableArraySeq) + // 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( + log"Spark ignores the partitioning ${MDC(CLASS_NAME, p.getClass.getSimpleName)}. " + + log"Please use KeyGroupedPartitioning for better performance") + None + } - d.copy(keyGroupedPartitioning = catalystPartitioning) - } + d.copy(keyGroupedPartitioning = catalystPartitioning) } private def ordering(plan: LogicalPlan) = plan.transformDownWithPruning( 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 d84979b0b300b..0cf12e8b5e1d2 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 45cca0c40a68d..04223078fe23d 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 df66cf6750879..ebdca9bf1fc82 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 @@ -324,6 +324,39 @@ 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", IntegerType), + 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. + 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(1, 42.0f), Row(1, 44.0f))) + } + } } @ExtendedSQLTest @@ -2600,9 +2633,10 @@ 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. Without - // allowKeysSubsetOfPartitionKeys the pruned `data` key drops the reported partitioning and both - // sides shuffle; with it, the partitioning is kept and projected onto `id`, so SPJ triggers. + // 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")) @@ -2668,6 +2702,99 @@ class KeyGroupedPartitioningSuite } } + 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, @@ -2789,41 +2916,6 @@ class KeyGroupedPartitioningSuite } } - test("SPARK-59248: a pruned source-reported ordering must not defeat exchange reuse") { - val table1 = "ord_reuse" - // Partitioned by `id` (kept in the output) and the source also reports an ordering on `data`. - // The self-join prunes `data`, leaving only the reported ordering dangling; that dangling - // ordering must not keep the two identical legs from being reused. - createTable(table1, columns, Array(identity("id")), - Array(sort(column("data"), SortDirection.ASCENDING))) - 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))") - - 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 - - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - val df = sql(query) - checkAnswer(df, Seq(Row(1, 1), Row(2, 2), Row(3, 3))) - val plan = df.queryExecution.executedPlan - // The scan still reports the (dangling) ordering; the point is that it does not block reuse. - collectScans(plan).foreach { scan => - assert(scan.ordering.exists(_.exists(_.child.references.exists(_.name == "data"))), - s"expected the scan to report an ordering on the pruned `data`:\n$plan") - } - val scans = collectScans(plan) - val reused = collect(plan) { case r: ReusedExchangeExec => r } - assert(scans.length == 1, s"the two identical legs should reuse a single scan:\n$plan") - assert(reused.length == 1, s"expected one reused exchange:\n$plan") - } - } - test("SPARK-47094: SPJ: Support compatible buckets") { val table1 = "tab1e1" val table2 = "table2" From b47eb9d488106248645819d50d5e825940cdb85a Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Mon, 7 Sep 2026 15:31:16 +0800 Subject: [PATCH 3/3] [SPARK-59248][SQL] Make the runtime-filtering test reach the pruned-key path The test declared purchases.item_id narrower than items.id, so the join cast it, translateRuntimeFilterV2 could not translate the cast, and no runtime filter reached the scan. The re-planning path never ran, and the test stayed green with filteredPartitions fed the projected partitioning. Match the two types, and record why the type has to match and why only the V2 fixture of the shared trait prunes a column at all. Assisted-by: Claude Opus 5 --- .../sql/connector/KeyGroupedPartitioningSuite.scala | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) 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 ebdca9bf1fc82..88a2420a3ebc9 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 @@ -341,7 +341,7 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui val pCols = Array( Column.create("store_id", IntegerType), - Column.create("item_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 " + @@ -351,10 +351,18 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui // 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(1, 42.0f), Row(1, 44.0f))) + checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f))) } } }