You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "LinSimon-901101 (via GitHub)" <gi...@apache.org> on 2026/10/09 09:07:21 UTC
[PR] fix: preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
LinSimon-901101 opened a new pull request, #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814
## Which issue does this PR close?
Closes #6667.
## Rationale for this change
On Spark 3.4, a single-split Iceberg scan can run a global aggregate without a shuffle. When dynamic partition pruning removes the only split, Comet drops Spark's empty execution partition, preventing the aggregate from running. As a result, `COUNT(*)` and `SUM(amount)` return no rows instead of `(0, NULL)`.
## What changes are included in this PR?
- Preserve Spark's empty execution partitions by serializing an empty per-partition Iceberg scan.
- Return an empty native stream with the expected schema before initializing storage or requesting catalog credentials when there are no file scan tasks.
- Add regression tests for the aggregate result with AQE enabled and disabled, and for an empty scan with an unavailable S3 credential-policy provider.
## How are these changes tested?
- Reproduced the original failure on Spark 3.4 with both AQE settings before the fix.
- DPP tests passed locally on Spark 3.4 and 3.5, including the new aggregate and credential-provider regressions.
- [[Full fork CI passed](https://github.com/LinSimon-901101/datafusion-comet/actions/runs/37885859935)] for commit `c9a808c`, including the Spark SQL and Iceberg test matrices.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814#discussion_r4232153856
##########
spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala:
##########
@@ -5150,6 +5152,156 @@ class CometIcebergNativeSuite
}
}
+ Seq(false, true).foreach { aqeEnabled =>
+ test(s"DPP - global aggregate after pruning a single split (AQE=$aqeEnabled)") {
+ assume(icebergAvailable, "Iceberg not available")
+ withTempIcebergDir { warehouseDir =>
+ val dimDir = new File(warehouseDir, "dim_parquet")
+ 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 -> aqeEnabled.toString,
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {
+ val factTable = "aqe_cat.db.single_split_dpp_fact"
+ spark.sql(s"""
+ CREATE TABLE $factTable (amount INT, store_id INT)
+ USING iceberg PARTITIONED BY (store_id)
+ """)
+ try {
+ spark
+ .range(100)
+ .selectExpr("CAST(id AS INT) AS amount", "5 AS store_id")
+ .coalesce(1)
+ .writeTo(factTable)
+ .append()
+
+ // The regression requires one original split: Spark 3.4 can then omit the
+ // exchange between the partial and final global aggregates.
+ val unprunedPlan = spark.table(factTable).queryExecution.executedPlan
+ assertSingleNativeScan(unprunedPlan)
+ val unprunedScan = collectIcebergNativeScans(unprunedPlan).head
+ assert(unprunedScan.numPartitions == 1, s"Expected one split:\n$unprunedPlan")
+ assert(
+ unprunedScan.perPartitionData
+ .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount)
+ .sum == 1,
+ s"Expected one file scan task before pruning:\n$unprunedPlan")
+
+ spark
+ .range(1)
+ .selectExpr("99 AS store_id", "'X' AS country")
+ .write
+ .parquet(dimDir.getAbsolutePath)
+ withTempView("single_split_dim") {
+ spark.read
+ .parquet(dimDir.getAbsolutePath)
+ .createOrReplaceTempView("single_split_dim")
+ val query = s"""SELECT /*+ BROADCAST(d) */ COUNT(*), SUM(f.amount)
+ |FROM $factTable f JOIN single_split_dim d ON f.store_id = d.store_id
+ |WHERE d.country = 'X'""".stripMargin
+ val expected = Seq(Row(0L, null))
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ checkAnswer(spark.sql(query), expected)
+ }
+ val df = spark.sql(query)
+ checkCometAnswer(df, expected)
+
+ val plan = df.queryExecution.executedPlan
+ assertSingleNativeScan(plan)
+ val scan = collectIcebergNativeScans(plan).head
+ if (!isSpark35Plus) {
+ assert(
+ collect(plan) { case exchange: ShuffleExchangeLike => exchange }.isEmpty,
+ s"Expected the single-split aggregate to run without a shuffle:\n$plan")
+ assert(scan.numPartitions == 1, s"Expected one empty execution partition:\n$plan")
+ }
+ assert(
+ collectIcebergDPPSubqueries(plan).nonEmpty,
+ s"Expected executable DPP on the native scan:\n$plan")
+ assert(
+ scan.perPartitionData
+ .forall(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount == 0),
+ s"Expected DPP to prune all file scan tasks:\n$plan")
+ assert(scan.metrics("num_splits").value == 0, "Pruned files must not be read")
+ }
+ } finally {
+ spark.sql(s"DROP TABLE $factTable")
+ }
+ }
+ }
+ }
+ }
+
+ test("DPP - empty Iceberg scan does not request unavailable S3 policy locations") {
+ // Exercise the empty executor scan emitted when DPP preserves Spark's last partition.
+ // A fresh catalog/bucket prevents a previously initialized FileIO from hiding the call.
+ val identity = UUID.randomUUID().toString
+ val common = OperatorOuterClass.IcebergScanCommon
+ .newBuilder()
+ .setMetadataLocation(s"s3://empty-scan-$identity/table/metadata/v1.metadata.json")
+ .setCatalogName(s"empty_scan_$identity")
+ .setDataFileConcurrencyLimit(1)
+ .putCatalogProperties(
+ "s3.comet.credential.provider.class",
+ classOf[TestCometS3LocationScopedCredentialProvider].getName)
+ .addRequiredSchema(
+ OperatorOuterClass.SparkStructField
+ .newBuilder()
+ .setName("amount")
+ .setNullable(true)
+ .setDataType(QueryPlanSerde.serializeDataType(IntegerType).get))
+ .build()
+ val planBytes = OperatorOuterClass.Operator
+ .newBuilder()
+ .setIcebergScan(OperatorOuterClass.IcebergScan.newBuilder().setCommon(common))
+ .build()
+ .toByteArray
+
+ val result = spark.sparkContext
+ .parallelize(Seq(1), 1)
+ .mapPartitions { _ =>
+ val provider = classOf[TestCometS3LocationScopedCredentialProvider]
+ val failureField = provider.getDeclaredField("throwOnNextLocationCall")
+ failureField.setAccessible(true)
+ val countField = provider.getDeclaredField("locationCallCount")
+ countField.setAccessible(true)
+ val calls = countField.get(null).asInstanceOf[AtomicInteger]
+ calls.set(0)
+ failureField.set(null, new IOException("empty scan policy source unavailable"))
Review Comment:
Could this test avoid reflecting into the package-private statics of `TestCometS3LocationScopedCredentialProvider`? `MinioLocationScopedCredentialProvider` in the same test package already has public `resetCounters()` and `locationCallCount()`. From reading the code I expect that asserting `locationCallCount() == 0` after the scan catches the regression too, because the policy locations are requested when `load_file_io` builds the `FileIO` (`iceberg_common.rs:473`). That would drop both `getDeclaredField` calls, the `setAccessible` calls and the injected `IOException`. I have not run it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814#discussion_r4237757953
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala:
##########
@@ -1283,14 +1283,19 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit
// Spark's BatchScanExec.inputRDD returns sparkContext.parallelize(empty, 1) when
// DPP filtering removes all input partitions. That ParallelCollectionRDD is the only
Review Comment:
Since this comment is being rewritten anyway, could the first sentence say exactly when Spark takes this path? `BatchScanExec.inputRDD` only returns `parallelize(empty, 1)` when DPP removes every split of a scan that reports `SinglePartition`. Only Spark 3.4 reports that, and only for a scan with one input partition. On 3.5+, and on 3.4 with more than one split, a fully pruned scan comes back as a `DataSourceRDD` with no partitions and goes through the branch above. Wording like "when DPP removes the only split of a scan that reports SinglePartition, which only Spark 3.4 does" would make it clear this branch is the 3.4 case.
##########
spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala:
##########
@@ -5150,6 +5152,144 @@ class CometIcebergNativeSuite
}
}
+ Seq(false, true).foreach { aqeEnabled =>
Review Comment:
This also fixes a second symptom of #6667 that the tests don't cover yet. On Spark 3.4 a sort-merge join between two single-split Iceberg tables has no exchange either, since both scans report `SinglePartition`, so the two scans run in one native block. When DPP prunes one side, main fails the query instead of returning wrong rows. With the pruned scan as the first input it fails with `All per-partition arrays must have length 0`, and as the second input with `Missing planning data for key`. I checked locally on 3.4.3 with AQE on and off, and this PR fixes both. Could you add a test for that shape next to this one? A second single-file table partitioned by `store_id` and a query like the one below would do it, run once with the DPP key on `a.store_id` and once on `b.store_id`.
```sql
SELECT /*+ BROADCAST(d), MERGE(a, b) */ COUNT(*), SUM(b.v)
FROM fact a JOIN fact2 b ON a.id = b.id
JOIN single_split_dim d ON a.store_id = d.store_id
WHERE d.country = 'X'
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "LinSimon-901101 (via GitHub)" <gi...@apache.org>.
LinSimon-901101 commented on PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814#issuecomment-6085839229
@comphead ,Thanks for the suggestion. I've switched to `MinioLocationScopedCredentialProvider` and removed the reflection and exception injection.
The full `CometIcebergNativeSuite` passed on Spark 3.4 in [[fork CI](https://github.com/LinSimon-901101/datafusion-comet/actions/runs/37963091003)], including all three regression tests. I also confirmed the revised test fails without the native guard.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "LinSimon-901101 (via GitHub)" <gi...@apache.org>.
LinSimon-901101 commented on PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814#issuecomment-6099405369
And here's one more question: an earlier CI run failed due to a Maven dependency download error.
If this happens again, should I tag a maintainer to request a rerun, or is there a preferred process?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: [spark34] preserve empty Iceberg partitions after dynamic pruning [datafusion-comet]
Posted by "LinSimon-901101 (via GitHub)" <gi...@apache.org>.
LinSimon-901101 commented on PR #6814:
URL: https://github.com/apache/datafusion-comet/pull/6814#issuecomment-6099401078
@andygrove, thanks for the detailed review! I've added four join cases (both inputs, AQE on/off) and clarified the comment. The full suite passes on Spark 3.4 and 4.1 in [[fork CI](https://github.com/LinSimon-901101/datafusion-comet/actions/runs/38059448639)]. Reverting the serde fix reproduces both errors.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org