You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/25 22:58:16 UTC
(datafusion-comet) 02/02: fix: inject Comet session extension rules only once per session (#6230)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 36ab57c689d4551b060637ec311f0e8da3b96a62
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Sep 25 20:19:21 2026 +0000
fix: inject Comet session extension rules only once per session (#6230)
* fix: inject Comet session extension rules only once per session
A session can receive CometSparkSessionExtensions more than once, for
example through both spark.sql.extensions and
SparkSession.Builder.withExtensions. Each application injected another
copy of every Comet rule, so the rules ran twice per plan. Skip the
injection when the extension has already been applied to the same
SparkSessionExtensions instance.
* ci: register CometSessionExtensionsSuite in the PR build workflows
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../apache/comet/CometSparkSessionExtensions.scala | 16 ++++++
.../sql/comet/CometSessionExtensionsSuite.scala | 57 ++++++++++++++++++++++
4 files changed, 75 insertions(+)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 2317106337..ea8a23d204 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -539,6 +539,7 @@ jobs:
org.apache.comet.QueryContextInternerSuite
org.apache.comet.CometSetOpWithGroupBySuite
org.apache.comet.CometSparkSessionExtensionsSuite
+ org.apache.spark.sql.comet.CometSessionExtensionsSuite
org.apache.spark.CometPluginsSuite
org.apache.spark.CometRuntimeShutdownSuite
org.apache.spark.CometTaskMemoryManagerSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 85ff13dfca..3a351d1e6f 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -187,6 +187,7 @@ jobs:
org.apache.comet.QueryContextInternerSuite
org.apache.comet.CometSetOpWithGroupBySuite
org.apache.comet.CometSparkSessionExtensionsSuite
+ org.apache.spark.sql.comet.CometSessionExtensionsSuite
org.apache.spark.CometPluginsSuite
org.apache.spark.CometRuntimeShutdownSuite
org.apache.spark.CometTaskMemoryManagerSuite
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 8959765111..782fe2beab 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -90,6 +90,13 @@ class CometSparkSessionExtensions
with Logging
with ShimCometSparkSessionExtensions {
override def apply(extensions: SparkSessionExtensions): Unit = {
+ // A session can be handed this extension more than once, for example through both
+ // spark.sql.extensions and SparkSession.Builder.withExtensions. Injecting twice would run
+ // every Comet rule twice per plan.
+ if (!CometSparkSessionExtensions.markConfigured(extensions)) {
+ logDebug("Comet extension already applied to these session extensions; skipping")
+ return
+ }
extensions.injectColumnar { session => CometColumnar(session) }
// Pre-3.5 only: tag AQE DPP regions so the conversion rules below leave them Spark-native.
// Registered before CometRule so tags are in place when conversion runs.
@@ -117,6 +124,15 @@ object CometSparkSessionExtensions extends Logging {
lazy val isBigEndian: Boolean = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN)
private val SHUFFLE_MANAGER_KEY = "spark.shuffle.manager"
+ /** Session extensions Comet has already been injected into. Weak so sessions can be GC'd. */
+ private val configuredExtensions =
+ java.util.Collections.synchronizedMap(
+ new java.util.WeakHashMap[SparkSessionExtensions, java.lang.Boolean]())
+
+ /** Records that Comet is being injected into `extensions`; false if it already was. */
+ private def markConfigured(extensions: SparkSessionExtensions): Boolean =
+ configuredExtensions.put(extensions, java.lang.Boolean.TRUE) == null
+
/**
* Checks whether Comet extension should be loaded for Spark.
*/
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometSessionExtensionsSuite.scala
new file mode 100644
index 0000000000..49a2f83b6e
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometSessionExtensionsSuite.scala
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import org.apache.spark.sql.{CometTestBase, SparkSessionExtensions}
+
+import org.apache.comet.CometSparkSessionExtensions
+import org.apache.comet.rules.CometRule
+
+class CometSessionExtensionsSuite extends CometTestBase {
+
+ private def countCometRules(extensions: SparkSessionExtensions): (Int, Int) = {
+ val columnar = extensions
+ .buildColumnarRules(spark)
+ .count(_.isInstanceOf[CometSparkSessionExtensions#CometColumnar])
+ val queryStagePrep = extensions
+ .buildQueryStagePrepRules(spark)
+ .count(_.isInstanceOf[CometRule])
+ (columnar, queryStagePrep)
+ }
+
+ test("Comet rules are injected once when the extension is applied twice") {
+ val extensions = new SparkSessionExtensions
+ // Separate instances, as when the extension is registered both through
+ // spark.sql.extensions and SparkSession.Builder.withExtensions.
+ new CometSparkSessionExtensions().apply(extensions)
+ new CometSparkSessionExtensions().apply(extensions)
+ assert(countCometRules(extensions) == ((1, 1)))
+ }
+
+ test("each session extensions instance gets its own Comet rules") {
+ val first = new SparkSessionExtensions
+ val second = new SparkSessionExtensions
+ val comet = new CometSparkSessionExtensions()
+ comet.apply(first)
+ comet.apply(second)
+ assert(countCometRules(first) == ((1, 1)))
+ assert(countCometRules(second) == ((1, 1)))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org