ulysses-you commented on code in PR #58522:
URL: https://github.com/apache/spark/pull/58522#discussion_r3946149353
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala:
##########
@@ -96,7 +96,17 @@ trait DataSourceV2ScanExecBase
val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes)
val partitionKeys =
inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
- KeyedPartitioning(exprs, partitionKeys)
+ val partitioning = KeyedPartitioning(exprs, partitionKeys)
+ // A partition key may reference a column that was pruned out of the
scan output (kept only
+ // when operation keys may be a subset of the partition keys, see
+ // V2ScanPartitioningAndOrdering). Project such unresolvable key
positions away so the
+ // reported partitioning only references output columns.
+ val resolvablePositions = exprs.indices.filter(i =>
exprs(i).references.subsetOf(outputSet))
+ if (resolvablePositions.isEmpty) {
+ super.outputPartitioning
+ } else {
+ partitioning.project(resolvablePositions)
Review Comment:
Same issue @sunchao raised; fixed in 973bde7e6ce.
`reportedKeyedPartitioning` carries the source's
report at full key width and is what `filteredPartitions` passes down;
`outputPartitioning` stays
projected for planning. `replanWithRuntimeFilters` now takes
`Option[KeyedPartitioning]`, so the
projected view cannot reach the raw key rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala:
##########
@@ -41,33 +42,45 @@ object V2ScanPartitioningAndOrdering extends
Rule[LogicalPlan] with Logging {
}
}
- private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning(
+ private def partitioning(plan: LogicalPlan) = {
+ val allowKeysSubsetOfPartitionKeys =
SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys
+ plan.transformDownWithPruning(
_.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) {
- case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _)
- if d.keyGroupedPartitioning.isEmpty =>
- val catalystPartitioning = scan.outputPartitioning() match {
- case kgp: KeyGroupedPartitioning =>
- val partitioning = sequenceToOption(
- kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation,
relation.funCatalog))
- .toImmutableArraySeq)
- if (partitioning.isEmpty) {
- None
- } else {
- if (partitioning.get.forall(p =>
p.references.subsetOf(d.outputSet))) {
- partitioning
- } else {
+ case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _)
+ if d.keyGroupedPartitioning.isEmpty =>
+ val catalystPartitioning = scan.outputPartitioning() match {
+ case kgp: KeyGroupedPartitioning =>
+ val partitioning = sequenceToOption(
+ kgp.keys().map(V2ExpressionUtils.toCatalystOpt(_, relation,
relation.funCatalog))
+ .toImmutableArraySeq)
+ if (partitioning.isEmpty) {
None
+ } else {
+ val inOutput = partitioning.get.map(p =>
p.references.subsetOf(d.outputSet))
+ if (inOutput.forall(identity)) {
+ partitioning
+ } else if (inOutput.exists(identity) &&
allowKeysSubsetOfPartitionKeys) {
Review Comment:
Agreed, gate removed in 973bde7e6ce.
Your mechanism holds: `project` marks `isCollapsed` exactly on the
collapsing case and
`mayGroupToSatisfy` keeps the config over it, so the rule only has to decide
whether any key
survives. The rule also has no partition key values in hand, so it cannot
tell the two cases apart
even in principle.
The case the conjunct was costing is now pinned by `a pruned key that
collapses nothing keeps SPJ
without the config`, with `allowKeysSubsetOfPartitionKeys` set to `false`
explicitly.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2526,6 +2526,233 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59248: join key subset of partition keys, extra partition key
pruned from the " +
+ "output") {
+ // Both tables are partitioned by (id, data). The join is only on `id`,
and `data` is not
+ // selected, so it is column-pruned out of both scan outputs. Without
+ // allowKeysSubsetOfPartitionKeys the pruned `data` key drops the reported
partitioning and both
+ // sides shuffle; with it, the partitioning is kept and projected onto
`id`, so SPJ triggers.
+ val table1 = "prune_t1"
+ val table2 = "prune_t2"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-01' as timestamp))")
+
+ createTable(table2, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'ee', cast('2020-01-01' as timestamp)), " +
+ "(4, 'ff', cast('2020-01-01' as timestamp))")
+
+ // Selecting only `id` prunes the other partition key `data` (and `ts`)
from both scans. The
+ // expected result is the within-`id` cross product (id=2 matches 2 x 2
rows, id=3 matches
+ // 1 x 1).
+ val expected = Seq(Row(2), Row(2), Row(2), Row(2), Row(3))
+ val query =
+ s"""
+ |${selectWithMergeJoinHint("t1", "t2")}
+ |t1.id AS id
+ |FROM testcat.ns.$table1 t1 JOIN testcat.ns.$table2 t2
+ |ON t1.id = t2.id ORDER BY id
+ |""".stripMargin
+
+ Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowKeysSubsetOfPartitionKeys.toString) {
+ val df = sql(query)
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ val groupPartitions =
collectGroupPartitions(df.queryExecution.executedPlan)
+ if (allowKeysSubsetOfPartitionKeys) {
+ assert(shuffles.isEmpty, "SPJ should be triggered even though `data`
is pruned")
+ assert(groupPartitions.nonEmpty, "GroupPartitionsExec should
coalesce on the join key")
+ // The reported partitioning is kept on the scan even though `data`
is pruned ...
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.nonEmpty)
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isDefined,
+ "partitioning should be kept despite the pruned key")
+ // ... but the physical output partitioning must only reference
output columns, so the
+ // pruned column reaches no consumer (shuffle spec, ordering, plan
equality).
+ scan.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+
assert(kp.expressions.forall(_.references.subsetOf(scan.outputSet)),
+ s"partitioning ${kp.expressions} references a column outside
${scan.output}")
+ case other =>
+ fail(s"expected KeyedPartitioning but got $other")
+ }
+ }
+ } else {
+ assert(shuffles.nonEmpty, "SPJ should not be triggered without the
config")
+ assert(groupPartitions.isEmpty)
+ }
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ test("SPARK-59248: scan reports no partitioning when all partition keys are
pruned") {
+ // The table is partitioned by (id, data), but the query selects only
`ts`, so both partition
+ // keys are column-pruned out of the scan output. Even with
allowKeysSubsetOfPartitionKeys on,
+ // no partition key survives in the output, so the scan must not keep a
dangling
+ // KeyedPartitioning and reports no (unknown) partitioning.
+ val table1 = "prune_all_keys"
+ createTable(table1, columns, Array(identity("id"), identity("data")))
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-02' as timestamp))")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(s"SELECT ts FROM testcat.ns.$table1")
+ checkAnswer(df, Seq(
+ Row(Timestamp.valueOf("2020-01-01 00:00:00")),
+ Row(Timestamp.valueOf("2020-01-02 00:00:00"))))
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.length == 1)
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isEmpty,
+ s"no partition key survives in the output, got
${scan.keyGroupedPartitioning}")
+ scan.outputPartitioning match {
+ case _: physical.UnknownPartitioning => // expected: nothing left to
partition by
+ case other => fail(s"expected UnknownPartitioning but got $other")
+ }
+ }
+ }
+ }
+
+ test("SPARK-59248: self-join with a pruned partition key keeps plans
canonicalizable") {
+ // Same-table join where the extra partition key `data` is pruned from
both scan instances. This
+ // exercises canonicalization/plan-equality over scans whose reported
partitioning carries a key
+ // that is not in the scan output; results must stay correct and planning
must not fail.
+ val table1 = "prune_self"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-01' as timestamp)), " +
+ "(2, 'cc', cast('2020-01-01' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-01' as timestamp))")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("a", "b")}
+ |a.id AS id
+ |FROM testcat.ns.$table1 a JOIN testcat.ns.$table1 b
+ |ON a.id = b.id ORDER BY id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty, "SPJ
should be triggered")
+ // id=1 yields 1 row, id=2 yields 2 x 2 = 4 rows, id=3 yields 1 row.
+ checkAnswer(df, Seq(Row(1), Row(2), Row(2), Row(2), Row(2), Row(3)))
+
+ // Both scan instances must survive (no incorrect dedup) and each must
report a partitioning
+ // that only references its own output, even though `data` is pruned:
this is what keeps the
+ // dangling key from reaching any consumer (shuffle spec, ordering,
canonicalized comparison).
+ val scans = collectScans(df.queryExecution.executedPlan)
+ assert(scans.length == 2, s"expected the two self-join scans, got:\n" +
+ s"${df.queryExecution.executedPlan}")
+ scans.foreach { scan =>
+ assert(scan.keyGroupedPartitioning.isDefined,
+ "partitioning should be kept despite the pruned key")
+ scan.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+
assert(kp.expressions.forall(_.references.subsetOf(scan.outputSet)),
+ s"partitioning ${kp.expressions} references a column outside
${scan.output}")
+ case other =>
+ fail(s"expected KeyedPartitioning but got $other")
+ }
+ // Canonicalization must be stable and must not throw with a dangling
key present.
+ assert(scan.canonicalized.sameResult(scan.canonicalized))
+ }
+ }
+ }
+
+ test("SPARK-59248: a pruned partition key must not defeat plan reuse") {
+ val table1 = "prune_reuse"
+ val partition = Array(identity("id"), identity("data"))
+ createTable(table1, columns, partition)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(1, 'aa', cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', cast('2020-01-02' as timestamp)), " +
+ "(3, 'dd', cast('2020-01-03' as timestamp))")
+
+ // Self-join on the non-partition column `ts`; the other partition key
`data` is pruned from
+ // both scan instances. The two legs are identical subtrees, so Spark
reuses one leg's exchange
+ // for the other. With allowKeysSubsetOfPartitionKeys the scan keeps its
reported partitioning,
+ // which then references the pruned `data`; that dangling key must not
leak into canonicalized
+ // plan comparison and break the reuse.
+ val query =
+ s"""
+ |SELECT a.id AS id1, b.id AS id2
+ |FROM testcat.ns.$table1 a JOIN testcat.ns.$table1 b
+ |ON a.ts = b.ts ORDER BY id1, id2
+ |""".stripMargin
+
+ Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
+ allowKeysSubsetOfPartitionKeys.toString) {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 2), Row(3, 3)))
+ val plan = df.queryExecution.executedPlan
+ val reused = collect(plan) { case r: ReusedExchangeExec => r }
+ val scans = collectScans(plan)
+ assert(scans.length == 1,
+ s"the two identical legs should reuse a single scan " +
+
s"(allowKeysSubsetOfPartitionKeys=$allowKeysSubsetOfPartitionKeys):\n$plan")
+ assert(reused.length == 1,
+ s"expected one reused exchange " +
+
s"(allowKeysSubsetOfPartitionKeys=$allowKeysSubsetOfPartitionKeys):\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59248: a pruned source-reported ordering must not defeat
exchange reuse") {
Review Comment:
Confirmed, and the test is removed in 973bde7e6ce.
The ordering half of `doCanonicalize` is pinned instead by
`MergeSubplansSuite`'s `identical DSv2
scans whose reported ordering is on a pruned column are deduplicated, not
fused`, which does fail
with the `takeWhile` reverted.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala:
##########
@@ -41,33 +42,45 @@ object V2ScanPartitioningAndOrdering extends
Rule[LogicalPlan] with Logging {
}
}
- private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning(
+ private def partitioning(plan: LogicalPlan) = {
+ val allowKeysSubsetOfPartitionKeys =
SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys
Review Comment:
Moot now that the gate is gone (finding 2): the config keeps the two roles
it already documents, so
`V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.doc()` is unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
