mbutrovich opened a new issue, #7441:
URL: https://github.com/apache/datafusion-comet/issues/7441
### Describe the bug
With AQE on, a query fails with `java.lang.AssertionError: assertion failed`
in `AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStage` when the build side
of a dynamic partition pruning (DPP) broadcast contains a storage-partitioned
join over native Iceberg scans. A storage-partitioned join is a sort-merge join
that Spark plans without a shuffle because both scans report
`KeyGroupedPartitioning`.
```
org.apache.spark.SparkException: Exception thrown in awaitResult:
at org.apache.spark.util.ThreadUtils$.awaitResult(ThreadUtils.scala:326)
at
org.apache.spark.sql.comet.CometSubqueryBroadcastExec.executeCollect(CometSubqueryBroadcastExec.scala:149)
at
org.apache.spark.sql.execution.InSubqueryExec.updateResult(subquery.scala:129)
...
Caused by: java.lang.AssertionError: assertion failed
at scala.Predef$.assert(Predef.scala:208)
at
org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStage(AdaptiveSparkPlanExec.scala:648)
at
org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.newQueryStage(AdaptiveSparkPlanExec.scala:622)
at
org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.createQueryStages(AdaptiveSparkPlanExec.scala:540)
...
```
The trace is from Spark 3.5. The same assertion fails on Spark 3.4, 4.0, and
4.1. On Spark 3.4 the DPP subquery is Spark's `SubqueryBroadcastExec` instead
of `CometSubqueryBroadcastExec`. The `spark-4.2` profile has no Iceberg
runtime, so the test below does not run there.
It is a regression from #7231. The query passes at 7f50037fd, the parent of
#7231, where the DPP subquery reuses the join's broadcast through a
`ReusedExchange`. It fails at ce3e665e8.
### Steps to reproduce
Add this test to `CometIcebergNativeSuite`:
```scala
test("AQE DPP - storage-partitioned join on the DPP build side") {
assume(icebergAvailable, "Iceberg not available")
withTempIcebergDir { warehouseDir =>
withSQLConf(
"spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog",
"spark.sql.catalog.aqe_cat.type" -> "hadoop",
"spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath,
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
SQLConf.V2_BUCKETING_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
"spark.sql.iceberg.planning.preserve-data-grouping" -> "true",
CometConf.COMET_ENABLED.key -> "true",
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {
spark.sql("""
CREATE TABLE aqe_cat.db.probe_fact (amount INT, store_id INT)
USING iceberg PARTITIONED BY (store_id)
""")
spark.sql("""
CREATE TABLE aqe_cat.db.probe_dim_a (id INT, country STRING)
USING iceberg PARTITIONED BY (id)
""")
spark.sql("""
CREATE TABLE aqe_cat.db.probe_dim_b (id INT, name STRING)
USING iceberg PARTITIONED BY (id)
""")
spark.sql(
"INSERT INTO aqe_cat.db.probe_fact SELECT CAST(id AS INT), CAST(id %
10 AS INT) " +
"FROM range(0, 100)")
spark.sql(
"INSERT INTO aqe_cat.db.probe_dim_a SELECT CAST(id AS INT), " +
"IF(id < 3, 'X', 'Y') FROM range(0, 10)")
spark.sql(
"INSERT INTO aqe_cat.db.probe_dim_b SELECT CAST(id AS INT),
concat('n', id) " +
"FROM range(0, 10)")
val query =
"""SELECT /*+ BROADCAST(d) */ f.store_id, f.amount, d.name
|FROM aqe_cat.db.probe_fact f JOIN (
| SELECT /*+ MERGE(a, b) */ a.id, b.name
| FROM aqe_cat.db.probe_dim_a a JOIN aqe_cat.db.probe_dim_b b ON
a.id = b.id
| WHERE a.country = 'X'
|) d ON f.store_id = d.id""".stripMargin
checkSparkAnswer(query)
spark.sql("DROP TABLE aqe_cat.db.probe_fact")
spark.sql("DROP TABLE aqe_cat.db.probe_dim_a")
spark.sql("DROP TABLE aqe_cat.db.probe_dim_b")
}
}
}
```
Then run it:
```sh
make core
./mvnw test -Pspark-3.5 -Dtest=none
-Dsuites="org.apache.comet.CometIcebergNativeSuite AQE DPP -
storage-partitioned join on the DPP build side"
```
### Expected behavior
The query returns the same rows as Spark, and the DPP subquery reuses the
join's broadcast, as it does at 7f50037fd.
### Additional context
To build the DPP subquery's plan, Comet wraps the build side it already
converted in a new `AdaptiveSparkPlanExec`
([`CometPlanAdaptiveDynamicPruningFilters.scala#L262-L302`](https://github.com/apache/datafusion-comet/blob/79b5ddd50ba2e3f605399d641e4c0b3be84c20ef/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala#L262-L302)).
On Spark 3.4, Spark's own rule does the same
([`PlanAdaptiveDynamicPruningFilters.scala#L48-L61`](https://github.com/apache/spark/blob/1eb558c3a6fbdd59e5a305bc3ab12ce748f6511f/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala#L48-L61)).
The `AdaptiveSparkPlanExec` constructor runs the query stage preparation rules
over that plan again, including `EnsureRequirements`
([`AdaptiveSparkPlanExec.scala#L113-L134`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec
.scala#L113-L134),
[`#L204-L208`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala#L204-L208)).
`BatchScanExec` reports `KeyGroupedPartitioning` for these scans
([`DataSourceV2ScanExecBase.scala#L93-L103`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala#L93-L103)),
but `CometIcebergNativeScanExec` reports `UnknownPartitioning(0)`
([`CometIcebergNativeScanExec.scala#L138`](https://github.com/apache/datafusion-comet/blob/79b5ddd50ba2e3f605399d641e4c0b3be84c20ef/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala#L138)).
Since #7231, `CometSortMergeJoinExec` declares Spark's required child
distribution
([`operators.scala#L746-L749`](https://github.com/apache/datafusion-comet/blob/79b5ddd50ba2e3f605399d641e4c0b3be84c20ef/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala#L746-L749),
[`#L3740-L3753`](https://github.com/apache/datafusion-comet/blob/79b5ddd50ba2e3f605399d641e4c0b3be84c20ef/spark/src/main/scala/o
rg/apache/spark/sql/comet/operators.scala#L3740-L3753)). So
`EnsureRequirements` inserts an exchange under each side of the join in the
subquery's plan, although the main plan has none:
```
CometSubqueryBroadcast dynamicpruning#93, [0], [id#61], [id=#835]
+- AdaptiveSparkPlan isFinalPlan=false
+- CometBroadcastExchange [id#61, name#64]
+- CometProject [id#61, name#64], [id#61, name#64]
+- CometSortMergeJoin [id#61], [id#63], Inner
:- CometExchange hashpartitioning(id#61, 10),
ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=830]
: +- CometSort [id#61], [id#61 ASC NULLS FIRST]
: +- CometProject [id#61], [id#61]
: +- CometFilter [id#61, country#62],
(isnotnull(country#62) AND (country#62 = X))
: +- CometIcebergNativeScan [id#61, country#62], ...
[filters=country IS NOT NULL, country = 'X', id IS NOT NULL, groupedBy=id], 3
+- CometExchange hashpartitioning(id#63, 10),
ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=820]
+- CometSort [id#63, name#64], [id#63 ASC NULLS FIRST]
+- CometIcebergNativeScan [id#63, name#64], ...
[filters=id IS NOT NULL, groupedBy=id],
runtimeFilters=[dynamicpruningexpression(id#63 IN dynamicpruning#92)], 10
```
AQE finds no logical link on the inserted exchange or below it, and the
assertion fails
([`AdaptiveSparkPlanExec.scala#L641-L650`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala#L641-L650)).
#7422 is the `SinglePartition` case of the same mismatch on Spark 3.4, and
#7424 fixes it by having the native scan report what `BatchScanExec` reports.
Could the native scan report `originalPlan.outputPartitioning` for
`KeyGroupedPartitioning` too? #6268 moved it to `UnknownPartitioning(0)`
because counting the scan's tasks went through `BatchScanExec.inputRDD`, which
translates the DPP filter before its subquery has a result. As I read Spark
3.5.9, `BatchScanExec.outputPartitioning` reads `groupedPartitions`, which
groups `inputPartitions` from `planInputPartitions()`, not the runtime-filtered
partitions
([`BatchScanExec.scala#L59`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala#L59),
[`#L116-L128`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.s
cala#L116-L128),
[`DataSourceV2ScanExecBase.scala#L106-L112`](https://github.com/apache/spark/blob/7c14a3c28b141cc97a330c4d0f5d2a6da7267f85/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala#L106-L112)).
### Willingness to contribute
I cannot contribute a fix for this bug at this time
--
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]
