peter-toth commented on code in PR #58522:
URL: https://github.com/apache/spark/pull/58522#discussion_r3933688559
##########
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:
**Finding 1.** `project` rebuilds the key rows onto `resolvablePositions`,
so from here the partitioning's `partitionKeys` are projected rows and its
`keyDataTypes` are the projected types. `BatchScanExec.filteredPartitions`
passes this same partitioning to `PushDownUtils.replanWithRuntimeFilters`,
which reads the raw, full-width `HasPartitionKey.partitionKey()` rows with it:
-
`parts.sortBy(_.asInstanceOf[HasPartitionKey].partitionKey())(k.keyRowOrdering)`
at
`sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:357`
- `getInternalRowComparableWrapperFactory(k.keyDataTypes)` applied to
`partitionKey()` at
`sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:316`
Those line up only when the surviving positions are the leading prefix.
Every existing test prunes the trailing key, so nothing catches it.
**Wrong results.** `GroupPartitionsExec.grouping` labels child RDD partition
`i` with `childKp.partitionKeys(i)`. `partitionKeys` is ordered by the full
key, while `filteredPartitions` is ordered by the projected prefix. The two
disagree wherever the full-key sort reorders inside a prefix group. Measured on
a 7-split table partitioned by `(store_id, dept_id)` with `store_id` pruned:
rdd partition keys (store,dept): (1,10), (1,20), (1,30), (2,5), (2,40),
(3,7), (3,1)
partitioning keys (dept): 10, 20, 30, 5, 40,
1, 7
Indices 5 and 6 carry each other's label. A join that lays the other side
out on those keys loses exactly those two:
```scala
val cols = Array(
Column.create("store_id", IntegerType),
Column.create("dept_id", IntegerType))
createTable("t", cols, Array(identity("store_id"), identity("dept_id")))
sql("INSERT INTO testcat.ns.t VALUES (1, 20), (1, 10), (1, 30), (2, 5), (2,
40), (3, 7), (3, 1)")
withTempView("other") {
spark.range(1, 41).selectExpr("cast(id as int) as
dept_id").createOrReplaceTempView("other")
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
val df = sql("SELECT /*+ MERGE(t, o) */ t.dept_id " +
"FROM testcat.ns.t t JOIN other o ON t.dept_id = o.dept_id")
assert(df.collect().map(_.getInt(0)).sorted.toSeq === Seq(1, 5, 7, 10,
20, 30, 40))
}
}
```
This branch returns `5, 10, 20, 30, 40`. With
`allowKeysSubsetOfPartitionKeys` off the same query is correct, so it is this
path.
**ClassCastException** when the pruned leading key has a different type. Two
tables partitioned by `(identity("data"), identity("id"))`, joined on `id` with
`data` pruned, throw `java.lang.ClassCastException: class
org.apache.spark.unsafe.types.UTF8String cannot be cast to class
java.lang.Integer` at `PushDownUtils.scala:357`. `keyRowOrdering` comes from
the projected `[IntegerType]` and is applied to a `(String, Int)` row.
**Fix.** Give `replanWithRuntimeFilters` the unprojected partitioning: its
keys are the raw rows, in the order it sorts the input partitions into.
Splitting the construction out of `outputPartitioning` is enough:
```scala
/**
* The partitioning as the source reported it: one key per input
partition, holding every reported
* key position, in the order a consumer must sort the input partitions
into.
*/
protected def reportedKeyedPartitioning: Option[KeyedPartitioning] = {
keyGroupedPartitioning match {
case Some(exprs) if conf.v2BucketingEnabled &&
KeyedPartitioning.supportsExpressions(exprs) &&
inputPartitions.nonEmpty &&
inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) =>
val dataTypes = exprs.map(_.dataType)
val rowOrdering =
RowOrdering.createNaturalAscendingOrdering(dataTypes)
val partitionKeys =
inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
Some(KeyedPartitioning(exprs, partitionKeys))
case _ => None
}
}
/**
* What orders the input partitions. `outputPartitioning` may have
projected key positions away,
* and its keys and key types then no longer describe the raw
`HasPartitionKey.partitionKey()`
* rows, so a consumer reading those rows must take the types and the
order from here.
*/
protected def inputPartitionOrdering: physical.Partitioning =
if (outputPartitioning.isInstanceOf[KeyedPartitioning]) {
reportedKeyedPartitioning.getOrElse(super.outputPartitioning)
} else {
super.outputPartitioning
}
override def outputPartitioning: physical.Partitioning =
reportedKeyedPartitioning match {
case Some(partitioning) =>
val exprs = partitioning.expressions
val resolvablePositions = exprs.indices.filter(i =>
exprs(i).references.subsetOf(outputSet))
if (resolvablePositions.isEmpty) super.outputPartitioning
else partitioning.project(resolvablePositions)
case _ => super.outputPartitioning
}
```
with `BatchScanExec.filteredPartitions` passing `inputPartitionOrdering`
instead of `outputPartitioning`. I ran `KeyGroupedPartitioningSuite`,
`DataSourceV2CatalystRuntimeFilterSuite` and `MergeSubplansSuite` with that
applied plus the two cases above: 245 tests, all green.
##########
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:
**Finding 2.** The opt-in is needed for the grouping, not for the report.
Projecting a key position away is a collapse only when two distinct source keys
land on the same projected key. `KeyedPartitioning.project` sets `isCollapsed`
exactly then, and `mayGroupToSatisfy` refuses to group a collapsed partitioning
without this config. The safety property is therefore already enforced one
layer down, on the case that needs it.
What this conjunct costs is the other case: a pruned key whose removal
collapses nothing. The projected keys stay unique, `isGrouped` is true, and
`satisfies(ClusteredDistribution)` holds whatever the config says.
That split is the one SPARK-46367 (`e656d04c157`) already drew for the
sibling narrowing in `PartitioningPreservingUnaryExecNode`: distinct projected
keys need no config, duplicate projected keys require
`allowKeysSubsetOfPartitionKeys`. `KeyedPartitioning`'s class doc argues the
same for not gating the report - "Reporting `UnknownPartitioning` would give up
all of them, and make the plan shape depend on a config."
I measured both halves with the conjunct dropped, each with
`allowKeysSubsetOfPartitionKeys=false`:
- partitioned by `(dept_id, store_id)`, one `store_id` per `dept_id`,
`store_id` pruned, join on `dept_id`: no shuffle, correct answer.
- partitioned by `(id, data)`, two `data` values per `id`, `data` pruned,
join on `id`: two shuffles, correct answer. `mayGroupToSatisfy` refuses, as it
should.
So `else if (inOutput.exists(identity))` looks like the right gate.
##########
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:
**Finding 3.** This one passes on master. `BatchScanExec` never looks at
`ordering` - `equals`, `hashCode` and `doCanonicalize` all leave it out - so a
dangling reported ordering could not block physical exchange reuse in the first
place.
Measured two ways. On `e261626f152` with only the test files applied, it
passes. With the whole PR applied but the `takeWhile` in
`DataSourceV2ScanRelation.doCanonicalize` reverted, it still passes. The other
tests behave as expected under the same treatment: "join key subset ..." and
"self-join ..." fail on base, and "a pruned partition key must not defeat plan
reuse" fails when only `BatchScanExec.scala` is reverted.
The `MergeSubplansSuite` test is the one that covers the `takeWhile` - it
fails with that hunk alone reverted, matching your description. So either drop
this test, or point it at the logical path (`sameResult` over two
`DataSourceV2ScanRelation`s), where the ordering really is compared.
##########
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:
**Finding 4.** If the gate stays (see finding 2),
`V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.doc()` needs a sentence for
it. The doc enumerates the config's roles today - allowing the operation keys
to be a subset of the partition keys, and gating the grouping of a collapsed
partitioning. This adds a third that differs in kind: whether the scan reports
a partitioning at all. That is the one a user meets as "the plan shape
changed", so it is worth naming.
--
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]
