andygrove opened a new pull request, #6230: URL: https://github.com/apache/datafusion-comet/pull/6230
## Which issue does this PR close? N/A ## Rationale for this change A `SparkSession` can be handed `CometSparkSessionExtensions` more than once, for example when it is listed in `spark.sql.extensions` and also passed to `SparkSession.Builder.withExtensions`. Each application injects another copy of every Comet rule: `CometColumnar`, the query-stage-prep `CometRule`, the AQE query-stage optimizer rules and the Iceberg write strategy. As a result every Comet rule runs twice per plan. The second pass is mostly redundant work, but it is not always harmless. For example, with `spark.comet.explain.planOnly.enabled`, the second `CometRule` sees a plan the first already tagged as reported, treats it as a re-application, and clears the per-query subquery dedup set. A subquery referenced twice is then reported twice. ## What changes are included in this PR? `CometSparkSessionExtensions.apply` records each `SparkSessionExtensions` instance it has configured, in a synchronized weak map so sessions can still be garbage collected. It does nothing on a repeat call for the same instance. Distinct `SparkSessionExtensions` instances, and therefore distinct sessions, are still configured independently. ## How are these changes tested? New `CometSessionExtensionsSuite`: - Applying two `CometSparkSessionExtensions` instances to the same `SparkSessionExtensions` yields one `CometColumnar` and one query-stage-prep `CometRule`. Without the fix this fails with `(2, 2) did not equal (1, 1)`. - Applying one instance to two different `SparkSessionExtensions` configures both. Ran locally on the `spark-3.4` and `spark-3.5` profiles, along with the existing `CometSparkSessionExtensionsSuite` and the `CometExecRuleSuite` plan-only tests. -- 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]
