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..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,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 2e3394ac2a082..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 @@ -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 cde8b1041a952..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,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 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 165425aea19dc..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 @@ -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 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 02fc8afb7d9de..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 @@ -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 @@ -324,6 +324,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 @@ -2597,6 +2638,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 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