This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 8645e2c08929 [SPARK-58549][SQL] Preserve key-grouped partitioning and
ordering across a DSv2 scan merge
8645e2c08929 is described below
commit 8645e2c08929bbc1ef2ef66b442f525814f92220
Author: Peter Toth <[email protected]>
AuthorDate: Sun Aug 9 11:53:52 2026 +0200
[SPARK-58549][SQL] Preserve key-grouped partitioning and ordering across a
DSv2 scan merge
### What changes were proposed in this pull request?
Follow-up to SPARK-40259 (subquery plan merge for DataSource V2 scans).
Today the DSv2 scan merge declines whenever either input scan reports
key-grouped partitioning or ordering. The rebuilt merged scan carries neither
on its own, so fusing the two scans could drop a partitioning/ordering the
original plan relied on and force an extra shuffle or sort. Declining is safe
but leaves the merge unused for any partitioned/ordered source.
This PR lets the merge proceed and re-derives the merged scan's own report
instead of declining up front:
- Drop the `keyGroupedPartitioning`/`ordering` conjuncts from the
`mergeable` gate in `PlanMerger.tryMergeScanRelations`.
- At the leaf, combine the two inputs' reports (np's remapped into cp's
relation space) into the single report the merge must preserve:
- key-grouped partitioning: the two must be equal
(`combineRequiredKeyGroupedPartitioning`);
- ordering: the stronger of the two, i.e. the one that satisfies the
other (`combineRequiredOrdering`).
Both compare reported expressions by transform semantics
(`TransformExpression.isSameFunction`, the way a storage-partitioned join
compares its two sides) rather than by canonical equality. `V2ExpressionUtils`
binds a transform function afresh on every derivation, and `canonicalName` is a
documented obligation (its default returns a random UUID to force an override)
while `equals` on `BoundFunction` is required nowhere -- so for a connector
meeting only the documented contract, cano [...]
If the inputs are incompatible (differing non-empty kGP, or neither
ordering satisfies the other), no rebuilt scan can keep both not-worse, so the
merge is declined right there -- before rebuilding -- saving a scan rebuild.
- After `V2ScanRelationPushDown.rebuildScan`, re-derive the merged scan's
partitioning/ordering by running `V2ScanPartitioningAndOrdering` on the single
merged scan node, then check it against the combined required report
(`mergeDegradesReporting`): decline if the merged scan's kGP does not match, or
its ordering does not satisfy, what was required. The check runs at the build's
callers, and per build attempt, so a report degradation stays distinguishable
from a strict filter that can [...]
Two new opt-in configs gate accepting a degradation instead of declining
(both default `false`), so with the defaults this is a pure improvement --
merge when not worse, decline on loss:
-
`spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.keyGroupedPartitioningDegradation.enabled`
-
`spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.orderingDegradation.enabled`
Only sources that declare the `SCAN_MERGING` table capability are affected.
### Why are the changes needed?
The parent feature (SPARK-40259) conservatively declines a merge whenever
an input reports key-grouped partitioning or ordering, to avoid forcing a
shuffle/sort that the original plan avoided. That is correct but overly broad:
when the rebuilt merged scan re-derives the same (or a stronger) report, fusing
the scans loses nothing and still removes a duplicate scan -- which matters for
partitioned/ordered sources, e.g. an SPJ join inside the merged subplan. This
PR keeps the safety (dec [...]
### Does this PR introduce _any_ user-facing change?
No. The two new configs are opt-in and default to `false`, and no built-in
source declares `SCAN_MERGING`, so there is no behavior change for existing
sources.
### How was this patch tested?
New and existing unit tests, all under `sql/core`:
- `MergeSubplansSuite`:
- default-decline when a single input reports kGP/ordering the rebuilt
scan drops (degradation);
- default-decline when the two inputs report incompatible kGP/ordering
(declined at the leaf, before rebuild);
- config-allows-degradation: the merge proceeds to the plain column union
under each `...Degradation.enabled=true`, both for a single reporting input
(the degradation is found after the rebuild) and for two incompatible ones
(declined before it);
- two independently derived `bucket(4, a)` reports merge and the rebuilt
scan keeps the report -- the only test that reaches the transform comparison.
It uses a `FunctionCatalog` fixture that binds a fresh `BoundFunction` per
call, the way a real connector does; Spark's own `UnboundBucketFunction`
returns a singleton, so a fixture built on that would pass with or without the
fix;
- the merge proceeds when the rebuilt scan re-derives the required
ordering, and declines when it re-derives less than the two inputs' combined
(stronger) ordering;
- the deferred (under-`Filter`) build path enforces the required report
too: decline by default, merge with the config on;
- the deferred build checks the report on EACH attempt, so a source that
reports only over the unpruned file set still merges on the strict-only attempt
-- which is what the leaf builds when no `Filter` is above the scan, so
checking only the first attempt would leave the deferred path weaker than the
leaf path;
- the pre-existing empty-report case still merges.
- `DSv2PlanMergingSuite`: an end-to-end test that a scan merge preserves
the sources' reported key-grouped partitioning, using a new `SCAN_MERGING`
fixture
(`InMemoryScanMergingReportingCatalog`/`InMemoryScanMergingReportingTable`)
that keeps its reported partitioning (no `NonReportingScan` wrapper), an
identity-partitioned table under `spark.sql.sources.v2.bucketing.enabled=true`,
and two scalar subqueries reading the partition column. It also asserts that
the preserved partitioning [...]
`build/sbt 'sql/testOnly *MergeSubplansSuite *DSv2PlanMergingSuite
*PlanMergingSuite'` -- 99 tests pass. `dev/lint-scala` clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
Closes #57753 from
peter-toth/SPARK-58549-preserve-kgp-ordering-dsv2-scan-merge.
Authored-by: Peter Toth <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 30 ++
.../InMemoryScanMergingPartitionFilterTable.scala | 59 +++-
.../sql/execution/planmerging/PlanMerger.scala | 209 ++++++++++--
.../planmerging/DSv2PlanMergingSuite.scala | 65 +++-
.../execution/planmerging/MergeSubplansSuite.scala | 379 ++++++++++++++++++++-
5 files changed, 686 insertions(+), 56 deletions(-)
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 4f1643d4725b..435aeb6e085f 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
@@ -7477,6 +7477,36 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION =
buildConf(
+
"spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.keyGroupedPartitioningDegradation.enabled")
+ .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. The merge is " +
+ "also declined, before any rebuild, when the two input scans report
different clustering " +
+ "expressions, since no single merged scan could preserve both. When
true, the merge " +
+ "proceeds anyway in both cases, 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.orderingDegradation.enabled")
+ .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. The merge " +
+ "is also declined, before any rebuild, when neither input's reported
ordering satisfies " +
+ "the other, since no single merged scan could preserve both. When true,
the merge proceeds " +
+ "anyway in both cases, 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 37de3b560b5c..6fd043c36464 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
@@ -34,6 +34,17 @@ class InMemoryScanMergingPartitionFilterCatalog
extends InMemoryTableEnhancedPartitionFilterCatalog {
import CatalogV2Implicits._
+ /**
+ * The table this catalog hands out. Overridden by
[[InMemoryScanMergingReportingCatalog]], which
+ * shares the `createTable` body below and differs only in the table it
creates.
+ */
+ protected def newScanMergingTable(
+ tableName: String,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ properties: util.Map[String, String]): Table =
+ new InMemoryScanMergingPartitionFilterTable(tableName, columns,
partitions, properties)
+
override def createTable(
ident: Identifier,
columns: Array[Column],
@@ -44,8 +55,7 @@ class InMemoryScanMergingPartitionFilterCatalog
}
InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
val tableName = s"$name.${ident.quoted}"
- val table =
- new InMemoryScanMergingPartitionFilterTable(tableName, columns,
partitions, properties)
+ val table = newScanMergingTable(tableName, columns, partitions, properties)
tables.put(ident, table)
namespaces.putIfAbsent(ident.namespace.toList, Map())
table
@@ -53,13 +63,28 @@ 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.
+ * 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 InMemoryScanMergingPartitionFilterTable(
+class InMemoryScanMergingReportingCatalog extends
InMemoryScanMergingPartitionFilterCatalog {
+ override protected def newScanMergingTable(
+ tableName: String,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ properties: util.Map[String, String]): Table =
+ new InMemoryScanMergingReportingTable(tableName, columns, partitions,
properties)
+}
+
+/**
+ * An [[InMemoryEnhancedPartitionFilterTable]] that opts into
`TableCapability.SCAN_MERGING`, so
+ * [[org.apache.spark.sql.execution.planmerging.PlanMerger]] may fuse two
scans of it. It does NOT
+ * strip the reported partitioning: a partitioned table's scan reports
`KeyGroupedPartitioning` as
+ * usual, so this is the fixture for checking that a merge preserves that
report on the rebuilt
+ * merged scan (re-derived by V2ScanPartitioningAndOrdering). It is also the
base of
+ * [[InMemoryScanMergingPartitionFilterTable]], which drops the report again.
+ */
+class InMemoryScanMergingReportingTable(
name: String,
columns: Array[Column],
partitioning: Array[Transform],
@@ -71,6 +96,20 @@ class InMemoryScanMergingPartitionFilterTable(
caps.add(TableCapability.SCAN_MERGING)
caps
}
+}
+
+/**
+ * An [[InMemoryScanMergingReportingTable]] whose scan is wrapped in a thin
[[NonReportingScan]], so
+ * a partitioned table does not 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 by the base
[[InMemoryScanMergingReportingTable]] instead.
+ */
+class InMemoryScanMergingPartitionFilterTable(
+ name: String,
+ columns: Array[Column],
+ partitioning: Array[Transform],
+ properties: util.Map[String, String])
+ extends InMemoryScanMergingReportingTable(name, columns, partitioning,
properties) {
override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
new InMemoryEnhancedPartitionFilterScanBuilder(schema()) {
@@ -82,8 +121,8 @@ 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()
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 df74bbe0036d..b614985bed9e 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, TransformExpression}
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
+ * `dsv2AllowKeyGroupedPartitioningDegradation` is set.
+ * @param requiredOrdering The output ordering the merged scan must satisfy
likewise; empty means
+ * no requirement. Enforced unless `dsv2AllowOrderingDegradation` 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.
@@ -278,8 +290,10 @@ class PlanMerger(
* - The `merge()` pattern requiring `dsv2DeferredScan = None` is the
fail-safe backstop: if a
* deferred scan somehow reached `merge()` unbuilt, it declines the merge
rather than emitting a
* plan with a placeholder relation.
- * - Eligibility gating stays at the leaf (read-only, inspects only the two
input scans). Only the
- * build moves up.
+ * - Eligibility gating stays at the leaf (read-only, inspects only the two
input scans) with one
+ * exception: whether the rebuilt merged scan degrades a
partitioning/ordering an input reported
+ * can only be decided once that scan exists, so on the deferred path that
one check runs at the
+ * Filter, next to the build (see `tryBuildFilterDSv2ScanChild`).
*/
case class MergeContext(filterPropagationSupported: Boolean,
filterAboveScan: Boolean = false)
@@ -407,9 +421,11 @@ class PlanMerger(
if (mappedNPCondition.canonicalized ==
cp.condition.canonicalized) {
// Identical conditions: the filter node adds no new
discrimination between the two
// sides, so keep it unchanged. If it sits above a deferred
merged DSv2 scan, build
- // that scan once here with this condition as a best-effort
filter. If the build
- // cannot re-enforce the strict filters, decline the whole
merge (the leaf's
- // strict-only build would have failed identically).
+ // that scan once here with this condition as a best-effort
filter. If it cannot be
+ // built to spec -- the strict filters do not come back
re-enforced, or the rebuilt
+ // scan's re-derived partitioning/ordering would degrade what
the inputs reported --
+ // decline the whole merge (the leaf's strict-only build would
have failed
+ // identically).
tryBuildFilterDSv2ScanChild(mergedChild, deferred,
Some(cp.condition))
.map { prunedChild =>
val mergedPlan = Filter(cp.condition, prunedChild)
@@ -617,8 +633,10 @@ class PlanMerger(
* columns. The connector opts in via
* `TableCapability.SCAN_MERGING`; Spark runs the real DSv2 pushdown
* ([[V2ScanRelationPushDown]]) on a synthetic `Filter` over the relation,
extracts the merged
- * scan, and verifies the (equal) strict filters remain fully enforced.
Eligibility gating
- * (`mergeable`) is read-only, inspecting only the two input scans.
+ * scan, and verifies the (equal) strict filters remain fully enforced. The
`mergeable` gate is
+ * read-only, inspecting only the two input scans; two further checks can
decline the merge -- the
+ * inputs' reported partitioning/ordering must combine into a single report
(checked here, before
+ * any rebuild), and the rebuilt scan must not degrade that report (checked
once it is built).
*
* When this scan sits under a [[Filter]] (`context.filterAboveScan`), the
build is DEFERRED: the
* merged scan is not built here but carried up as a [[DSv2DeferredScan]]
and built exactly once
@@ -665,16 +683,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 mergeDegradesReporting 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.
@@ -689,7 +702,8 @@ class PlanMerger(
ExpressionSet(np.pushedFilters.map(mapAttributes(_,
npRelationMapping))) ==
ExpressionSet(cp.pushedFilters)
- // Eligibility is settled above (a read-only gate); everything below
constructs the merge.
+ // The read-only gate above is settled. The report combine below can still
decline (and so can
+ // the degradation check once the scan is built); everything else below
constructs the merge.
if (!mergeable) {
return None
}
@@ -700,17 +714,44 @@ 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 normally agree -- the
clustering expressions come
+ // from the table, and for partitioning, pruning only ever drops a report
wholesale (an empty
+ // side, which never constrains) -- but a source is free to report per
scan, e.g. an ordering
+ // that holds only for the file set the filters pushed into that scan left
behind. Combine them
+ // into the single report the merge must keep (kGP: 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)
+ .filterNot(mergeDegradesReporting(_, expectedKeyGroupedPartitioning,
expectedOrdering))
.map(TryMergeResult(_, npMapping, dsv2Merged = true))
}
}
@@ -728,6 +769,10 @@ class PlanMerger(
* non-deterministic predicate the source prunes on would drop rows the
enclosing Filter cannot
* recover) and references only the relation's own columns (propagated
boolean filter attributes
* are not columns of the relation).
+ *
+ * The rebuilt scan's reported partitioning/ordering is re-derived here, but
NOT checked against
+ * what the inputs reported -- callers do that via
[[mergeDegradesReporting]], so a degradation
+ * stays distinguishable from a filter-enforcement failure.
*/
private def tryBuildMergedDSv2Scan(
relation: DataSourceV2Relation,
@@ -760,6 +805,12 @@ 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]
}
}
@@ -768,10 +819,13 @@ class PlanMerger(
* [[DSv2DeferredScan]], build the scan once here -- at the enclosing
Filter, with the strict
* filters plus the Filter's `bestEffortFilter` -- and splice it in place of
the placeholder
* relation. Tries strict + best-effort first, then strict-only (the
best-effort filter is
- * droppable); returns `None` only if the strict filters cannot be
re-enforced at all, in which
- * case the caller must decline the merge (the leaf's strict-only build
would have failed
- * identically). If there is no deferred scan (a non-DSv2 child, or a scan
not under a Filter),
- * there is nothing to build and the child is returned unchanged.
+ * droppable); a build returns `None` only if the strict filters cannot be
re-enforced at all (the
+ * leaf's strict-only build would have failed identically). Each built scan
is also checked
+ * against the required report the leaf computed, and rejected if it would
degrade that -- per
+ * attempt, so a source whose report depends on which filters were pushed
can still satisfy it
+ * strict-only. Either way the caller must decline the merge. If there is no
deferred scan (a
+ * non-DSv2 child, or a scan not under a Filter), there is nothing to build
and the child is
+ * returned unchanged.
*/
private def tryBuildFilterDSv2ScanChild(
child: LogicalPlan,
@@ -784,14 +838,111 @@ 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))
+ // Check the report per attempt, and at the caller rather than inside
the build: the strict
+ // filters and the report are independent reasons to reject a build,
so a source whose
+ // report depends on what got pushed still gets its second chance from
the strict-only
+ // attempt, and `tryBuildMergedDSv2Scan`'s `None` keeps its single
meaning. The strict-only
+ // attempt is what the leaf builds when no Filter is above the scan,
so checking only the
+ // first attempt would leave the deferred path weaker than the leaf
path.
+ def build(offeredBestEffortFilter: Option[Expression]) =
+ tryBuildMergedDSv2Scan(
+ relation, d.unionAttrs, d.strictFilters, offeredBestEffortFilter)
+ .filterNot(
+ mergeDegradesReporting(_, d.requiredKeyGroupedPartitioning,
d.requiredOrdering))
+
+ build(bestEffortFilter)
+ .orElse(build(None))
.map { built =>
child.transformUp { case r: DataSourceV2Relation if r eq relation
=> built }
}
}
}
+ // Compares two reported partitioning expression lists. A
TransformExpression must be compared by
+ // transform semantics -- `isSameFunction`, i.e. the connector's
`canonicalName` and bucket count
+ // -- and NOT by the BoundFunction instance it holds. V2ExpressionUtils
binds the function afresh
+ // on every derivation, and while `canonicalName` is a documented obligation
(its default returns
+ // a random UUID precisely to force an override), `equals` on BoundFunction
is not required
+ // anywhere. So for a connector that meets only the documented contract,
`canonicalized ==` calls
+ // two identical `bucket(4, id)` reports different and declines the merge.
Connectors that do
+ // define `equals` semantically are unaffected -- Iceberg's
`BaseScalarFunction` compares on
+ // `canonicalName` -- but that is a courtesy we should not depend on.
`isSameFunction` is the same
+ // notion a storage-partitioned join uses to compare its two sides, but only
that part of it is
+ // borrowed: `KeyedShuffleSpec.isExpressionCompatible` matches any two
leaves (it recovers the
+ // positions separately), ignores transform children and honours
+ // `v2BucketingAllowCompatibleTransforms`, none of which is sound here --
preserving a report
+ // means reproducing it exactly, so leaves and children must match too.
+ private def sameReportedExpressions(a: Seq[Expression], b: Seq[Expression]):
Boolean = {
+ a.length == b.length && a.zip(b).forall {
+ case (l: TransformExpression, r: TransformExpression) =>
+ // Recurse, so a nested transform's children are compared by transform
semantics too,
+ // rather than falling back to the instance-sensitive equality this
method exists to avoid.
+ l.isSameFunction(r) && sameReportedExpressions(l.children, r.children)
+ case (l, r) => l semanticEquals r
+ }
+ }
+
+ // [[SortOrder.orderingSatisfies]] with the sort keys compared by
`sameReportedExpressions` (see
+ // above) rather than semanticEquals. `sameOrderExpressions` needs no
handling here: a reported
+ // ordering comes from V2ExpressionUtils.toCatalystOrdering, which always
leaves it empty.
+ private def reportedOrderingSatisfies(o1: Seq[SortOrder], o2:
Seq[SortOrder]): Boolean = {
+ o2.length <= o1.length && o2.zip(o1).forall { case (required, provided) =>
+ provided.direction == required.direction && provided.nullOrdering ==
required.nullOrdering &&
+ sameReportedExpressions(Seq(provided.child), Seq(required.child))
+ }
+ }
+
+ // The key-grouped partitioning the merged scan must reproduce to keep both
inputs not-worse: the
+ // two must be equal, so a differing non-empty pair is INCOMPATIBLE (None);
an empty side imposes
+ // no constraint. Compared 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 (sameReportedExpressions(a, b)) 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 (reportedOrderingSatisfies(a, b)) Some(a)
+ else if (reportedOrderingSatisfies(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) -- that can force a shuffle/sort the original plan avoided.
Gated per dimension by
+ // the dsv2ScanMerge degradation configs; an empty required report imposes
no constraint. Compared
+ // in cp's relation space.
+ //
+ // Only the reported EXPRESSIONS are compared, which is all a
DataSourceV2ScanRelation carries;
+ // 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
+ // 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(
+ merged: DataSourceV2ScanRelation,
+ requiredKeyGroupedPartitioning: Seq[Expression],
+ requiredOrdering: Seq[SortOrder]): Boolean = {
+ val kgpDegraded = !dsv2AllowKeyGroupedPartitioningDegradation &&
+ requiredKeyGroupedPartitioning.nonEmpty &&
+ !merged.keyGroupedPartitioning.exists(
+ sameReportedExpressions(_, requiredKeyGroupedPartitioning))
+ val orderingDegraded = !dsv2AllowOrderingDegradation &&
+ requiredOrdering.nonEmpty &&
+ !reportedOrderingSatisfies(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 ad40f8c9e1a9..977e115d16b0 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
@@ -20,9 +20,10 @@ package org.apache.spark.sql.execution.planmerging
import org.scalatest.BeforeAndAfter
import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
import org.apache.spark.sql.connector.FakeV2ProviderWithCustomSchema
-import
org.apache.spark.sql.connector.catalog.InMemoryScanMergingPartitionFilterCatalog
-import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
DataSourceV2ScanRelation}
+import
org.apache.spark.sql.connector.catalog.{InMemoryScanMergingPartitionFilterCatalog,
InMemoryScanMergingReportingCatalog}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2Relation, DataSourceV2ScanRelation}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -44,11 +45,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 +173,61 @@ 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) {
+ // V2_BUCKETING_ENABLED is what makes the preserved report reach the
physical plan (only
+ // DataSourceV2ScanExecBase.outputPartitioning reads it), which the last
assertion checks; AQE
+ // is off so that `executedPlan` is the plan those assertions can walk.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ 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)
+
+ // The preserved report is not just carried on the logical node: the
merged scan still
+ // reports it as its physical output partitioning, which is what lets
a storage-partitioned
+ // join above it skip the shuffle.
+ val batchScans = df.queryExecution.executedPlan.collectWithSubqueries {
+ case b: BatchScanExec => b
+ }
+ assert(batchScans.nonEmpty, s"expected a
BatchScanExec:\n${df.queryExecution.executedPlan}")
+ batchScans.foreach { b =>
+ b.outputPartitioning match {
+ case k: KeyedPartitioning =>
+ assert(k.expressions.flatMap(_.references).exists(_.name ==
"c1"),
+ s"the merged scan should be key-grouped on c1; got
${k.expressions}")
+ case other =>
+ fail(s"expected a KeyedPartitioning on the merged scan, got
$other")
+ }
+ }
+ }
+ }
+ }
}
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 5e4b7fef314b..bf689c73476f 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
@@ -19,18 +19,20 @@ package org.apache.spark.sql.execution.planmerging
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending,
Attribute, AttributeReference, CreateNamedStruct, ExprId, GetStructField, If,
Literal, Or, ScalarSubquery, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending,
Attribute, AttributeReference, CreateNamedStruct, ExprId, GetStructField, If,
Literal, Or, ScalarSubquery, SortOrder, TransformExpression}
import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules._
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
-import org.apache.spark.sql.connector.catalog.{SupportsRead, Table,
TableCapability}
-import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder}
+import org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier,
SupportsRead, Table, TableCapability}
+import org.apache.spark.sql.connector.catalog.functions.{BoundFunction,
ScalarFunction, UnboundFunction}
+import org.apache.spark.sql.connector.expressions.{Expressions,
FieldReference, SortDirection => V2SortDirection, SortOrder => V2SortOrder}
import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownLimit, SupportsPushDownOffset, SupportsPushDownRequiredColumns,
SupportsPushDownTableSample, SupportsPushDownTopN, SupportsPushDownV2Filters}
-import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
DataSourceV2ScanRelation, V2ScanRelationPushDown}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownLimit, SupportsPushDownOffset, SupportsPushDownRequiredColumns,
SupportsPushDownTableSample, SupportsPushDownTopN, SupportsPushDownV2Filters,
SupportsReportOrdering, SupportsReportPartitioning}
+import
org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning =>
V2KeyGroupedPartitioning, Partitioning => V2Partitioning, UnknownPartitioning
=> V2UnknownPartitioning}
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+import org.apache.spark.sql.types.{DataType, IntegerType, StringType,
StructField, StructType}
import org.apache.spark.sql.util.CaseInsensitiveStringMap
class MergeSubplansSuite extends PlanTest {
@@ -1985,6 +1987,32 @@ class MergeSubplansSuite extends PlanTest {
StructField("a", IntegerType), StructField("b", IntegerType),
StructField("c", StringType))),
supportsScanMerging = false)
+ /**
+ * Same shape as [[v2Table]], but its scans report `a ASC` as their output
ordering, so a scan
+ * rebuilt over this table re-derives a report instead of none.
+ */
+ private val v2TableReportingA = new TestV2Table(StructType(Seq(
+ StructField("a", IntegerType), StructField("b", IntegerType),
StructField("c", StringType))),
+ reportedOrderingCols = Seq("a"))
+
+ /**
+ * Same shape as [[v2Table]], but its scans report `bucket(4, a)` as their
key-grouped
+ * partitioning -- a TRANSFORM report, which resolves to a
`TransformExpression` rather than a
+ * plain attribute (see [[v2ScanReportingOn]]).
+ */
+ private val v2TableBucketedOnA = new TestV2Table(StructType(Seq(
+ StructField("a", IntegerType), StructField("b", IntegerType),
StructField("c", StringType))),
+ reportedBucketPartition = Some((4, "a")))
+
+ /**
+ * Same shape as [[v2Table]], but its scans report `a ASC` ONLY when no
filter was pushed into
+ * them, so the deferred build's two attempts (strict + best-effort, then
strict-only) re-derive
+ * different reports.
+ */
+ private val v2TablePushSensitiveReport = new TestV2Table(StructType(Seq(
+ StructField("a", IntegerType), StructField("b", IntegerType),
StructField("c", StringType))),
+ reportedOrderingCols = Seq("a"), reportOnlyWhenNothingPushed = true)
+
/** A `DataSourceV2ScanRelation` over `table` projecting only the given
columns. */
private def v2ScanReadingOn(table: TestV2Table, cols: Seq[String]):
DataSourceV2ScanRelation = {
val fullOutput = toAttributes(table.schema())
@@ -2000,6 +2028,27 @@ class MergeSubplansSuite extends PlanTest {
private def v2ScanReading(cols: String*): DataSourceV2ScanRelation =
v2ScanReadingOn(v2Table, cols)
+ /**
+ * A `DataSourceV2ScanRelation` over `table` whose reported
partitioning/ordering is derived the
+ * way the optimizer derives it: by running `V2ScanPartitioningAndOrdering`
over the table's own
+ * scan, resolving transforms through a function catalog. Every call
therefore binds the transform
+ * function afresh, so the two sides of a merge hold DIFFERENT
`BoundFunction` instances --
+ * exactly as in production, where each scan relation is annotated by its
own derivation.
+ */
+ private def v2ScanReportingOn(
+ table: TestV2Table, cols: Seq[String]): DataSourceV2ScanRelation = {
+ val fullOutput = toAttributes(table.schema())
+ val relation = DataSourceV2Relation(
+ table, fullOutput, Some(TestFreshBindFunctionCatalog), None,
CaseInsensitiveStringMap.empty())
+ val output = cols.map(c => fullOutput.find(_.name == c).get)
+ val builder = table.newScanBuilder(CaseInsensitiveStringMap.empty())
+ builder.asInstanceOf[SupportsPushDownRequiredColumns].pruneColumns(
+ StructType(output.map(a => StructField(a.name, a.dataType, a.nullable))))
+ V2ScanPartitioningAndOrdering(
+ DataSourceV2ScanRelation(relation, builder.build(), output,
mergeableScan = true))
+ .asInstanceOf[DataSourceV2ScanRelation]
+ }
+
/** A `DataSourceV2ScanRelation` over a table without the `SCAN_MERGING`
capability. */
private def v2ScanReadingNoMerge(cols: String*): DataSourceV2ScanRelation =
v2ScanReadingOn(v2TableNoMerge, cols)
@@ -2599,11 +2648,13 @@ 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.)
+ test("SPARK-40259: do not merge DSv2 scans when the merge would degrade a
reported " +
+ "key-grouped partitioning or ordering") {
+ // 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 +2669,217 @@ 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,
+ bothSides: Boolean = false): Unit = {
+ val sub1 =
ScalarSubquery(withField(v2ScanReading("a")).groupBy()(sum($"a").as("sum_a")))
+ val cpScan = if (bothSides) withField(v2ScanReading("b")) else
v2ScanReading("b")
+ val sub2 = ScalarSubquery(cpScan.groupBy()(sum($"b").as("sum_b")))
+ val originalQuery = testRelation.select(sub1, sub2)
+ withSQLConf(confKey -> "true") {
+ comparePlans(Optimize.execute(originalQuery.analyze),
correctAnswer.analyze)
+ }
+ }
+
+ // One side reports: the report combines fine, and the config lets through
the degradation the
+ // post-rebuild check finds (the rebuilt scan re-derives nothing).
+ 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)
+ // Both sides report, each on the column it reads: INCOMPATIBLE, so here
the config is what
+ // skips the EARLY decline at the leaf, before any rebuild (the mirror of
the incompatible
+ // default-decline test above).
+ assertMerges(s => s.copy(keyGroupedPartitioning =
Some(Seq(s.output.head))),
+
SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key,
+ bothSides = true)
+ assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head,
Ascending)))),
+ SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key, bothSides =
true)
+ }
+
+ test("SPARK-58549: merge DSv2 scans when the rebuilt scan re-derives the
required ordering") {
+ // The NOT-WORSE side of the check: both inputs report `a ASC` and the
table reports the same,
+ // so the rebuilt merged scan re-derives it and nothing is degraded -- the
merge proceeds with
+ // the configs off. Every other test here rebuilds over a plain
TestV2Scan, which reports
+ // nothing, so any non-empty requirement is degraded by construction and
only the declining side
+ // gets covered.
+ val withOrdering = (s: DataSourceV2ScanRelation) =>
+ s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))
+ val q = testRelation.select(
+ ScalarSubquery(withOrdering(v2ScanReadingOn(v2TableReportingA, Seq("a",
"b")))
+ .groupBy()(sum($"b").as("sum_b"))),
+ ScalarSubquery(withOrdering(v2ScanReadingOn(v2TableReportingA, Seq("a",
"c")))
+ .groupBy()(sum($"c").as("sum_c"))))
+ val optimized = Optimize.execute(q.analyze)
+
+ val scans = v2Scans(optimized)
+ assert(scans.length == 1, s"the two scans should be fused into
one:\n$optimized")
+ assert(scans.head.output.map(_.name).toSet == Set("a", "b", "c"),
+ s"the merged scan should read the union of both columns; got
${scans.head.output}")
+ assert(scans.head.ordering.exists(_.map(_.child).collect {
+ case a: Attribute => a.name
+ } == Seq("a")),
+ s"the merged scan should re-derive the reported ordering; got
${scans.head.ordering}")
+ }
+
+ test("SPARK-58549: do not merge when the rebuilt scan re-derives less than
the combined " +
+ "ordering") {
+ // One input reports `a ASC`, the other `a ASC, c ASC`, so the combined
requirement is the
+ // STRONGER of the two. The table reports only `a ASC`, so the rebuilt
merged scan satisfies the
+ // weaker input's ordering but not the combined one -- a degradation,
declined with the configs
+ // off. Guards that combineRequiredOrdering keeps the stronger side:
keeping the weaker one
+ // would let this merge through and lose the second sort key.
+ val scanAB = v2ScanReadingOn(v2TableReportingA, Seq("a", "b"))
+ val scanAC = v2ScanReadingOn(v2TableReportingA, Seq("a", "c"))
+ val q = testRelation.select(
+ ScalarSubquery(scanAB.copy(ordering =
Some(Seq(SortOrder(scanAB.output.head, Ascending))))
+ .groupBy()(sum($"b").as("sum_b"))),
+ ScalarSubquery(scanAC.copy(ordering = Some(Seq(
+ SortOrder(scanAC.output.head, Ascending),
SortOrder(scanAC.output.last, Ascending))))
+ .groupBy()(sum($"c").as("sum_c"))))
+ comparePlans(Optimize.execute(q.analyze), q.analyze)
+ }
+
+ test("SPARK-58549: merge DSv2 scans reporting the same bucket transform
partitioning") {
+ // Both inputs report `bucket(4, a)`, each derived independently, so each
holds its OWN
+ // BoundFunction instance -- what production does, since V2ExpressionUtils
binds the function
+ // afresh per derivation and a BoundFunction defines no `equals`.
Comparing the two reports by
+ // canonical equality would therefore call them different and decline the
merge; they must be
+ // compared by transform semantics (isSameFunction), and then the merge
proceeds and the rebuilt
+ // scan re-derives the same report. This is the only test that reaches the
transform arm of
+ // sameReportedExpressions: an identity-partitioned source reports a plain
attribute instead.
+ val q = testRelation.select(
+ ScalarSubquery(v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "b"))
+ .groupBy()(sum($"b").as("sum_b"))),
+ ScalarSubquery(v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "c"))
+ .groupBy()(sum($"c").as("sum_c"))))
+ val optimized = Optimize.execute(q.analyze)
+
+ val scans = v2Scans(optimized)
+ assert(scans.length == 1, s"the two scans should be fused into
one:\n$optimized")
+ val kgp = scans.head.keyGroupedPartitioning
+ assert(kgp.exists(_.forall(_.isInstanceOf[TransformExpression])),
+ s"the merged scan should preserve the reported bucket transform; got
$kgp")
+ assert(kgp.get.flatMap(_.references).exists(_.name == "a"),
+ s"the preserved partitioning should be on a; got $kgp")
+ }
+
+ 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, <col>} 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-58549: the deferred build checks the required report per
attempt") {
+ // The deferred build tries strict + best-effort first, then strict-only.
The report is checked
+ // on EACH attempt, not once on whichever came back, so the deferred path
stays no weaker than
+ // the leaf path: the strict-only attempt is exactly what the leaf builds
when no Filter sits
+ // above the scan, so a merge it can satisfy must not be lost just because
the first attempt
+ // could not. Here both inputs report `a ASC` and neither has a strict
pushed filter, and the
+ // source reports only when nothing was pushed: attempt 1 (offered the
Filter's condition)
+ // re-derives no ordering and is rejected, attempt 2 re-derives `a ASC`
and the merge lands.
+ val withOrdering = (s: DataSourceV2ScanRelation) =>
+ s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))
+ val q = testRelation.select(
+ ScalarSubquery(withOrdering(v2ScanReadingOn(v2TablePushSensitiveReport,
Seq("a", "b")))
+ .where($"a" > 1).groupBy()(sum($"b").as("sum_b"))),
+ ScalarSubquery(withOrdering(v2ScanReadingOn(v2TablePushSensitiveReport,
Seq("a", "c")))
+ .where($"a" > 1).groupBy()(sum($"c").as("sum_c"))))
+ val optimized = Optimize.execute(q.analyze)
+
+ val scans = v2Scans(optimized)
+ assert(scans.length == 1, s"the two scans should be fused into
one:\n$optimized")
+ assert(scans.head.ordering.exists(_.map(_.child).collect {
+ case a: Attribute => a.name
+ } == Seq("a")),
+ s"the strict-only attempt should have satisfied the required ordering; "
+
+ s"got ${scans.head.ordering}")
+
assert(scans.head.scan.asInstanceOf[TestV2ReportingScan].inner.pushed.isEmpty,
+ "the surviving build should be the strict-only one, with nothing pushed")
+ }
+
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
@@ -2785,7 +3047,10 @@ private class TestV2Table(
tableSchema: StructType,
acceptsFilters: Boolean = true,
strictColumns: Set[String] = Set.empty,
- supportsScanMerging: Boolean = true)
+ supportsScanMerging: Boolean = true,
+ reportedOrderingCols: Seq[String] = Nil,
+ reportedBucketPartition: Option[(Int, String)] = None,
+ reportOnlyWhenNothingPushed: Boolean = false)
extends Table with SupportsRead {
override def name(): String = "test_v2_table"
override def schema(): StructType = tableSchema
@@ -2795,13 +3060,17 @@ private class TestV2Table(
caps
}
override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
- new TestV2ScanBuilder(tableSchema, acceptsFilters, strictColumns)
+ new TestV2ScanBuilder(tableSchema, acceptsFilters, strictColumns,
reportedOrderingCols,
+ reportedBucketPartition, reportOnlyWhenNothingPushed)
}
private class TestV2ScanBuilder(
tableSchema: StructType,
acceptsFilters: Boolean,
- strictColumns: Set[String])
+ strictColumns: Set[String],
+ reportedOrderingCols: Seq[String] = Nil,
+ reportedBucketPartition: Option[(Int, String)] = None,
+ reportOnlyWhenNothingPushed: Boolean = false)
extends ScanBuilder with SupportsPushDownRequiredColumns with
SupportsPushDownV2Filters
with SupportsPushDownLimit with SupportsPushDownOffset with
SupportsPushDownTopN
with SupportsPushDownTableSample {
@@ -2838,7 +3107,15 @@ private class TestV2ScanBuilder(
p.references().forall(r =>
strictColumns.contains(r.fieldNames().mkString(".")))
override def pushedPredicates(): Array[Predicate] = accepted
- override def build(): Scan = TestV2Scan(prunedSchema,
accepted.map(_.describe()).toSeq)
+ override def build(): Scan = {
+ val scan = TestV2Scan(prunedSchema, accepted.map(_.describe()).toSeq)
+ if (reportedOrderingCols.isEmpty && reportedBucketPartition.isEmpty) {
+ scan
+ } else {
+ TestV2ReportingScan(scan, reportedOrderingCols, reportedBucketPartition,
+ reportOnlyWhenNothingPushed)
+ }
+ }
}
/** `pushed` records the `describe()` of the predicates pushed to the scan,
for test assertions. */
@@ -2846,3 +3123,75 @@ private case class TestV2Scan(schema: StructType,
pushed: Seq[String] = Seq.empt
extends Scan {
override def readSchema(): StructType = schema
}
+
+/**
+ * A [[TestV2Scan]] that also reports an output ordering and/or a key-grouped
partitioning, so a
+ * scan rebuilt by the merge re-derives a report through
`V2ScanPartitioningAndOrdering`. The plain
+ * [[TestV2Scan]] reports nothing, which only ever exercises the degrading
side of the not-worse
+ * check.
+ */
+private case class TestV2ReportingScan(
+ inner: TestV2Scan,
+ orderingCols: Seq[String] = Nil,
+ bucketPartition: Option[(Int, String)] = None,
+ onlyWhenNothingPushed: Boolean = false)
+ extends Scan with SupportsReportOrdering with SupportsReportPartitioning {
+ override def readSchema(): StructType = inner.readSchema()
+
+ // A source may report per scan: `outputOrdering`/`outputPartitioning` are
answered by the Scan
+ // the connector built AFTER seeing what it accepted, so an ordering that
only holds over the
+ // unpruned file set can legitimately disappear once a filter is pushed.
`onlyWhenNothingPushed`
+ // models that, so a build offered a best-effort filter and one without it
report differently.
+ private def reports: Boolean = !onlyWhenNothingPushed || inner.pushed.isEmpty
+
+ override def outputOrdering(): Array[V2SortOrder] = if (reports) {
+ orderingCols.map(c => Expressions.sort(FieldReference(c),
V2SortDirection.ASCENDING)).toArray
+ } else {
+ Array.empty
+ }
+
+ override def outputPartitioning(): V2Partitioning = bucketPartition.filter(_
=> reports) match {
+ case Some((numBuckets, col)) =>
+ new V2KeyGroupedPartitioning(Array(Expressions.bucket(numBuckets, col)),
0)
+ case None => new V2UnknownPartitioning(0)
+ }
+}
+
+/**
+ * A `FunctionCatalog` that resolves `bucket`, binding it to a FRESH function
instance every time --
+ * as a real connector does, since `V2ExpressionUtils.loadV2FunctionOpt` binds
the function afresh
+ * for every derivation of a reported transform. Spark's own
`UnboundBucketFunction` hands back a
+ * singleton, which would hide the fact that two independently derived reports
have to be
+ * compared by transform semantics rather than by function instance.
+ */
+private object TestFreshBindFunctionCatalog extends FunctionCatalog {
+ override def initialize(name: String, options: CaseInsensitiveStringMap):
Unit = {}
+ override def name(): String = "test_fresh_bind"
+ override def listFunctions(namespace: Array[String]): Array[Identifier] =
+ Array(Identifier.of(Array.empty, "bucket"))
+ override def loadFunction(ident: Identifier): UnboundFunction = ident.name()
match {
+ case "bucket" => TestUnboundBucketFunction
+ // loadV2FunctionOpt treats this as "function not found" and reports no
partitioning.
+ case other => throw new UnsupportedOperationException(s"no such function:
$other")
+ }
+}
+
+private object TestUnboundBucketFunction extends UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction = new
TestBucketFunction
+ override def description(): String = name()
+ override def name(): String = "bucket"
+}
+
+/**
+ * A bound `bucket` with no `equals`, so two instances are never `==` -- which
is what makes the
+ * fixture able to catch an instance-sensitive comparison. `canonicalName` is
stable, as the
+ * `BoundFunction` contract requires (its default returns a fresh random
UUID); that stable name is
+ * what identifies two separately bound instances as the same transform, for a
storage-partitioned
+ * join and for the merge's not-worse check alike.
+ */
+private class TestBucketFunction extends ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "testcat.bucket"
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
