diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index b35ad05091c7a..ae946932b1540 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -7293,6 +7293,32 @@ object SQLConf { .booleanConf .createWithDefault(false) + val MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION = buildConf( + "spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.allowKeyGroupedPartitioningDegradation") + .doc("When false, a DataSource V2 scan merge is declined if the rebuilt merged scan would " + + "report weaker key-grouped partitioning than an input reported (no longer clustering by " + + "the same expressions), which can force a shuffle the original plan avoided. When true, " + + "the merge proceeds anyway, trading the partitioning for a single scan. Reported " + + "partitioning is re-derived on the merged scan, so a merge that does not weaken it is " + + "always allowed. Only affects sources that declare the SCAN_MERGING table capability.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + + val MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION = buildConf( + "spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.allowOrderingDegradation") + .doc("When false, a DataSource V2 scan merge is declined if the rebuilt merged scan would " + + "report a weaker output ordering than an input reported (an input ordering is no longer a " + + "prefix of the merged scan's), which can force a sort the original plan avoided. When " + + "true, the merge proceeds anyway, trading the ordering for a single scan. Reported " + + "ordering is re-derived on the merged scan, so a merge that does not weaken it is always " + + "allowed. Only affects sources that declare the SCAN_MERGING table capability.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED = buildConf("spark.sql.optimizer.mergeSubplans.filterPropagation.throughJoin.enabled") .doc("When set to true, filter attributes can propagate through Join nodes during subplan " + diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala index 37de3b560b5c8..5a1e0e79bf0b6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala @@ -56,8 +56,9 @@ class InMemoryScanMergingPartitionFilterCatalog * An [[InMemoryEnhancedPartitionFilterTable]] that returns the `TableCapability.SCAN_MERGING` * capability, so [[org.apache.spark.sql.execution.planmerging.PlanMerger]] may fuse two scans of * this table. Its scan is wrapped in a thin [[NonReportingScan]] so a partitioned table does not - * set the scan relation's `keyGroupedPartitioning` (whose preservation across a merge is a separate - * follow-up); this keeps the fixture focused on the iterative-pushdown behavior under test. + * set the scan relation's `keyGroupedPartitioning`, keeping this fixture focused on the + * iterative-pushdown behavior under test. Preserving a reported partitioning across a merge is + * exercised separately by [[InMemoryScanMergingReportingTable]]. */ class InMemoryScanMergingPartitionFilterTable( name: String, @@ -82,11 +83,57 @@ class InMemoryScanMergingPartitionFilterTable( * Thin scan decorator that exposes only `readSchema`, `toBatch` and `description`, dropping the * base scan's `SupportsReportPartitioning`/`SupportsReportStatistics`. So the scan relation carries * no reported partitioning/ordering/statistics -- for a partitioned table this keeps - * `keyGroupedPartitioning` unset, which the scan merge requires (preserving reported partitioning - * across a merge is a separate follow-up). + * `keyGroupedPartitioning` unset, so the fixture stays focused on pushdown; preserving reported + * partitioning across a merge is exercised by [[InMemoryScanMergingReportingTable]]. */ case class NonReportingScan(inner: Scan) extends Scan { override def readSchema(): StructType = inner.readSchema() override def toBatch: Batch = inner.toBatch override def description(): String = inner.description() } + +/** + * Like [[InMemoryScanMergingPartitionFilterCatalog]] but hands out tables that KEEP their reported + * partitioning/ordering (no [[NonReportingScan]] wrapper), so a scan merge that must preserve the + * reported key-grouped partitioning across the merge can be exercised. + */ +class InMemoryScanMergingReportingCatalog + extends InMemoryTableEnhancedPartitionFilterCatalog { + import CatalogV2Implicits._ + + override def createTable( + ident: Identifier, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = { + if (tables.containsKey(ident)) { + throw new TableAlreadyExistsException(ident.asMultipartIdentifier) + } + InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) + val tableName = s"$name.${ident.quoted}" + val table = new InMemoryScanMergingReportingTable(tableName, columns, partitions, properties) + tables.put(ident, table) + namespaces.putIfAbsent(ident.namespace.toList, Map()) + table + } +} + +/** + * An [[InMemoryEnhancedPartitionFilterTable]] that opts into `TableCapability.SCAN_MERGING` but, + * unlike [[InMemoryScanMergingPartitionFilterTable]], does NOT strip the reported partitioning: a + * partitioned table's scan reports `KeyGroupedPartitioning` as usual. Used to check that a merge + * preserves that report on the rebuilt merged scan (re-derived by V2ScanPartitioningAndOrdering). + */ +class InMemoryScanMergingReportingTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String]) + extends InMemoryEnhancedPartitionFilterTable(name, columns, partitioning, properties) { + + override def capabilities(): util.Set[TableCapability] = { + val caps = new util.HashSet[TableCapability](super.capabilities()) + caps.add(TableCapability.SCAN_MERGING) + caps + } +} 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 df74bbe0036df..7b8117587e8db 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 @@ -19,13 +19,13 @@ package org.apache.spark.sql.execution.planmerging import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, AttributeSet, Expression, ExpressionSet, If, Literal, NamedExpression, Or} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, AttributeSet, Expression, ExpressionSet, If, Literal, NamedExpression, Or, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.{Cross, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project} import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.connector.catalog.TableCapability -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanRelationPushDown} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown} import org.apache.spark.sql.internal.SQLConf /** @@ -139,7 +139,11 @@ class PlanMerger( filterPropagationThroughJoinEnabled: Boolean = SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED), dsv2SymmetricFilterPropagationEnabled: Boolean = - SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED)) { + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED), + dsv2AllowKeyGroupedPartitioningDegradation: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION), + dsv2AllowOrderingDegradation: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION)) { val cache = mutable.ArrayBuffer.empty[MergedPlan] /** @@ -251,10 +255,18 @@ class PlanMerger( * * @param unionAttrs The union of both sides' projected columns the merged scan must produce. * @param strictFilters The strict pushed filters that must be re-enforced by the rebuilt scan. + * @param requiredKeyGroupedPartitioning The key-grouped partitioning the merged scan must + * reproduce to keep both inputs not-worse (the inputs' combined report, in the merged + * relation's attribute space); empty means no requirement. Enforced unless + * `allowKeyGroupedPartitioningDegradation` is set. + * @param requiredOrdering The output ordering the merged scan must satisfy likewise; empty means + * no requirement. Enforced unless `allowOrderingDegradation` is set. */ case class DSv2DeferredScan( unionAttrs: Seq[Attribute], - strictFilters: Seq[Expression]) + strictFilters: Seq[Expression], + requiredKeyGroupedPartitioning: Seq[Expression], + requiredOrdering: Seq[SortOrder]) /** * Context threaded DOWN through [[tryMergePlans]] recursion. @@ -665,16 +677,11 @@ class PlanMerger( // pushdown (aggregate, join, variant, limit, offset, top-N, sample) or built by any other // rule is not mergeable by default. np.mergeableScan && cp.mergeableScan && - // Reported partitioning/ordering (e.g. storage-partitioned join, reported sort) is not - // reconstructed by the rebuilt scan -- it never carries reported partitioning or ordering, - // as V2ScanPartitioningAndOrdering is a separate early rule that rebuildScan does not run. - // So decline the merge when either input reports a NON-EMPTY one rather than silently - // dropping it; merging these can be added as a follow-up. A source that implements - // SupportsReportOrdering/SupportsReportPartitioning but reports nothing yields Some(Nil), - // so test the inner Seq (forall) rather than the Option (isEmpty), which would decline it. - np.keyGroupedPartitioning.forall(_.isEmpty) && - cp.keyGroupedPartitioning.forall(_.isEmpty) && - np.ordering.forall(_.isEmpty) && cp.ordering.forall(_.isEmpty) && + // Reported partitioning/ordering is not reconstructed by the rebuilt scan + // (V2ScanPartitioningAndOrdering is a separate early rule that rebuildScan does not run). + // Rather than decline here, the merged scan re-derives its own when built (see + // tryBuildMergedDSv2Scan), and a not-worse check there declines the merge if that would + // degrade what an input reported (unless a dsv2ScanMerge config allows it). // The table opts in to Spark-side merging (a table capability, so a V1-fallback source // whose scan Spark wraps can still opt in). Both relations are the same table (canonically // equal, checked above), but check each to be safe. @@ -700,17 +707,42 @@ class PlanMerger( // order (npMapping.values would be exprId-hash-ordered). val unionAttrs = cp.output ++ np.output.map(npMapping).filterNot(cp.outputSet.contains) + // The reported key-grouped partitioning / ordering the merged scan must preserve so BOTH inputs + // stay not-worse. Each input reports its own, remapped into cp's relation space (cp's already + // is; np's via npRelationMapping). The two usually agree (same table) but need not -- differing + // best-effort filters can prune different files, and a source may report per file set. Combine + // them into the single report the merge must keep (kGP: they must be equal; ordering: the + // stronger, which satisfies both). None from combine* means the inputs are INCOMPATIBLE -- no + // rebuilt scan could keep both not-worse -- so decline HERE, before rebuilding, unless the + // matching config accepts degrading that dimension. + val combinedKeyGroupedPartitioning = combineRequiredKeyGroupedPartitioning( + np.keyGroupedPartitioning.map(_.map(mapAttributes(_, npRelationMapping))).getOrElse(Nil), + cp.keyGroupedPartitioning.getOrElse(Nil)) + val combinedOrdering = combineRequiredOrdering( + np.ordering.map(_.map(mapAttributes(_, npRelationMapping))).getOrElse(Nil), + cp.ordering.getOrElse(Nil)) + if ((combinedKeyGroupedPartitioning.isEmpty && !dsv2AllowKeyGroupedPartitioningDegradation) || + (combinedOrdering.isEmpty && !dsv2AllowOrderingDegradation)) { + return None + } + // Empty = no requirement (both inputs reported none, or they were incompatible but the config + // accepts the degradation). Otherwise the single report the merged scan must reproduce/satisfy. + val expectedKeyGroupedPartitioning = combinedKeyGroupedPartitioning.getOrElse(Nil) + val expectedOrdering = combinedOrdering.getOrElse(Nil) + if (context.filterAboveScan) { // Defer the build to the enclosing Filter so the scan is built once with strict + // best-effort filters. The placeholder mergedPlan is the bare relation (its output is a // superset of unionAttrs). rebuildScan reuses the relation's attributes, so mapping // np.output to cp's relation attributes is consistent with the eventual built scan. Some(TryMergeResult(cp.relation, npMapping, - dsv2DeferredScan = Some(DSv2DeferredScan(unionAttrs, cp.pushedFilters)), dsv2Merged = true)) + dsv2DeferredScan = Some(DSv2DeferredScan(unionAttrs, cp.pushedFilters, + expectedKeyGroupedPartitioning, expectedOrdering)), dsv2Merged = true)) } else { // No enclosing Filter: build the merged scan here enforcing the (equal) strict filters over // the union of columns, with no best-effort filter (no post-scan Filter to prune on). - tryBuildMergedDSv2Scan(cp.relation, unionAttrs, cp.pushedFilters, bestEffortFilter = None) + tryBuildMergedDSv2Scan(cp.relation, unionAttrs, cp.pushedFilters, None, + expectedKeyGroupedPartitioning, expectedOrdering) .map(TryMergeResult(_, npMapping, dsv2Merged = true)) } } @@ -733,7 +765,9 @@ class PlanMerger( relation: DataSourceV2Relation, unionAttrs: Seq[Attribute], strictFilters: Seq[Expression], - bestEffortFilter: Option[Expression]): Option[DataSourceV2ScanRelation] = { + bestEffortFilter: Option[Expression], + requiredKeyGroupedPartitioning: Seq[Expression], + requiredOrdering: Seq[SortOrder]): Option[DataSourceV2ScanRelation] = { val relationOut = relation.outputSet // Defensive: strict filters come from `pushedFilters`, which reference only relation columns, // so this holds today. If a future caller offers a filter over non-relation attributes, decline @@ -760,6 +794,16 @@ class PlanMerger( // re-checks it), and the scan must produce exactly the requested union of columns. strictFilters.forall(ExpressionSet(scan.pushedFilters).contains) && scan.outputSet == AttributeSet(unionAttrs) + }.map { scan => + // rebuildScan returns the merged scan with reported partitioning/ordering unset + // (V2ScanPartitioningAndOrdering is a separate early rule the rebuild does not run), so + // re-derive them on this single node. Safe on one node: the partitioning pass is idempotent + // and the ordering pass is applied once to a fresh node. + V2ScanPartitioningAndOrdering(scan).asInstanceOf[DataSourceV2ScanRelation] + }.filterNot { merged => + // Decline if the merged scan degrades a partitioning/ordering an input reported -- that can + // force a shuffle/sort the original plan avoided -- unless the matching config opts in. + mergeDegradesReporting(merged, requiredKeyGroupedPartitioning, requiredOrdering) } } @@ -784,14 +828,56 @@ class PlanMerger( // turned every other relation into a DataSourceV2ScanRelation), so recover it by type here // rather than carrying it on DSv2DeferredScan. child.collectFirst { case r: DataSourceV2Relation => r }.flatMap { relation => - tryBuildMergedDSv2Scan(relation, d.unionAttrs, d.strictFilters, bestEffortFilter) - .orElse(tryBuildMergedDSv2Scan(relation, d.unionAttrs, d.strictFilters, None)) + tryBuildMergedDSv2Scan(relation, d.unionAttrs, d.strictFilters, bestEffortFilter, + d.requiredKeyGroupedPartitioning, d.requiredOrdering) + .orElse(tryBuildMergedDSv2Scan(relation, d.unionAttrs, d.strictFilters, None, + d.requiredKeyGroupedPartitioning, d.requiredOrdering)) .map { built => child.transformUp { case r: DataSourceV2Relation if r eq relation => built } } } } + // The key-grouped partitioning the merged scan must reproduce to keep both inputs not-worse: they + // must be equal (bucketing is a table property), so a differing non-empty pair is INCOMPATIBLE + // (None); an empty side imposes no constraint. Compared canonically in cp's relation space (np's + // report was remapped into it by the caller). + private def combineRequiredKeyGroupedPartitioning( + a: Seq[Expression], b: Seq[Expression]): Option[Seq[Expression]] = { + if (a.isEmpty) Some(b) + else if (b.isEmpty) Some(a) + else if (a.map(_.canonicalized) == b.map(_.canonicalized)) Some(a) + else None + } + + // The ordering the merged scan must satisfy to keep both inputs not-worse: the stronger of the + // two (the one that satisfies the other -- satisfying it implies satisfying the weaker). If + // neither satisfies the other they are INCOMPATIBLE (None). An empty ordering never constrains. + private def combineRequiredOrdering( + a: Seq[SortOrder], b: Seq[SortOrder]): Option[Seq[SortOrder]] = { + if (SortOrder.orderingSatisfies(a, b)) Some(a) + else if (SortOrder.orderingSatisfies(b, a)) Some(b) + else None + } + + // True when the rebuilt merged scan does not reproduce the required key-grouped partitioning, or + // does not satisfy the required ordering (the combined report the merge must preserve, computed + // at the leaf). Gated per dimension by the dsv2ScanMerge degradation configs; an empty required + // report imposes no constraint. Compared in cp's relation space. + private def mergeDegradesReporting( + merged: DataSourceV2ScanRelation, + requiredKeyGroupedPartitioning: Seq[Expression], + requiredOrdering: Seq[SortOrder]): Boolean = { + val kgpDegraded = !dsv2AllowKeyGroupedPartitioningDegradation && + requiredKeyGroupedPartitioning.nonEmpty && + !merged.keyGroupedPartitioning.exists( + _.map(_.canonicalized) == requiredKeyGroupedPartitioning.map(_.canonicalized)) + val orderingDegraded = !dsv2AllowOrderingDegradation && + requiredOrdering.nonEmpty && + !SortOrder.orderingSatisfies(merged.ordering.getOrElse(Nil), requiredOrdering) + kgpDegraded || orderingDegraded + } + // Returns true when a filter attribute originating from `fromLeft` child of a join with // `joinType` can be safely propagated through that join to a parent Aggregate. // diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala index ad40f8c9e1a9f..4663b5e3cc707 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala @@ -21,7 +21,7 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.sql.{DataFrame, QueryTest, Row} import org.apache.spark.sql.connector.FakeV2ProviderWithCustomSchema -import org.apache.spark.sql.connector.catalog.InMemoryScanMergingPartitionFilterCatalog +import org.apache.spark.sql.connector.catalog.{InMemoryScanMergingPartitionFilterCatalog, InMemoryScanMergingReportingCatalog} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -44,11 +44,14 @@ class DSv2PlanMergingSuite extends QueryTest with SharedSparkSession before { spark.conf.set("spark.sql.catalog.scanmerge", classOf[InMemoryScanMergingPartitionFilterCatalog].getName) + spark.conf.set("spark.sql.catalog.scanmergereport", + classOf[InMemoryScanMergingReportingCatalog].getName) } after { spark.sessionState.catalogManager.reset() spark.conf.unset("spark.sql.catalog.scanmerge") + spark.conf.unset("spark.sql.catalog.scanmergereport") } private def v2Scans(df: DataFrame): Seq[DataSourceV2ScanRelation] = @@ -169,4 +172,39 @@ class DSv2PlanMergingSuite extends QueryTest with SharedSparkSession } } } + + test("SPARK-58549: a scan merge preserves the sources' reported key-grouped partitioning") { + val t = "scanmergereport.t2" + withTable(t) { + withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "true") { + sql(s"CREATE TABLE $t (c1 int, c2 int) USING $v2Source PARTITIONED BY (c1)") + sql(s"INSERT INTO $t VALUES (1, 10), (2, 20), (3, 30)") + + // Both scalar subqueries read the partition column c1, so each scan reports + // KeyGroupedPartitioning on c1; they differ in the extra column read, so PlanMerger fuses + // them into one scan reading {c1, c2}. c1 survives in the union, so the merged scan + // re-derives the same partitioning -- not a degradation, so the merge proceeds by default. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $t) AS m1, + | (SELECT max(c1 + c2) FROM $t) AS m2 + |""".stripMargin) + checkAnswer(df, Row(3, 33)) + + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 1, + s"the two scans should be fused into one:\n${df.queryExecution.optimizedPlan}") + val scan = scans.head + assert(scan.output.map(_.name).toSet == Set("c1", "c2"), + s"the merged scan should read the union of both columns; got ${scan.output}") + assert(scan.keyGroupedPartitioning.exists(_.nonEmpty), + s"the merged scan should preserve the reported key-grouped partitioning; " + + s"got ${scan.keyGroupedPartitioning}") + assert(scan.keyGroupedPartitioning.get.flatMap(_.references).exists(_.name == "c1"), + s"the preserved partitioning should be on c1; got ${scan.keyGroupedPartitioning}") + assertNoPlaceholderRelation(df) + } + } + } } 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 5e4b7fef314b2..59b6c72decd94 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 @@ -2600,10 +2600,11 @@ class MergeSubplansSuite extends PlanTest { } test("SPARK-40259: do not merge DSv2 scans that report key-grouped partitioning or ordering") { - // The rebuilt merged scan does not reconstruct reported partitioning/ordering, so a scan - // reporting either declines the merge (checked on both the np and cp side) -- the plan is left - // unchanged -- rather than silently dropping it. Preserving them across a merge is a deferred - // follow-up. (The plain-scan merge is already covered by the projected-columns test above.) + // An input reports key-grouped partitioning or ordering, but the merged scan -- rebuilt over a + // TestV2Scan that reports neither -- re-derives nothing, so merging would degrade what the + // input reported. With the degradation configs off (the default) the merge is declined on both + // the np and cp side and the plan is left unchanged, rather than forcing a shuffle/sort the + // original plan avoided. def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = Seq( (withField(v2ScanReading("a")), v2ScanReading("b")), @@ -2618,6 +2619,109 @@ class MergeSubplansSuite extends PlanTest { assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) } + test("SPARK-58549: do not merge DSv2 scans reporting incompatible kGP/ordering") { + // Both inputs report a partitioning/ordering, but on the different column each reads, so no + // single rebuilt scan could keep both not-worse. combineRequired* returns None (incompatible), + // so the merge is declined at the leaf -- before any rebuild -- with the degradation configs + // off (the default). This is distinct from the single-side case above (which rebuilds and then + // finds the re-derived report degraded): here the two inputs disagree with each other up front. + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = { + // withField applied to the "a" scan reports on a, applied to the "b" scan reports on b. + val q = testRelation.select( + ScalarSubquery(withField(v2ScanReading("a")).groupBy()(sum($"a").as("sa"))), + ScalarSubquery(withField(v2ScanReading("b")).groupBy()(sum($"b").as("sb")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + } + + test("SPARK-58549: merge DSv2 scans reporting kGP/ordering when the degradation config allows") { + // With the matching degradation config on, a merge that would drop a reported partitioning or + // ordering proceeds anyway (trading it for a single scan). The merged scan re-derives no report + // from the non-reporting TestV2Scan, so the fused plan is the plain column union. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val mergedSubquery = mergedScan + .groupBy()(sum($"a").as("sum_a"), sum($"b").as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + def assertMerges( + withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation, confKey: String): Unit = { + val sub1 = ScalarSubquery(withField(v2ScanReading("a")).groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(v2ScanReading("b").groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + withSQLConf(confKey -> "true") { + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + } + + assertMerges(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key) + assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key) + } + + 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 + // DEFERRED to the enclosing Filter, and the required report is carried there through + // DSv2DeferredScan. This test drives that deferred path: each scan reads {a, } and reports + // on a, both filter on `a > 1` (so the two fuse into {a, b, c} under one Filter). The merged + // scan is rebuilt over a non-reporting TestV2Scan, so it re-derives no report -- a degradation. + def reportsOnA(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation) + : (ScalarSubquery, ScalarSubquery) = ( + ScalarSubquery( + withField(v2ScanReading("a", "b")).where($"a" > 1).groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery( + withField(v2ScanReading("a", "c")).where($"a" > 1).groupBy()(sum($"c").as("sum_c")))) + + // Default configs: the deferred build declines on the degradation, leaving the plan unchanged. + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = { + val (sub1, sub2) = reportsOnA(withField) + val q = testRelation.select(sub1, sub2) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + + // With the matching config on, the deferred build proceeds: the two scans fuse into {a, b, c} + // with the identical `a > 1` re-pushed for pruning (as in the identical-filter merge test). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b", "c")) + val mergedSubquery = mergedScan.where($"a" > 1) + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("sum_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("sum_c"), $"sum_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + def assertMerges( + withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation, confKey: String): Unit = { + val (sub1, sub2) = reportsOnA(withField) + val q = testRelation.select(sub1, sub2) + withSQLConf(confKey -> "true") { + comparePlans(Optimize.execute(q.analyze), correctAnswer.analyze) + } + } + assertMerges(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key) + assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key) + } + test("SPARK-40259: merge DSv2 scans that report empty key-grouped partitioning or ordering") { // A source implementing SupportsReportPartitioning/SupportsReportOrdering but reporting nothing // yields Some(Nil), not None (V2ScanPartitioningAndOrdering sets the field unconditionally). An