From f34dee18a50764da842abd7dc52f47349f2d2814 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 14:47:34 -0700
Subject: [PATCH 1/6] chore: upgrade Spark 4.x support from 4.0.1 to 4.1.0
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Update the spark-4.0 profile to use Spark 4.1.0 and add shims for
API changes:

- Sum.evalMode -> Sum.evalContext.evalMode (SPARK-53968)
- BinaryOutputStyle config now returns enum directly
- MapStatus.apply() requires new checksumVal parameter (SPARK-51756)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 .github/workflows/spark_sql_test.yml          |  4 +--
 dev/release/create-tarball.sh                 |  2 +-
 docs/source/user-guide/latest/installation.md |  2 +-
 pom.xml                                       |  2 +-
 .../CometBypassMergeSortShuffleWriter.java    |  7 +++--
 .../shuffle/CometUnsafeShuffleWriter.java     |  5 ++--
 .../org/apache/comet/serde/aggregates.scala   |  7 +++--
 .../org/apache/comet/shims/CometAggShim.scala | 27 +++++++++++++++++
 .../spark/sql/comet/shims/ShimMapStatus.scala | 29 +++++++++++++++++++
 .../org/apache/comet/shims/CometAggShim.scala | 27 +++++++++++++++++
 .../apache/comet/shims/CometExprShim.scala    |  4 +--
 .../spark/sql/comet/shims/ShimMapStatus.scala | 29 +++++++++++++++++++
 12 files changed, 129 insertions(+), 16 deletions(-)
 create mode 100644 spark/src/main/spark-3.x/org/apache/comet/shims/CometAggShim.scala
 create mode 100644 spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
 create mode 100644 spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala
 create mode 100644 spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala

diff --git a/.github/workflows/spark_sql_test.yml b/.github/workflows/spark_sql_test.yml
index 3d7aa2e2f9..09b85515ea 100644
--- a/.github/workflows/spark_sql_test.yml
+++ b/.github/workflows/spark_sql_test.yml
@@ -50,7 +50,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-24.04]
-        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.0', full: '4.0.1', java: 17}]
+        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.0', full: '4.1.0', java: 17}]
         module:
           - {name: "catalyst", args1: "catalyst/test", args2: ""}
           - {name: "sql_core-1", args1: "", args2: sql/testOnly * -- -l org.apache.spark.tags.ExtendedSQLTest -l org.apache.spark.tags.SlowSQLTest}
@@ -61,7 +61,7 @@ jobs:
           - {name: "sql_hive-3", args1: "", args2: "hive/testOnly * -- -n org.apache.spark.tags.SlowHiveTest"}
         # Skip sql_hive-1 for Spark 4.0 due to https://github.com/apache/datafusion-comet/issues/2946
         exclude:
-          - spark-version: {short: '4.0', full: '4.0.1', java: 17}
+          - spark-version: {short: '4.0', full: '4.1.0', java: 17}
             module: {name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest"}
       fail-fast: false
     name: spark-sql-${{ matrix.module.name }}/${{ matrix.os }}/spark-${{ matrix.spark-version.full }}/java-${{ matrix.spark-version.java }}
diff --git a/dev/release/create-tarball.sh b/dev/release/create-tarball.sh
index 6387cf485a..8fe5af0607 100755
--- a/dev/release/create-tarball.sh
+++ b/dev/release/create-tarball.sh
@@ -110,7 +110,7 @@ echo "---------------------------------------------------------"
 
 
 # create <tarball> containing the files in git at $release_hash
-# the files in the tarball are prefixed with {version} (e.g. 4.0.1)
+# the files in the tarball are prefixed with {version} (e.g. 4.1.0)
 mkdir -p ${distdir}
 (cd "${DEV_RELEASE_TOP_DIR}" && git archive ${release_hash} --prefix ${release}/ | gzip > ${tarball})
 
diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md
index d74fee36b8..c250c98aed 100644
--- a/docs/source/user-guide/latest/installation.md
+++ b/docs/source/user-guide/latest/installation.md
@@ -49,7 +49,7 @@ use only and should not be used in production yet.
 
 | Spark Version | Java Version | Scala Version | Comet Tests in CI | Spark SQL Tests in CI |
 | ------------- | ------------ | ------------- | ----------------- | --------------------- |
-| 4.0.1         | 17           | 2.13          | Yes               | Yes                   |
+| 4.1.0         | 17           | 2.13          | Yes               | Yes                   |
 
 Note that Comet may not fully work with proprietary forks of Apache Spark such as the Spark versions offered by
 Cloud Service Providers.
diff --git a/pom.xml b/pom.xml
index ab42aa773e..4e7848ae99 100644
--- a/pom.xml
+++ b/pom.xml
@@ -637,7 +637,7 @@ under the License.
         <!-- Use Scala 2.13 by default -->
         <scala.version>2.13.16</scala.version>
         <scala.binary.version>2.13</scala.binary.version>
-        <spark.version>4.0.1</spark.version>
+        <spark.version>4.1.0</spark.version>
         <spark.version.short>4.0</spark.version.short>
         <parquet.version>1.15.2</parquet.version>
         <semanticdb.version>4.13.6</semanticdb.version>
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
index a58ec7851b..6dc3fd6d7c 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
@@ -41,7 +41,6 @@
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper;
 import org.apache.spark.scheduler.MapStatus;
-import org.apache.spark.scheduler.MapStatus$;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
 import org.apache.spark.shuffle.ShuffleWriter;
@@ -54,6 +53,7 @@
 import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocatorTrait;
 import org.apache.spark.shuffle.sort.CometShuffleExternalSorter;
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.comet.shims.ShimMapStatus$;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.storage.FileSegment;
@@ -172,7 +172,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
                 .commitAllPartitions(ShuffleChecksumHelper.EMPTY_CHECKSUM_VALUE)
                 .getPartitionLengths();
         mapStatus =
-            MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+            ShimMapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
         return;
       }
       final long openStartTime = System.nanoTime();
@@ -261,7 +261,8 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
 
       // TODO: We probably can move checksum generation here when concatenating partition files
       partitionLengths = writePartitionedData(mapOutputWriter);
-      mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+      mapStatus =
+          ShimMapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
     } catch (Exception e) {
       try {
         mapOutputWriter.abort(e);
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
index a845e743d4..d736512a86 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
@@ -50,7 +50,6 @@
 import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper;
 import org.apache.spark.network.util.LimitedInputStream;
 import org.apache.spark.scheduler.MapStatus;
-import org.apache.spark.scheduler.MapStatus$;
 import org.apache.spark.serializer.SerializationStream;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.shuffle.BaseShuffleHandle;
@@ -67,6 +66,7 @@
 import org.apache.spark.shuffle.sort.SortShuffleManager;
 import org.apache.spark.shuffle.sort.UnsafeShuffleWriter;
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.comet.shims.ShimMapStatus$;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.storage.TimeTrackingOutputStream;
@@ -288,7 +288,8 @@ void closeAndWriteOutput() throws IOException {
         }
       }
     }
-    mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+    mapStatus =
+        ShimMapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
   }
 
   @VisibleForTesting
diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
index 8ab568dc83..7e3e978d8f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala
@@ -30,7 +30,7 @@ import org.apache.comet.CometConf
 import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType}
-import org.apache.comet.shims.CometEvalModeUtil
+import org.apache.comet.shims.{CometAggShim, CometEvalModeUtil}
 
 object CometMin extends CometAggregateExpressionSerde[Min] {
 
@@ -214,7 +214,7 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
 object CometSum extends CometAggregateExpressionSerde[Sum] {
 
   override def getSupportLevel(sum: Sum): SupportLevel = {
-    sum.evalMode match {
+    CometAggShim.getSumEvalMode(sum) match {
       case EvalMode.ANSI if !sum.dataType.isInstanceOf[DecimalType] =>
         Incompatible(Some("ANSI mode for non decimal inputs is not supported"))
       case EvalMode.TRY if !sum.dataType.isInstanceOf[DecimalType] =>
@@ -243,7 +243,8 @@ object CometSum extends CometAggregateExpressionSerde[Sum] {
       val builder = ExprOuterClass.Sum.newBuilder()
       builder.setChild(childExpr.get)
       builder.setDatatype(dataType.get)
-      builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(sum.evalMode)))
+      val evalMode = CometEvalModeUtil.fromSparkEvalMode(CometAggShim.getSumEvalMode(sum))
+      builder.setEvalMode(evalModeToProto(evalMode))
 
       Some(
         ExprOuterClass.AggExpr
diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/CometAggShim.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/CometAggShim.scala
new file mode 100644
index 0000000000..2c311a6373
--- /dev/null
+++ b/spark/src/main/spark-3.x/org/apache/comet/shims/CometAggShim.scala
@@ -0,0 +1,27 @@
+/*
+ * 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.comet.shims
+
+import org.apache.spark.sql.catalyst.expressions.EvalMode
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+
+object CometAggShim {
+  def getSumEvalMode(sum: Sum): EvalMode.Value = sum.evalMode
+}
diff --git a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimMapStatus.scala b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
new file mode 100644
index 0000000000..280a086602
--- /dev/null
+++ b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
@@ -0,0 +1,29 @@
+/*
+ * 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.shims
+
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.storage.BlockManagerId
+
+object ShimMapStatus {
+  def apply(loc: BlockManagerId, uncompressedSizes: Array[Long], mapTaskId: Long): MapStatus = {
+    MapStatus(loc, uncompressedSizes, mapTaskId)
+  }
+}
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala
new file mode 100644
index 0000000000..d34915389c
--- /dev/null
+++ b/spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala
@@ -0,0 +1,27 @@
+/*
+ * 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.comet.shims
+
+import org.apache.spark.sql.catalyst.expressions.EvalMode
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+
+object CometAggShim {
+  def getSumEvalMode(sum: Sum): EvalMode.Value = sum.evalContext.evalMode
+}
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
index fc3db183b3..21d18f1e1e 100644
--- a/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
+++ b/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
@@ -39,9 +39,7 @@ trait CometExprShim extends CommonStringExprs {
     CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 
   protected def binaryOutputStyle: BinaryOutputStyle = {
-    SQLConf.get
-      .getConf(SQLConf.BINARY_OUTPUT_STYLE)
-      .map(SQLConf.BinaryOutputStyle.withName) match {
+    SQLConf.get.getConf(SQLConf.BINARY_OUTPUT_STYLE) match {
       case Some(SQLConf.BinaryOutputStyle.UTF8) => BinaryOutputStyle.UTF8
       case Some(SQLConf.BinaryOutputStyle.BASIC) => BinaryOutputStyle.BASIC
       case Some(SQLConf.BinaryOutputStyle.BASE64) => BinaryOutputStyle.BASE64
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
new file mode 100644
index 0000000000..f805a9b542
--- /dev/null
+++ b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
@@ -0,0 +1,29 @@
+/*
+ * 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.shims
+
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.storage.BlockManagerId
+
+object ShimMapStatus {
+  def apply(loc: BlockManagerId, uncompressedSizes: Array[Long], mapTaskId: Long): MapStatus = {
+    MapStatus(loc, uncompressedSizes, mapTaskId, 0L)
+  }
+}

From 118ca281158066e9867ee7de3b464f19a914a053 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 14:53:21 -0700
Subject: [PATCH 2/6] chore: rename spark-4.0 profile and shims directory to
 spark-4.1
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 pom.xml                                                     | 6 +++---
 spark/pom.xml                                               | 4 ++--
 .../org/apache/comet/shims/CometAggShim.scala               | 0
 .../org/apache/comet/shims/CometExprShim.scala              | 0
 .../apache/comet/shims/ShimCometBroadcastExchangeExec.scala | 0
 .../apache/comet/shims/ShimCometShuffleExchangeExec.scala   | 0
 .../comet/shims/ShimCometSparkSessionExtensions.scala       | 0
 .../org/apache/comet/shims/ShimSQLConf.scala                | 0
 .../apache/spark/comet/shims/ShimCometDriverPlugin.scala    | 0
 .../apache/spark/sql/comet/shims/ShimCometScanExec.scala    | 0
 .../sql/comet/shims/ShimCometShuffleWriteProcessor.scala    | 0
 .../org/apache/spark/sql/comet/shims/ShimMapStatus.scala    | 0
 .../sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala    | 0
 13 files changed, 5 insertions(+), 5 deletions(-)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/CometAggShim.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/CometExprShim.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/ShimCometBroadcastExchangeExec.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/ShimCometShuffleExchangeExec.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/comet/shims/ShimSQLConf.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/spark/comet/shims/ShimCometDriverPlugin.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/spark/sql/comet/shims/ShimCometScanExec.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/spark/sql/comet/shims/ShimCometShuffleWriteProcessor.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/spark/sql/comet/shims/ShimMapStatus.scala (100%)
 rename spark/src/main/{spark-4.0 => spark-4.1}/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala (100%)

diff --git a/pom.xml b/pom.xml
index 4e7848ae99..91c64e6f01 100644
--- a/pom.xml
+++ b/pom.xml
@@ -632,17 +632,17 @@ under the License.
 
     <profile>
       <!-- FIXME: this is WIP. Tests may fail https://github.com/apache/datafusion-comet/issues/551 -->
-      <id>spark-4.0</id>
+      <id>spark-4.1</id>
       <properties>
         <!-- Use Scala 2.13 by default -->
         <scala.version>2.13.16</scala.version>
         <scala.binary.version>2.13</scala.binary.version>
         <spark.version>4.1.0</spark.version>
-        <spark.version.short>4.0</spark.version.short>
+        <spark.version.short>4.1</spark.version.short>
         <parquet.version>1.15.2</parquet.version>
         <semanticdb.version>4.13.6</semanticdb.version>
         <slf4j.version>2.0.16</slf4j.version>
-        <shims.majorVerSrc>spark-4.0</shims.majorVerSrc>
+        <shims.majorVerSrc>spark-4.1</shims.majorVerSrc>
         <shims.minorVerSrc>not-needed-yet</shims.minorVerSrc>
         <!-- Use jdk17 by default -->
         <java.version>17</java.version>
diff --git a/spark/pom.xml b/spark/pom.xml
index 3b832e37a2..13f7772fe6 100644
--- a/spark/pom.xml
+++ b/spark/pom.xml
@@ -233,7 +233,7 @@ under the License.
     </profile>
 
     <profile>
-      <id>spark-4.0</id>
+      <id>spark-4.1</id>
       <dependencies>
         <dependency>
           <groupId>org.apache.iceberg</groupId>
@@ -241,7 +241,7 @@ under the License.
           <version>1.10.0</version>
           <scope>test</scope>
         </dependency>
-        <!-- Jetty 11.x for Spark 4.0 (jakarta.servlet) -->
+        <!-- Jetty 11.x for Spark 4.1 (jakarta.servlet) -->
         <dependency>
           <groupId>org.eclipse.jetty</groupId>
           <artifactId>jetty-server</artifactId>
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/CometAggShim.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/CometAggShim.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/CometAggShim.scala
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/CometExprShim.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/CometExprShim.scala
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometBroadcastExchangeExec.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometBroadcastExchangeExec.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometBroadcastExchangeExec.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometBroadcastExchangeExec.scala
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometShuffleExchangeExec.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometShuffleExchangeExec.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometShuffleExchangeExec.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometShuffleExchangeExec.scala
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimSQLConf.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimSQLConf.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/comet/shims/ShimSQLConf.scala
rename to spark/src/main/spark-4.1/org/apache/comet/shims/ShimSQLConf.scala
diff --git a/spark/src/main/spark-4.0/org/apache/spark/comet/shims/ShimCometDriverPlugin.scala b/spark/src/main/spark-4.1/org/apache/spark/comet/shims/ShimCometDriverPlugin.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/spark/comet/shims/ShimCometDriverPlugin.scala
rename to spark/src/main/spark-4.1/org/apache/spark/comet/shims/ShimCometDriverPlugin.scala
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometScanExec.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometScanExec.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometScanExec.scala
rename to spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometScanExec.scala
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometShuffleWriteProcessor.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometShuffleWriteProcessor.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometShuffleWriteProcessor.scala
rename to spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometShuffleWriteProcessor.scala
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
rename to spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimMapStatus.scala
diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala
similarity index 100%
rename from spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala
rename to spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimStreamSourceAwareSparkPlan.scala

From fdca3161a71727b6bd62bd60f2ce20fc950f1be0 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 14:54:50 -0700
Subject: [PATCH 3/6] chore: update spark short version 4.0 to 4.1 in
 spark_sql_test workflow
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 .github/workflows/spark_sql_test.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/spark_sql_test.yml b/.github/workflows/spark_sql_test.yml
index 09b85515ea..cdfb233e97 100644
--- a/.github/workflows/spark_sql_test.yml
+++ b/.github/workflows/spark_sql_test.yml
@@ -50,7 +50,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-24.04]
-        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.0', full: '4.1.0', java: 17}]
+        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.1', full: '4.1.0', java: 17}]
         module:
           - {name: "catalyst", args1: "catalyst/test", args2: ""}
           - {name: "sql_core-1", args1: "", args2: sql/testOnly * -- -l org.apache.spark.tags.ExtendedSQLTest -l org.apache.spark.tags.SlowSQLTest}
@@ -59,9 +59,9 @@ jobs:
           - {name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest"}
           - {name: "sql_hive-2", args1: "", args2: "hive/testOnly * -- -n org.apache.spark.tags.ExtendedHiveTest"}
           - {name: "sql_hive-3", args1: "", args2: "hive/testOnly * -- -n org.apache.spark.tags.SlowHiveTest"}
-        # Skip sql_hive-1 for Spark 4.0 due to https://github.com/apache/datafusion-comet/issues/2946
+        # Skip sql_hive-1 for Spark 4.1 due to https://github.com/apache/datafusion-comet/issues/2946
         exclude:
-          - spark-version: {short: '4.0', full: '4.1.0', java: 17}
+          - spark-version: {short: '4.1', full: '4.1.0', java: 17}
             module: {name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest"}
       fail-fast: false
     name: spark-sql-${{ matrix.module.name }}/${{ matrix.os }}/spark-${{ matrix.spark-version.full }}/java-${{ matrix.spark-version.java }}

From c4e7bbe3dab6767c18943cad3ec6f9b5a1db18e3 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 16:01:45 -0700
Subject: [PATCH 4/6] fix

---
 spark/src/test/scala/org/apache/comet/CometCastSuite.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
index a7bd6febf8..c20b4ef824 100644
--- a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
@@ -31,9 +31,9 @@ import org.apache.spark.sql.catalyst.expressions.Cast
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DataTypes, DecimalType, IntegerType, LongType, ShortType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DataTypes, DecimalType, IntegerType, LongType, ShortType, StringType, StructField, StructType}import org.apache.comet.expressions.{CometCast, CometEvalMode}
 
-import org.apache.comet.expressions.{CometCast, CometEvalMode}
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.rules.CometScanTypeChecker
 import org.apache.comet.serde.Compatible
 

From d1be4feee4bdeac3549016555f5006cc486538a0 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 16:05:09 -0700
Subject: [PATCH 5/6] rename diff

---
 dev/diffs/{4.0.1.diff => 4.1.0.diff} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename dev/diffs/{4.0.1.diff => 4.1.0.diff} (100%)

diff --git a/dev/diffs/4.0.1.diff b/dev/diffs/4.1.0.diff
similarity index 100%
rename from dev/diffs/4.0.1.diff
rename to dev/diffs/4.1.0.diff

From e1d5297038a79997ce05ace0a0f3b82561c3fcd3 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 22 Dec 2025 16:20:31 -0700
Subject: [PATCH 6/6] fix diff

---
 dev/diffs/4.1.0.diff | 379 +++++++++++++++++++++----------------------
 1 file changed, 184 insertions(+), 195 deletions(-)

diff --git a/dev/diffs/4.1.0.diff b/dev/diffs/4.1.0.diff
index a9315db005..241c864683 100644
--- a/dev/diffs/4.1.0.diff
+++ b/dev/diffs/4.1.0.diff
@@ -1,8 +1,8 @@
 diff --git a/pom.xml b/pom.xml
-index 22922143fc3..477d4ec4194 100644
+index 1824a28614b..0d0ce6d27ce 100644
 --- a/pom.xml
 +++ b/pom.xml
-@@ -148,6 +148,8 @@
+@@ -152,6 +152,8 @@
      <kryo.version>4.0.3</kryo.version>
      <ivy.version>2.5.3</ivy.version>
      <oro.version>2.0.8</oro.version>
@@ -11,7 +11,7 @@ index 22922143fc3..477d4ec4194 100644
      <!--
      If you change codahale.metrics.version, you also need to change
      the link to metrics.dropwizard.io in docs/monitoring.md.
-@@ -2596,6 +2598,25 @@
+@@ -2594,6 +2596,25 @@
          <artifactId>arpack</artifactId>
          <version>${netlib.ludovic.dev.version}</version>
        </dependency>
@@ -38,10 +38,10 @@ index 22922143fc3..477d4ec4194 100644
        <dependency>
          <groupId>org.apache.datasketches</groupId>
 diff --git a/sql/core/pom.xml b/sql/core/pom.xml
-index dcf6223a98b..0458a5bb640 100644
+index e3aa8f1f3bb..8817d48b272 100644
 --- a/sql/core/pom.xml
 +++ b/sql/core/pom.xml
-@@ -90,6 +90,10 @@
+@@ -97,6 +97,10 @@
        <groupId>org.apache.spark</groupId>
        <artifactId>spark-tags_${scala.binary.version}</artifactId>
      </dependency>
@@ -53,10 +53,10 @@ index dcf6223a98b..0458a5bb640 100644
      <!--
        This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
-index 0015d7ff99e..c9dd85e72c4 100644
+index c47e8454162..f553cbe78ec 100644
 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
 +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
-@@ -1042,6 +1042,23 @@ object SparkSession extends SparkSessionCompanion with Logging {
+@@ -1218,6 +1218,23 @@ object SparkSession extends SparkSessionCompanion with Logging {
      extensions
    }
  
@@ -80,7 +80,7 @@ index 0015d7ff99e..c9dd85e72c4 100644
    /**
     * Initialize extensions specified in [[StaticSQLConf]]. The classes will be applied to the
     * extensions passed into this function.
-@@ -1051,7 +1068,8 @@ object SparkSession extends SparkSessionCompanion with Logging {
+@@ -1227,7 +1244,8 @@ object SparkSession extends SparkSessionCompanion with Logging {
        extensions: SparkSessionExtensions): SparkSessionExtensions = {
      val extensionConfClassNames = sparkContext.conf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS)
        .getOrElse(Seq.empty)
@@ -111,11 +111,11 @@ index 4410fe50912..43bcce2a038 100644
      }
      val childrenInfo = children.flatMap {
 diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/listagg-collations.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/listagg-collations.sql.out
-index 7aca17dcb25..8afeb3b4a2f 100644
+index 7a7d7aa3f04..05aec429d7e 100644
 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/listagg-collations.sql.out
 +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/listagg-collations.sql.out
 @@ -64,15 +64,6 @@ WithCTE
-       +- CTERelationRef xxxx, true, [c1#x], false, false
+       +- CTERelationRef xxxx, true, [c1#x], false, false, 1
  
  
 --- !query
@@ -130,17 +130,6 @@ index 7aca17dcb25..8afeb3b4a2f 100644
  -- !query
  WITH t(c1) AS (SELECT replace(listagg(DISTINCT col1 COLLATE unicode_rtrim) COLLATE utf8_binary, ' ', '') FROM (VALUES ('xbc  '), ('xbc '), ('a'), ('xbc'))) SELECT len(c1), regexp_count(c1, 'a'), regexp_count(c1, 'xbc') FROM t
  -- !query analysis
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/collations.sql b/sql/core/src/test/resources/sql-tests/inputs/collations.sql
-index 17815ed5dde..baad440b1ce 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/collations.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/collations.sql
-@@ -1,3 +1,6 @@
-+-- TODO: https://github.com/apache/datafusion-comet/issues/551
-+--SET spark.comet.enabled = false
-+
- -- test cases for collation support
- 
- -- Create a test table with data
 diff --git a/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql b/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql
 index 13bbd9d81b7..541cdfb1e04 100644
 --- a/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql
@@ -179,7 +168,7 @@ index eeb2180f7a5..afd1b5ec289 100644
  CREATE TABLE explain_temp1(a INT, b INT) USING PARQUET;
  CREATE TABLE explain_temp2(c INT, d INT) USING PARQUET;
 diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain.sql b/sql/core/src/test/resources/sql-tests/inputs/explain.sql
-index 698ca009b4f..57d774a3617 100644
+index 96dddafd82a..efedbbb3c0e 100644
 --- a/sql/core/src/test/resources/sql-tests/inputs/explain.sql
 +++ b/sql/core/src/test/resources/sql-tests/inputs/explain.sql
 @@ -1,6 +1,7 @@
@@ -332,19 +321,19 @@ index 1f8c5822e7d..b7de4e28813 100644
  WITH t(c1) AS (SELECT replace(listagg(DISTINCT col1 COLLATE unicode_rtrim) COLLATE utf8_binary, ' ', '') FROM (VALUES ('xbc  '), ('xbc '), ('a'), ('xbc'))) SELECT len(c1), regexp_count(c1, 'a'), regexp_count(c1, 'xbc') FROM t
  -- !query schema
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
-index 0f42502f1d9..f616024a9c2 100644
+index 12d26c4e195..b529c283412 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
-@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants
- import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan, SparkPlanInfo}
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation}
+@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp
  import org.apache.spark.sql.execution.columnar._
+ import org.apache.spark.sql.execution.command.CommandUtils
+ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
  import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate
  import org.apache.spark.sql.functions._
  import org.apache.spark.sql.internal.SQLConf
-@@ -520,7 +520,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
+@@ -534,7 +534,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
        df.collect()
      }
      assert(
@@ -354,7 +343,7 @@ index 0f42502f1d9..f616024a9c2 100644
    }
  
    test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
-@@ -1626,7 +1627,9 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
+@@ -1640,7 +1641,9 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
      }
    }
  
@@ -365,7 +354,7 @@ index 0f42502f1d9..f616024a9c2 100644
      withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2",
        SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
-@@ -1661,7 +1664,12 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
+@@ -1675,7 +1678,12 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
            _.nodeName.contains("AdaptiveSparkPlan"))
          val aqePlanRoot = findNodeInSparkPlanInfo(inMemoryScanNode.get,
            _.nodeName.contains("ResultQueryStage"))
@@ -380,7 +369,7 @@ index 0f42502f1d9..f616024a9c2 100644
  
        withTempView("t0", "t1", "t2") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
-index 9db406ff12f..abbc91f5c11 100644
+index da2fbceae97..e0b4bf39ba5 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
 @@ -30,7 +30,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
@@ -401,7 +390,7 @@ index 9db406ff12f..abbc91f5c11 100644
        }
        assert(exchangePlans.length == 1)
      }
-@@ -2241,7 +2241,9 @@ class DataFrameAggregateSuite extends QueryTest
+@@ -2942,7 +2942,9 @@ class DataFrameAggregateSuite extends QueryTest
      }
    }
  
@@ -488,7 +477,7 @@ index 5b88eeefeca..d4f07bc182a 100644
      assert(exchanges.size == 2)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
-index 01e72daead4..0a8d1e8b9b9 100644
+index 6e9f3385571..15615f3e417 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
 @@ -24,8 +24,9 @@ import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression
@@ -526,7 +515,7 @@ index 01e72daead4..0a8d1e8b9b9 100644
            }
          case _ => false
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
-index 81713c777bc..b5f92ed9742 100644
+index 6df8d66ee7f..35e270c7241 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
 @@ -46,7 +46,7 @@ import org.apache.spark.sql.catalyst.trees.DataFrameQueryContext
@@ -535,10 +524,10 @@ index 81713c777bc..b5f92ed9742 100644
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
 +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
- import org.apache.spark.sql.execution.streaming.MemoryStream
+ import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
  import org.apache.spark.sql.expressions.UserDefinedFunction
  import org.apache.spark.sql.functions._
-@@ -2415,7 +2415,7 @@ class DatasetSuite extends QueryTest
+@@ -2420,7 +2420,7 @@ class DatasetSuite extends QueryTest
  
      // Assert that no extra shuffle introduced by cogroup.
      val exchanges = collect(df3.queryExecution.executedPlan) {
@@ -548,7 +537,7 @@ index 81713c777bc..b5f92ed9742 100644
      assert(exchanges.size == 2)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
-index 2c24cc7d570..21d36ebc6f5 100644
+index e1a2fd33c7c..58cd92daea6 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 @@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen
@@ -597,9 +586,9 @@ index 2c24cc7d570..21d36ebc6f5 100644
 +    "canonicalization and exchange reuse",
 +    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/1839")) {
      withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-         val df = sql(
-@@ -1424,7 +1431,8 @@ abstract class DynamicPartitionPruningSuiteBase
+       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+           SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+@@ -1425,7 +1432,8 @@ abstract class DynamicPartitionPruningSuiteBase
      }
    }
  
@@ -609,7 +598,7 @@ index 2c24cc7d570..21d36ebc6f5 100644
      withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
        val df = sql(
          """ WITH v as (
-@@ -1455,7 +1463,8 @@ abstract class DynamicPartitionPruningSuiteBase
+@@ -1456,7 +1464,8 @@ abstract class DynamicPartitionPruningSuiteBase
      }
    }
  
@@ -619,7 +608,7 @@ index 2c24cc7d570..21d36ebc6f5 100644
      val df = sql(
        """
          |SELECT s.store_id, f.product_id
-@@ -1730,6 +1739,8 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
+@@ -1731,6 +1740,8 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
                case s: BatchScanExec =>
                  // we use f1 col for v2 tables due to schema pruning
                  s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1")))
@@ -629,7 +618,7 @@ index 2c24cc7d570..21d36ebc6f5 100644
              }
            assert(scanOption.isDefined)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
-index 9c90e0105a4..fadf2f0f698 100644
+index b27122a8de2..a4c5aac8212 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
 @@ -470,7 +470,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
@@ -654,7 +643,7 @@ index 9c90e0105a4..fadf2f0f698 100644
  
    test("SPARK-35884: Explain Formatted") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
-index 9c529d14221..2f1bc3880fd 100644
+index 95e86fe4311..24f99384bb7 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
 @@ -33,6 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GreaterTha
@@ -665,7 +654,7 @@ index 9c529d14221..2f1bc3880fd 100644
  import org.apache.spark.sql.execution.{FileSourceScanLike, SimpleMode}
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.execution.datasources.FilePartition
-@@ -967,6 +968,7 @@ class FileBasedDataSourceSuite extends QueryTest
+@@ -954,6 +955,7 @@ class FileBasedDataSourceSuite extends QueryTest
              assert(bJoinExec.isEmpty)
              val smJoinExec = collect(joinedDF.queryExecution.executedPlan) {
                case smJoin: SortMergeJoinExec => smJoin
@@ -673,7 +662,7 @@ index 9c529d14221..2f1bc3880fd 100644
              }
              assert(smJoinExec.nonEmpty)
            }
-@@ -1027,6 +1029,7 @@ class FileBasedDataSourceSuite extends QueryTest
+@@ -1014,6 +1016,7 @@ class FileBasedDataSourceSuite extends QueryTest
  
            val fileScan = df.queryExecution.executedPlan collectFirst {
              case BatchScanExec(_, f: FileScan, _, _, _, _) => f
@@ -681,7 +670,7 @@ index 9c529d14221..2f1bc3880fd 100644
            }
            assert(fileScan.nonEmpty)
            assert(fileScan.get.partitionFilters.nonEmpty)
-@@ -1068,6 +1071,7 @@ class FileBasedDataSourceSuite extends QueryTest
+@@ -1055,6 +1058,7 @@ class FileBasedDataSourceSuite extends QueryTest
  
            val fileScan = df.queryExecution.executedPlan collectFirst {
              case BatchScanExec(_, f: FileScan, _, _, _, _) => f
@@ -689,7 +678,7 @@ index 9c529d14221..2f1bc3880fd 100644
            }
            assert(fileScan.nonEmpty)
            assert(fileScan.get.partitionFilters.isEmpty)
-@@ -1252,6 +1256,8 @@ class FileBasedDataSourceSuite extends QueryTest
+@@ -1239,6 +1243,8 @@ class FileBasedDataSourceSuite extends QueryTest
            val filters = df.queryExecution.executedPlan.collect {
              case f: FileSourceScanLike => f.dataFilters
              case b: BatchScanExec => b.scan.asInstanceOf[FileScan].dataFilters
@@ -802,7 +791,7 @@ index 53e47f428c3..a55d8f0c161 100644
      assert(shuffleMergeJoins.size == 1)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
-index aaac0ebc9aa..fbef0774d46 100644
+index 885512d4d19..8499592c451 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
 @@ -29,7 +29,8 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
@@ -825,7 +814,7 @@ index aaac0ebc9aa..fbef0774d46 100644
      withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1",
        SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_IN_MEMORY_THRESHOLD.key -> "0",
        SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_SPILL_THRESHOLD.key -> "1") {
-@@ -931,10 +933,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -943,10 +945,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
        val physical = df.queryExecution.sparkPlan
        val physicalJoins = physical.collect {
          case j: SortMergeJoinExec => j
@@ -838,7 +827,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        }
        // This only applies to the above tested queries, in which a child SortMergeJoin always
        // contains the SortOrder required by its parent SortMergeJoin. Thus, SortExec should never
-@@ -1180,9 +1184,11 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1192,9 +1196,11 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
        val plan = df1.join(df2.hint("SHUFFLE_HASH"), $"k1" === $"k2", joinType)
          .groupBy($"k1").count()
          .queryExecution.executedPlan
@@ -852,7 +841,7 @@ index aaac0ebc9aa..fbef0774d46 100644
      })
    }
  
-@@ -1199,10 +1205,11 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1211,10 +1217,11 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
          .join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
          .queryExecution
          .executedPlan
@@ -866,7 +855,7 @@ index aaac0ebc9aa..fbef0774d46 100644
      })
  
      // Test shuffled hash join
-@@ -1212,10 +1219,13 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1224,10 +1231,13 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
          .join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
          .queryExecution
          .executedPlan
@@ -883,7 +872,7 @@ index aaac0ebc9aa..fbef0774d46 100644
      })
    }
  
-@@ -1306,12 +1316,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1318,12 +1328,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
      inputDFs.foreach { case (df1, df2, joinExprs) =>
        val smjDF = df1.join(df2.hint("SHUFFLE_MERGE"), joinExprs, "full")
        assert(collect(smjDF.queryExecution.executedPlan) {
@@ -898,7 +887,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        // Same result between shuffled hash join and sort merge join
        checkAnswer(shjDF, smjResult)
      }
-@@ -1370,12 +1380,14 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1382,12 +1392,14 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            val smjDF = df1.hint("SHUFFLE_MERGE").join(df2, joinExprs, "leftouter")
            assert(collect(smjDF.queryExecution.executedPlan) {
              case _: SortMergeJoinExec => true
@@ -913,7 +902,7 @@ index aaac0ebc9aa..fbef0774d46 100644
            }.size === 1)
            // Same result between shuffled hash join and sort merge join
            checkAnswer(shjDF, smjResult)
-@@ -1386,12 +1398,14 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1398,12 +1410,14 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            val smjDF = df2.join(df1.hint("SHUFFLE_MERGE"), joinExprs, "rightouter")
            assert(collect(smjDF.queryExecution.executedPlan) {
              case _: SortMergeJoinExec => true
@@ -928,7 +917,7 @@ index aaac0ebc9aa..fbef0774d46 100644
            }.size === 1)
            // Same result between shuffled hash join and sort merge join
            checkAnswer(shjDF, smjResult)
-@@ -1435,13 +1449,20 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1447,13 +1461,20 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
          assert(shjCodegenDF.queryExecution.executedPlan.collect {
            case WholeStageCodegenExec(_ : ShuffledHashJoinExec) => true
            case WholeStageCodegenExec(ProjectExec(_, _ : ShuffledHashJoinExec)) => true
@@ -950,7 +939,7 @@ index aaac0ebc9aa..fbef0774d46 100644
            checkAnswer(shjNonCodegenDF, Seq.empty)
          }
        }
-@@ -1489,7 +1510,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1501,7 +1522,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            val plan = sql(getAggQuery(selectExpr, joinType)).queryExecution.executedPlan
            assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
            // Have shuffle before aggregation
@@ -960,7 +949,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        }
  
        def getJoinQuery(selectExpr: String, joinType: String): String = {
-@@ -1518,9 +1540,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1530,9 +1552,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            }
            val plan = sql(getJoinQuery(selectExpr, joinType)).queryExecution.executedPlan
            assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
@@ -975,7 +964,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        }
  
        // Test output ordering is not preserved
-@@ -1529,9 +1554,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1541,9 +1566,12 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            val selectExpr = "/*+ BROADCAST(left_t) */ k1 as k0"
            val plan = sql(getJoinQuery(selectExpr, joinType)).queryExecution.executedPlan
            assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
@@ -990,7 +979,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        }
  
        // Test singe partition
-@@ -1541,7 +1569,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1553,7 +1581,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
             |FROM range(0, 10, 1, 1) t1 FULL OUTER JOIN range(0, 10, 1, 1) t2
             |""".stripMargin)
        val plan = fullJoinDF.queryExecution.executedPlan
@@ -1000,7 +989,7 @@ index aaac0ebc9aa..fbef0774d46 100644
        checkAnswer(fullJoinDF, Row(100))
      }
    }
-@@ -1614,6 +1643,9 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1626,6 +1655,9 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
            Seq(semiJoinDF, antiJoinDF).foreach { df =>
              assert(collect(df.queryExecution.executedPlan) {
                case j: ShuffledHashJoinExec if j.ignoreDuplicatedKey == ignoreDuplicatedKey => true
@@ -1010,7 +999,7 @@ index aaac0ebc9aa..fbef0774d46 100644
              }.size == 1)
            }
        }
-@@ -1658,14 +1690,20 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
+@@ -1670,14 +1702,20 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
  
    test("SPARK-43113: Full outer join with duplicate stream-side references in condition (SMJ)") {
      def check(plan: SparkPlan): Unit = {
@@ -1033,7 +1022,7 @@ index aaac0ebc9aa..fbef0774d46 100644
      }
      dupStreamSideColTest("SHUFFLE_HASH", check)
    }
-@@ -1801,7 +1839,8 @@ class ThreadLeakInSortMergeJoinSuite
+@@ -1813,7 +1851,8 @@ class ThreadLeakInSortMergeJoinSuite
        sparkConf.set(SHUFFLE_SPILL_NUM_ELEMENTS_FORCE_SPILL_THRESHOLD, 20))
    }
  
@@ -1044,10 +1033,10 @@ index aaac0ebc9aa..fbef0774d46 100644
      withSQLConf(
        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
-index ad424b3a7cc..4ece0117a34 100644
+index e4b5e10f7c3..c6efde09c8a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
-@@ -69,7 +69,7 @@ import org.apache.spark.tags.ExtendedSQLTest
+@@ -69,7 +69,7 @@ import org.apache.spark.util.Utils
   * }}}
   */
  // scalastyle:on line.size.limit
@@ -1057,10 +1046,10 @@ index ad424b3a7cc..4ece0117a34 100644
    protected val baseResourcePath = {
      // use the same way as `SQLQueryTestSuite` to get the resource path
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
-index f294ff81021..8a3b818ee94 100644
+index 74cdee49e55..e9781798431 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
-@@ -1524,7 +1524,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
+@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
      checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001")))
    }
  
@@ -1070,7 +1059,7 @@ index f294ff81021..8a3b818ee94 100644
      AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") {
        sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
      }
-@@ -4449,7 +4450,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
+@@ -4446,7 +4447,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
    }
  
    test("SPARK-39166: Query context of binary arithmetic should be serialized to executors" +
@@ -1080,7 +1069,7 @@ index f294ff81021..8a3b818ee94 100644
      withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
-@@ -4470,7 +4472,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
+@@ -4467,7 +4469,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
    }
  
    test("SPARK-39175: Query context of Cast should be serialized to executors" +
@@ -1090,7 +1079,7 @@ index f294ff81021..8a3b818ee94 100644
      withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
-@@ -4490,14 +4493,20 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
+@@ -4487,14 +4490,20 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
            assert(ex.isInstanceOf[SparkNumberFormatException] ||
              ex.isInstanceOf[SparkDateTimeException] ||
              ex.isInstanceOf[SparkRuntimeException])
@@ -1114,7 +1103,7 @@ index f294ff81021..8a3b818ee94 100644
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
-index c1c041509c3..7d463e4b85e 100644
+index 6ee0029b683..741532fe639 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
 @@ -235,6 +235,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
@@ -1145,7 +1134,7 @@ index c1c041509c3..7d463e4b85e 100644
          extensions.injectColumnar(session =>
            MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
-index 0df7f806272..52d33d67328 100644
+index a1de322ac29..69c9d2d5d6a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
 @@ -17,6 +17,8 @@
@@ -1212,11 +1201,11 @@ index 0df7f806272..52d33d67328 100644
  
    test("non-matching optional group") {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
-index 2e33f6505ab..e1e93ab3bad 100644
+index 3ba48da0e32..2aa3da75ad4 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
 @@ -23,10 +23,11 @@ import org.apache.spark.SparkRuntimeException
- import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
+ import org.apache.spark.sql.catalyst.expressions.{EqualTo, NamedExpression, OuterReference, SubqueryExpression}
  import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi}
  import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union}
 +import org.apache.spark.sql.comet.CometScanExec
@@ -1308,11 +1297,11 @@ index fee375db10a..8c2c24e2c5f 100644
      val v = VariantBuilder.parseJson(s, false)
      new VariantVal(v.getValue, v.getMetadata)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
-index 11e9547dfc5..be9ae40ab3d 100644
+index 6cdf681d65c..4ea007422dd 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
-@@ -20,10 +20,11 @@ package org.apache.spark.sql.collation
- import scala.jdk.CollectionConverters.MapHasAsJava
+@@ -22,10 +22,11 @@ import scala.jdk.CollectionConverters.MapHasAsJava
+ import com.ibm.icu.util.VersionInfo.ICU_VERSION
  
  import org.apache.spark.SparkException
 -import org.apache.spark.sql.{AnalysisException, Row}
@@ -1322,9 +1311,9 @@ index 11e9547dfc5..be9ae40ab3d 100644
  import org.apache.spark.sql.catalyst.util.CollationFactory
 +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec}
  import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema}
- import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTable}
+ import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable}
  import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper
-@@ -55,7 +56,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
+@@ -57,7 +58,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
      assert(
        collectFirst(queryPlan) {
          case _: SortMergeJoinExec => assert(isSortMergeForced)
@@ -1334,7 +1323,7 @@ index 11e9547dfc5..be9ae40ab3d 100644
        }.nonEmpty
      )
    }
-@@ -1505,7 +1508,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
+@@ -1508,7 +1511,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
      }
    }
  
@@ -1345,7 +1334,7 @@ index 11e9547dfc5..be9ae40ab3d 100644
      val t1 = "T_1"
      val t2 = "T_2"
  
-@@ -1611,6 +1616,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
+@@ -1614,6 +1619,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
              } else {
                assert(!collectFirst(queryPlan) {
                  case b: BroadcastHashJoinExec => b.leftKeys.head
@@ -1353,7 +1342,7 @@ index 11e9547dfc5..be9ae40ab3d 100644
                }.head.isInstanceOf[ArrayTransform])
              }
            }
-@@ -1676,6 +1682,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
+@@ -1679,6 +1685,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
              } else {
                assert(!collectFirst(queryPlan) {
                  case b: BroadcastHashJoinExec => b.leftKeys.head
@@ -1361,7 +1350,7 @@ index 11e9547dfc5..be9ae40ab3d 100644
                }.head.isInstanceOf[ArrayTransform])
              }
            }
-@@ -1815,7 +1822,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
+@@ -1818,7 +1825,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
      }
    }
  
@@ -1373,27 +1362,27 @@ index 11e9547dfc5..be9ae40ab3d 100644
      val t2 = "T_2"
  
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
-index 3eeed2e4175..9f21d547c1c 100644
+index a09b7e0827c..a3c137bf659 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
-@@ -26,6 +26,7 @@ import test.org.apache.spark.sql.connector._
- import org.apache.spark.SparkUnsupportedOperationException
- import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row}
+@@ -30,6 +30,7 @@ import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row}
  import org.apache.spark.sql.catalyst.InternalRow
+ import org.apache.spark.sql.catalyst.expressions.ScalarSubquery
+ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Project}
 +import org.apache.spark.sql.comet.CometSortExec
  import org.apache.spark.sql.connector.catalog.{PartitionInternalRow, SupportsRead, Table, TableCapability, TableProvider}
  import org.apache.spark.sql.connector.catalog.TableCapability._
  import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, Literal, NamedReference, NullOrdering, SortDirection, SortOrder, Transform}
-@@ -36,7 +37,7 @@ import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning,
- import org.apache.spark.sql.execution.SortExec
+@@ -41,7 +42,7 @@ import org.apache.spark.sql.execution.SortExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation}
+ import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering}
+ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
 -import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec}
 +import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec, ShuffleExchangeLike}
  import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.functions._
-@@ -278,13 +279,13 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
+@@ -283,13 +284,13 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
            val groupByColJ = df.groupBy($"j").agg(sum($"i"))
            checkAnswer(groupByColJ, Seq(Row(2, 8), Row(4, 2), Row(6, 5)))
            assert(collectFirst(groupByColJ.queryExecution.executedPlan) {
@@ -1409,7 +1398,7 @@ index 3eeed2e4175..9f21d547c1c 100644
            }.isDefined)
          }
        }
-@@ -344,10 +345,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
+@@ -349,10 +350,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
  
                  val (shuffleExpected, sortExpected) = groupByExpects
                  assert(collectFirst(groupBy.queryExecution.executedPlan) {
@@ -1422,7 +1411,7 @@ index 3eeed2e4175..9f21d547c1c 100644
                  }.isDefined === sortExpected)
                }
  
-@@ -362,10 +364,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
+@@ -367,10 +369,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
  
                  val (shuffleExpected, sortExpected) = windowFuncExpects
                  assert(collectFirst(windowPartByColIOrderByColJ.queryExecution.executedPlan) {
@@ -1461,7 +1450,7 @@ index 2a0ab21ddb0..e8a5a891105 100644
          } finally {
            spark.listenerManager.unregister(listener)
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
-index c73e8e16fbb..88cd0d47da3 100644
+index 7c830bf6c6e..a3c5efb75e2 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
 @@ -20,10 +20,11 @@ import java.sql.Timestamp
@@ -1484,9 +1473,9 @@ index c73e8e16fbb..88cd0d47da3 100644
 -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 +import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
  import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+ import org.apache.spark.sql.functions.{col, max}
  import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.internal.SQLConf._
-@@ -305,13 +306,14 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
+@@ -306,13 +307,14 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
      }
    }
  
@@ -1503,7 +1492,7 @@ index c73e8e16fbb..88cd0d47da3 100644
        })
    }
  
-@@ -370,7 +372,9 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
+@@ -371,7 +373,9 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
      checkAnswer(df.sort("res"), Seq(Row(10.0), Row(15.5), Row(41.0)))
    }
  
@@ -1515,7 +1504,7 @@ index c73e8e16fbb..88cd0d47da3 100644
      createTable(items, itemsColumns, items_partitions)
      sql(s"INSERT INTO testcat.ns.$items VALUES " +
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
-index f62e092138a..c0404bfe85e 100644
+index 7c4852c5e22..d1a34456bdc 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
 @@ -21,7 +21,7 @@ package org.apache.spark.sql.connector
@@ -1538,7 +1527,7 @@ index f62e092138a..c0404bfe85e 100644
  
    before {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
-index 04d33ecd3d5..450df347297 100644
+index fcecaf25d4c..e5a511022cc 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
 @@ -31,7 +31,7 @@ import org.mockito.Mockito.{mock, spy, when}
@@ -1655,7 +1644,7 @@ index 743ec41dbe7..9f30d6c8e04 100644
      case _ => false
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
-index 1400ee25f43..5b016c3f9c5 100644
+index b926cc192bd..77510ccdb09 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
 @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution
@@ -1664,10 +1653,10 @@ index 1400ee25f43..5b016c3f9c5 100644
  import org.apache.spark.rdd.RDD
 -import org.apache.spark.sql.{execution, DataFrame, Row}
 +import org.apache.spark.sql.{execution, DataFrame, IgnoreCometSuite, Row}
+ import org.apache.spark.sql.AnalysisException
  import org.apache.spark.sql.catalyst.InternalRow
  import org.apache.spark.sql.catalyst.expressions._
- import org.apache.spark.sql.catalyst.plans._
-@@ -36,7 +36,9 @@ import org.apache.spark.sql.internal.SQLConf
+@@ -37,7 +37,9 @@ import org.apache.spark.sql.internal.SQLConf
  import org.apache.spark.sql.test.SharedSparkSession
  import org.apache.spark.sql.types._
  
@@ -1679,19 +1668,19 @@ index 1400ee25f43..5b016c3f9c5 100644
  
    setupTestData()
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
-index 47d5ff67b84..8dc8f65d4b1 100644
+index 0d0a0f2f310..101573a2a5f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
 @@ -20,7 +20,7 @@ import scala.collection.mutable
  import scala.io.Source
  import scala.util.Try
  
--import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator}
-+import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, IgnoreComet}
+-import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, SaveMode}
++import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, IgnoreComet, SaveMode}
  import org.apache.spark.sql.catalyst.{QueryPlanningTracker, QueryPlanningTrackerCallback, TableIdentifier}
  import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedFunction, UnresolvedRelation}
  import org.apache.spark.sql.catalyst.expressions.{Alias, UnsafeRow}
-@@ -400,7 +400,7 @@ class QueryExecutionSuite extends SharedSparkSession {
+@@ -436,7 +436,7 @@ class QueryExecutionSuite extends SharedSparkSession {
      }
    }
  
@@ -2072,10 +2061,10 @@ index a3cfdc5a240..1b08a1f42ee 100644
        })
        checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-index 272be70f9fe..06957694002 100644
+index 3e7d26f74bd..e50322898f9 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
-@@ -28,12 +28,14 @@ import org.apache.spark.SparkException
+@@ -27,12 +27,14 @@ import org.apache.spark.SparkException
  import org.apache.spark.rdd.RDD
  import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
  import org.apache.spark.shuffle.sort.SortShuffleManager
@@ -2091,7 +2080,7 @@ index 272be70f9fe..06957694002 100644
  import org.apache.spark.sql.execution._
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
  import org.apache.spark.sql.execution.columnar.{InMemoryTableScanExec, InMemoryTableScanLike}
-@@ -119,6 +121,7 @@ class AdaptiveQueryExecSuite
+@@ -121,6 +123,7 @@ class AdaptiveQueryExecSuite
    private def findTopLevelBroadcastHashJoin(plan: SparkPlan): Seq[BroadcastHashJoinExec] = {
      collect(plan) {
        case j: BroadcastHashJoinExec => j
@@ -2099,7 +2088,7 @@ index 272be70f9fe..06957694002 100644
      }
    }
  
-@@ -131,30 +134,39 @@ class AdaptiveQueryExecSuite
+@@ -133,30 +136,39 @@ class AdaptiveQueryExecSuite
    private def findTopLevelSortMergeJoin(plan: SparkPlan): Seq[SortMergeJoinExec] = {
      collect(plan) {
        case j: SortMergeJoinExec => j
@@ -2139,7 +2128,7 @@ index 272be70f9fe..06957694002 100644
      }
    }
  
-@@ -204,6 +216,7 @@ class AdaptiveQueryExecSuite
+@@ -206,6 +218,7 @@ class AdaptiveQueryExecSuite
        val parts = rdd.partitions
        assert(parts.forall(rdd.preferredLocations(_).nonEmpty))
      }
@@ -2147,7 +2136,7 @@ index 272be70f9fe..06957694002 100644
      assert(numShuffles === (numLocalReads.length + numShufflesWithoutLocalRead))
    }
  
-@@ -212,7 +225,7 @@ class AdaptiveQueryExecSuite
+@@ -214,7 +227,7 @@ class AdaptiveQueryExecSuite
      val plan = df.queryExecution.executedPlan
      assert(plan.isInstanceOf[AdaptiveSparkPlanExec])
      val shuffle = plan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan.collect {
@@ -2156,7 +2145,7 @@ index 272be70f9fe..06957694002 100644
      }
      assert(shuffle.size == 1)
      assert(shuffle(0).outputPartitioning.numPartitions == numPartition)
-@@ -228,7 +241,8 @@ class AdaptiveQueryExecSuite
+@@ -230,7 +243,8 @@ class AdaptiveQueryExecSuite
        assert(smj.size == 1)
        val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
        assert(bhj.size == 1)
@@ -2166,7 +2155,7 @@ index 272be70f9fe..06957694002 100644
      }
    }
  
-@@ -255,7 +269,8 @@ class AdaptiveQueryExecSuite
+@@ -257,7 +271,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2176,7 +2165,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
-@@ -287,7 +302,8 @@ class AdaptiveQueryExecSuite
+@@ -289,7 +304,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2186,7 +2175,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
-@@ -301,7 +317,8 @@ class AdaptiveQueryExecSuite
+@@ -303,7 +319,8 @@ class AdaptiveQueryExecSuite
        val localReads = collect(adaptivePlan) {
          case read: AQEShuffleReadExec if read.isLocalRead => read
        }
@@ -2196,7 +2185,7 @@ index 272be70f9fe..06957694002 100644
        val localShuffleRDD0 = localReads(0).execute().asInstanceOf[ShuffledRowRDD]
        val localShuffleRDD1 = localReads(1).execute().asInstanceOf[ShuffledRowRDD]
        // the final parallelism is math.max(1, numReduces / numMappers): math.max(1, 5/2) = 2
-@@ -326,7 +343,9 @@ class AdaptiveQueryExecSuite
+@@ -328,7 +345,9 @@ class AdaptiveQueryExecSuite
            .groupBy($"a").count()
          checkAnswer(testDf, Seq())
          val plan = testDf.queryExecution.executedPlan
@@ -2207,7 +2196,7 @@ index 272be70f9fe..06957694002 100644
          val coalescedReads = collect(plan) {
            case r: AQEShuffleReadExec => r
          }
-@@ -340,7 +359,9 @@ class AdaptiveQueryExecSuite
+@@ -342,7 +361,9 @@ class AdaptiveQueryExecSuite
            .groupBy($"a").count()
          checkAnswer(testDf, Seq())
          val plan = testDf.queryExecution.executedPlan
@@ -2218,7 +2207,7 @@ index 272be70f9fe..06957694002 100644
          val coalescedReads = collect(plan) {
            case r: AQEShuffleReadExec => r
          }
-@@ -350,7 +371,7 @@ class AdaptiveQueryExecSuite
+@@ -352,7 +373,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2227,7 +2216,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -365,7 +386,7 @@ class AdaptiveQueryExecSuite
+@@ -367,7 +388,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2236,7 +2225,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -381,7 +402,7 @@ class AdaptiveQueryExecSuite
+@@ -383,7 +404,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2245,7 +2234,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -426,7 +447,7 @@ class AdaptiveQueryExecSuite
+@@ -428,7 +449,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2254,7 +2243,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -471,7 +492,7 @@ class AdaptiveQueryExecSuite
+@@ -473,7 +494,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2263,7 +2252,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
-@@ -517,7 +538,7 @@ class AdaptiveQueryExecSuite
+@@ -519,7 +540,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2272,7 +2261,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -536,7 +557,7 @@ class AdaptiveQueryExecSuite
+@@ -538,7 +559,7 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2281,7 +2270,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -567,7 +588,9 @@ class AdaptiveQueryExecSuite
+@@ -569,7 +590,9 @@ class AdaptiveQueryExecSuite
        assert(smj.size == 1)
        val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
        assert(bhj.size == 1)
@@ -2292,7 +2281,7 @@ index 272be70f9fe..06957694002 100644
        // Even with local shuffle read, the query stage reuse can also work.
        val ex = findReusedExchange(adaptivePlan)
        assert(ex.nonEmpty)
-@@ -588,7 +611,9 @@ class AdaptiveQueryExecSuite
+@@ -590,7 +613,9 @@ class AdaptiveQueryExecSuite
        assert(smj.size == 1)
        val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
        assert(bhj.size == 1)
@@ -2303,7 +2292,7 @@ index 272be70f9fe..06957694002 100644
        // Even with local shuffle read, the query stage reuse can also work.
        val ex = findReusedExchange(adaptivePlan)
        assert(ex.isEmpty)
-@@ -597,7 +622,8 @@ class AdaptiveQueryExecSuite
+@@ -599,7 +624,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2313,7 +2302,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000",
-@@ -692,7 +718,8 @@ class AdaptiveQueryExecSuite
+@@ -694,7 +720,8 @@ class AdaptiveQueryExecSuite
        val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
        assert(bhj.size == 1)
        // There is still a SMJ, and its two shuffles can't apply local read.
@@ -2323,7 +2312,7 @@ index 272be70f9fe..06957694002 100644
      }
    }
  
-@@ -814,7 +841,8 @@ class AdaptiveQueryExecSuite
+@@ -822,7 +849,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2333,7 +2322,7 @@ index 272be70f9fe..06957694002 100644
      Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint =>
        def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") {
          findTopLevelSortMergeJoin(plan)
-@@ -1087,7 +1115,8 @@ class AdaptiveQueryExecSuite
+@@ -1257,7 +1285,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2343,7 +2332,7 @@ index 272be70f9fe..06957694002 100644
      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
        val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
          "SELECT key FROM testData GROUP BY key")
-@@ -1721,7 +1750,7 @@ class AdaptiveQueryExecSuite
+@@ -1891,7 +1920,7 @@ class AdaptiveQueryExecSuite
          val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
            "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id")
          assert(collect(adaptivePlan) {
@@ -2352,7 +2341,7 @@ index 272be70f9fe..06957694002 100644
          }.length == 1)
        }
      }
-@@ -1801,7 +1830,8 @@ class AdaptiveQueryExecSuite
+@@ -1972,7 +2001,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2362,7 +2351,7 @@ index 272be70f9fe..06957694002 100644
      def hasRepartitionShuffle(plan: SparkPlan): Boolean = {
        find(plan) {
          case s: ShuffleExchangeLike =>
-@@ -1986,6 +2016,9 @@ class AdaptiveQueryExecSuite
+@@ -2159,6 +2189,9 @@ class AdaptiveQueryExecSuite
      def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = {
        assert(collect(ds.queryExecution.executedPlan) {
          case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
@@ -2372,7 +2361,7 @@ index 272be70f9fe..06957694002 100644
        }.size == 1)
        ds.collect()
        val plan = ds.queryExecution.executedPlan
-@@ -1994,6 +2027,9 @@ class AdaptiveQueryExecSuite
+@@ -2167,6 +2200,9 @@ class AdaptiveQueryExecSuite
        }.isEmpty)
        assert(collect(plan) {
          case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
@@ -2382,7 +2371,7 @@ index 272be70f9fe..06957694002 100644
        }.size == 1)
        checkAnswer(ds, testData)
      }
-@@ -2150,7 +2186,8 @@ class AdaptiveQueryExecSuite
+@@ -2331,7 +2367,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2392,7 +2381,7 @@ index 272be70f9fe..06957694002 100644
      withTempView("t1", "t2") {
        def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
          Seq("100", "100000").foreach { size =>
-@@ -2236,7 +2273,8 @@ class AdaptiveQueryExecSuite
+@@ -2417,7 +2454,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2402,7 +2391,7 @@ index 272be70f9fe..06957694002 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2335,7 +2373,7 @@ class AdaptiveQueryExecSuite
+@@ -2516,7 +2554,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
              val shuffles1 = collect(adaptive1) {
@@ -2411,7 +2400,7 @@ index 272be70f9fe..06957694002 100644
              }
              assert(shuffles1.size == 3)
              // shuffles1.head is the top-level shuffle under the Aggregate operator
-@@ -2348,7 +2386,7 @@ class AdaptiveQueryExecSuite
+@@ -2529,7 +2567,7 @@ class AdaptiveQueryExecSuite
                runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
                  s"JOIN skewData2 ON key1 = key2")
              val shuffles2 = collect(adaptive2) {
@@ -2420,7 +2409,7 @@ index 272be70f9fe..06957694002 100644
              }
              if (hasRequiredDistribution) {
                assert(shuffles2.size == 3)
-@@ -2382,7 +2420,8 @@ class AdaptiveQueryExecSuite
+@@ -2563,7 +2601,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2430,7 +2419,7 @@ index 272be70f9fe..06957694002 100644
      CostEvaluator.instantiate(
        classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
      intercept[IllegalArgumentException] {
-@@ -2513,7 +2552,8 @@ class AdaptiveQueryExecSuite
+@@ -2694,7 +2733,8 @@ class AdaptiveQueryExecSuite
    }
  
    test("SPARK-48037: Fix SortShuffleWriter lacks shuffle write related metrics " +
@@ -2440,7 +2429,7 @@ index 272be70f9fe..06957694002 100644
      withTable("t3") {
        withSQLConf(
          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2548,6 +2588,7 @@ class AdaptiveQueryExecSuite
+@@ -2729,6 +2769,7 @@ class AdaptiveQueryExecSuite
            val (_, adaptive) = runAdaptiveAndVerifyResult(query)
            assert(adaptive.collect {
              case sort: SortExec => sort
@@ -2448,7 +2437,7 @@ index 272be70f9fe..06957694002 100644
            }.size == 1)
            val read = collect(adaptive) {
              case read: AQEShuffleReadExec => read
-@@ -2565,7 +2606,8 @@ class AdaptiveQueryExecSuite
+@@ -2746,7 +2787,8 @@ class AdaptiveQueryExecSuite
      }
    }
  
@@ -2458,7 +2447,7 @@ index 272be70f9fe..06957694002 100644
      withTempView("v") {
        withSQLConf(
          SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
-@@ -2677,7 +2719,7 @@ class AdaptiveQueryExecSuite
+@@ -2858,7 +2900,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value2 = value3")
          val shuffles1 = collect(adaptive1) {
@@ -2467,7 +2456,7 @@ index 272be70f9fe..06957694002 100644
          }
          assert(shuffles1.size == 4)
          val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2688,7 +2730,7 @@ class AdaptiveQueryExecSuite
+@@ -2869,7 +2911,7 @@ class AdaptiveQueryExecSuite
            runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
              "JOIN skewData3 ON value1 = value3")
          val shuffles2 = collect(adaptive2) {
@@ -2476,7 +2465,7 @@ index 272be70f9fe..06957694002 100644
          }
          assert(shuffles2.size == 4)
          val smj2 = findTopLevelSortMergeJoin(adaptive2)
-@@ -2946,6 +2988,7 @@ class AdaptiveQueryExecSuite
+@@ -3127,6 +3169,7 @@ class AdaptiveQueryExecSuite
          }.size == (if (firstAccess) 1 else 0))
          assert(collect(initialExecutedPlan) {
            case s: SortExec => s
@@ -2484,7 +2473,7 @@ index 272be70f9fe..06957694002 100644
          }.size == (if (firstAccess) 2 else 0))
          assert(collect(initialExecutedPlan) {
            case i: InMemoryTableScanLike => i
-@@ -2958,6 +3001,7 @@ class AdaptiveQueryExecSuite
+@@ -3139,6 +3182,7 @@ class AdaptiveQueryExecSuite
          }.isEmpty)
          assert(collect(finalExecutedPlan) {
            case s: SortExec => s
@@ -2493,7 +2482,7 @@ index 272be70f9fe..06957694002 100644
          assert(collect(initialExecutedPlan) {
            case i: InMemoryTableScanLike => i
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
-index 0a0b23d1e60..5685926250f 100644
+index 269990d7d14..56f200c322a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
 @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.Concat
@@ -2504,7 +2493,7 @@ index 0a0b23d1e60..5685926250f 100644
  import org.apache.spark.sql.execution.FileSourceScanExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  import org.apache.spark.sql.functions._
-@@ -868,6 +869,7 @@ abstract class SchemaPruningSuite
+@@ -884,6 +885,7 @@ abstract class SchemaPruningSuite
      val fileSourceScanSchemata =
        collect(df.queryExecution.executedPlan) {
          case scan: FileSourceScanExec => scan.requiredSchema
@@ -2513,7 +2502,7 @@ index 0a0b23d1e60..5685926250f 100644
      assert(fileSourceScanSchemata.size === expectedSchemaCatalogStrings.size,
        s"Found ${fileSourceScanSchemata.size} file sources in dataframe, " +
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
-index 80d771428d9..9327dca6c21 100644
+index a46afcef3cd..93239aede6c 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
 @@ -17,9 +17,10 @@
@@ -2528,7 +2517,7 @@ index 80d771428d9..9327dca6c21 100644
  import org.apache.spark.sql.execution.{QueryExecution, SortExec}
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
  import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
-@@ -226,6 +227,7 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
+@@ -244,6 +245,7 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
            // assert the outer most sort in the executed plan
            assert(plan.collectFirst {
              case s: SortExec => s
@@ -2536,7 +2525,7 @@ index 80d771428d9..9327dca6c21 100644
            }.exists {
              case SortExec(Seq(
                SortOrder(AttributeReference("key", IntegerType, _, _), Ascending, NullsFirst, _),
-@@ -273,6 +275,7 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
+@@ -291,6 +293,7 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
          // assert the outer most sort in the executed plan
          assert(plan.collectFirst {
            case s: SortExec => s
@@ -2544,7 +2533,7 @@ index 80d771428d9..9327dca6c21 100644
          }.exists {
            case SortExec(Seq(
              SortOrder(AttributeReference("value", StringType, _, _), Ascending, NullsFirst, _),
-@@ -306,7 +309,8 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
+@@ -324,7 +327,8 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
      }
    }
  
@@ -2555,7 +2544,7 @@ index 80d771428d9..9327dca6c21 100644
        withTable("t") {
          sql(
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
-index 62f2f2cb10a..feef4bb2928 100644
+index 7d2166beb2d..1ca5bc0d729 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
 @@ -28,7 +28,7 @@ import org.apache.hadoop.fs.{FileStatus, FileSystem, GlobFilter, Path}
@@ -2602,7 +2591,7 @@ index cd6f41b4ef4..4b6a17344bc 100644
        ParquetOutputFormat.WRITER_VERSION -> ParquetProperties.WriterVersion.PARQUET_2_0.toString
      )
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
-index 6080a5e8e4b..9aa8f49a62b 100644
+index 6b73cc8618d..69ab49344ce 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
 @@ -1102,7 +1102,11 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
@@ -2693,7 +2682,7 @@ index 6080a5e8e4b..9aa8f49a62b 100644
      withTempPath { dir =>
        val path = dir.getCanonicalPath
        spark.range(100).selectExpr("id * 2 AS id")
-@@ -2305,7 +2321,11 @@ class ParquetV1FilterSuite extends ParquetFilterSuite {
+@@ -2351,7 +2367,11 @@ class ParquetV1FilterSuite extends ParquetFilterSuite {
            assert(pushedParquetFilters.exists(_.getClass === filterClass),
              s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
  
@@ -2706,7 +2695,7 @@ index 6080a5e8e4b..9aa8f49a62b 100644
          } else {
            assert(selectedFilters.isEmpty, "There is filter pushed down")
          }
-@@ -2368,7 +2388,11 @@ class ParquetV2FilterSuite extends ParquetFilterSuite {
+@@ -2414,7 +2434,11 @@ class ParquetV2FilterSuite extends ParquetFilterSuite {
            assert(pushedParquetFilters.exists(_.getClass === filterClass),
              s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
  
@@ -2720,10 +2709,10 @@ index 6080a5e8e4b..9aa8f49a62b 100644
          case _ => assert(false, "Can not match ParquetTable in the query.")
        }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
-index 4474ec1fd42..97910c4fc3a 100644
+index 3072657a095..b10fde44730 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
-@@ -1344,7 +1344,8 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
+@@ -1567,7 +1567,8 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
      }
    }
  
@@ -2734,7 +2723,7 @@ index 4474ec1fd42..97910c4fc3a 100644
        checkAnswer(
          // "fruit" column in this file is encoded using DELTA_LENGTH_BYTE_ARRAY.
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
-index bba71f1c48d..38c60ee2584 100644
+index c530dc0d3df..d23069689da 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
 @@ -996,7 +996,11 @@ abstract class ParquetQuerySuite extends QueryTest with ParquetTest with SharedS
@@ -2890,7 +2879,7 @@ index 5c0b7def039..151184bc98c 100644
      assert(fileSourceScanSchemata.size === expectedSchemaCatalogStrings.size,
        s"Found ${fileSourceScanSchemata.size} file sources in dataframe, " +
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
-index 0acb21f3e6f..3a7bb73f03c 100644
+index 56076175d60..eec8b442da1 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
 @@ -27,7 +27,7 @@ import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
@@ -2969,7 +2958,7 @@ index b8f3ea3c6f3..bbd44221288 100644
        val workDirPath = workDir.getAbsolutePath
        val input = spark.range(5).toDF("id")
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
-index 0dd90925d3c..7d53ec845ef 100644
+index 402365a59ec..82b0868b263 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
 @@ -46,8 +46,10 @@ import org.apache.spark.sql.util.QueryExecutionListener
@@ -2984,7 +2973,7 @@ index 0dd90925d3c..7d53ec845ef 100644
    import testImplicits._
  
    /**
-@@ -765,7 +767,8 @@ class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
+@@ -786,7 +788,8 @@ class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
      }
    }
  
@@ -2995,7 +2984,7 @@ index 0dd90925d3c..7d53ec845ef 100644
        spark.range(10).selectExpr("id", "id % 3 as p")
          .write.partitionBy("p").saveAsTable("testDataForScan")
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
-index 0ab8691801d..d9125f658ad 100644
+index 79e157c4db6..a91bca2e9a7 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
 @@ -18,6 +18,7 @@
@@ -3050,10 +3039,10 @@ index 0ab8691801d..d9125f658ad 100644
            assert(scanNodes.length == 1)
            // $"a" is not null and $"a" > 1
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
-index 7838e62013d..8fa09652921 100644
+index e31e0e70cf3..034ad5b953e 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
-@@ -37,8 +37,10 @@ import org.apache.spark.sql.streaming.{StreamingQuery, StreamingQueryException,
+@@ -38,8 +38,10 @@ import org.apache.spark.sql.streaming.{StreamingQuery, StreamingQueryException,
  import org.apache.spark.sql.streaming.util.StreamManualClock
  import org.apache.spark.util.Utils
  
@@ -3066,7 +3055,7 @@ index 7838e62013d..8fa09652921 100644
    import testImplicits._
  
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
-index c4b09c4b289..a2f8ca47ffb 100644
+index 83ebd24384c..eeaa1d94bc0 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
 @@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.expressions
@@ -3261,10 +3250,10 @@ index c5c56f081d8..197cd241f48 100644
      }
  
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
-index 9742a004545..4e0417d730a 100644
+index 4c06a0109e3..0b0e4990c0a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
-@@ -34,6 +34,7 @@ import org.apache.spark.paths.SparkPath
+@@ -35,6 +35,7 @@ import org.apache.spark.paths.SparkPath
  import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
  import org.apache.spark.sql.{AnalysisException, DataFrame}
  import org.apache.spark.sql.catalyst.util.stringToFile
@@ -3272,7 +3261,7 @@ index 9742a004545..4e0417d730a 100644
  import org.apache.spark.sql.execution.DataSourceScanExec
  import org.apache.spark.sql.execution.datasources._
  import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
-@@ -786,6 +787,8 @@ class FileStreamSinkV2Suite extends FileStreamSinkSuite {
+@@ -804,6 +805,8 @@ class FileStreamSinkV2Suite extends FileStreamSinkSuite {
        val fileScan = df.queryExecution.executedPlan.collect {
          case batch: BatchScanExec if batch.scan.isInstanceOf[FileScan] =>
            batch.scan.asInstanceOf[FileScan]
@@ -3282,10 +3271,10 @@ index 9742a004545..4e0417d730a 100644
          fail(s"No FileScan in query\n${df.queryExecution}")
        }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
-index b0967d5ffdf..3d567f913de 100644
+index 11b761f1e4a..e54bf2a996a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
-@@ -40,6 +40,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
+@@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
  import org.apache.spark.sql.catalyst.util.DateTimeUtils
  import org.apache.spark.sql.classic.{DataFrame, Dataset}
  import org.apache.spark.sql.classic.ClassicConversions._
@@ -3293,8 +3282,8 @@ index b0967d5ffdf..3d567f913de 100644
  import org.apache.spark.sql.execution.{LocalLimitExec, SimpleMode, SparkPlan}
  import org.apache.spark.sql.execution.command.ExplainCommand
  import org.apache.spark.sql.execution.streaming._
-@@ -1118,11 +1119,12 @@ class StreamSuite extends StreamTest {
-       val localLimits = execPlan.collect {
+@@ -1119,11 +1120,12 @@ class StreamSuite extends StreamTest {
+       val localLimits = StreamingQueryPlanTraverseHelper.collectFromUnfoldedPlan(execPlan) {
          case l: LocalLimitExec => l
          case l: StreamingLocalLimitExec => l
 +        case l: CometLocalLimitExec => l
@@ -3307,7 +3296,7 @@ index b0967d5ffdf..3d567f913de 100644
  
        if (expectStreamingLimit) {
          assert(
-@@ -1130,7 +1132,8 @@ class StreamSuite extends StreamTest {
+@@ -1131,7 +1133,8 @@ class StreamSuite extends StreamTest {
            s"Local limit was not StreamingLocalLimitExec:\n$execPlan")
        } else {
          assert(
@@ -3318,17 +3307,17 @@ index b0967d5ffdf..3d567f913de 100644
        }
      }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-index b4c4ec7acbf..20579284856 100644
+index 465da3cd469..92ac998929d 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
-@@ -23,6 +23,7 @@ import org.apache.commons.io.FileUtils
+@@ -22,6 +22,7 @@ import java.io.File
  import org.scalatest.Assertions
  
  import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
 +import org.apache.spark.sql.comet.CometHashAggregateExec
  import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
- import org.apache.spark.sql.execution.streaming.{MemoryStream, StateStoreRestoreExec, StateStoreSaveExec}
- import org.apache.spark.sql.functions.count
+ import org.apache.spark.sql.execution.streaming.operators.stateful.{StateStoreRestoreExec, StateStoreSaveExec}
+ import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
 @@ -67,6 +68,7 @@ class StreamingAggregationDistributionSuite extends StreamTest
          // verify aggregations in between, except partial aggregation
          val allAggregateExecs = query.lastExecution.executedPlan.collect {
@@ -3346,19 +3335,19 @@ index b4c4ec7acbf..20579284856 100644
  
          val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
-index d3c44dcead3..8096bce4436 100644
+index 22028a585e2..20c6b7c796a 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
-@@ -33,7 +33,7 @@ import org.apache.spark.sql.{DataFrame, Row, SparkSession}
+@@ -34,7 +34,7 @@ import org.apache.spark.sql.{DataFrame, Row, SparkSession}
  import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression}
  import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
  import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions
 -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
- import org.apache.spark.sql.execution.streaming.{MemoryStream, StatefulOperatorStateInfo, StreamingSymmetricHashJoinExec, StreamingSymmetricHashJoinHelper}
- import org.apache.spark.sql.execution.streaming.state.{RocksDBStateStoreProvider, StateStore, StateStoreProviderId}
- import org.apache.spark.sql.functions._
-@@ -642,14 +642,28 @@ class StreamingInnerJoinSuite extends StreamingJoinSuite {
+ import org.apache.spark.sql.execution.streaming.checkpointing.CheckpointFileManager
+ import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorStateInfo
+ import org.apache.spark.sql.execution.streaming.operators.stateful.join.{StreamingSymmetricHashJoinExec, StreamingSymmetricHashJoinHelper}
+@@ -713,14 +713,28 @@ class StreamingInnerJoinSuite extends StreamingJoinSuite {
  
          val numPartitions = spark.sessionState.conf.getConf(SQLConf.SHUFFLE_PARTITIONS)
  
@@ -3396,19 +3385,19 @@ index d3c44dcead3..8096bce4436 100644
    }
  
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
-index e33d4f1f6ab..ce0a21d1e9d 100644
+index 7adf98b7920..fd1696da07b 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
-@@ -44,7 +44,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
+@@ -43,7 +43,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
  import org.apache.spark.sql.classic.{DataFrame, Dataset}
  import org.apache.spark.sql.connector.read.InputPartition
  import org.apache.spark.sql.connector.read.streaming.{Offset => OffsetV2, ReadLimit}
 -import org.apache.spark.sql.execution.exchange.{REQUIRED_BY_STATEFUL_OPERATOR, ReusedExchangeExec, ShuffleExchangeExec}
 +import org.apache.spark.sql.execution.exchange.{REQUIRED_BY_STATEFUL_OPERATOR, ReusedExchangeExec, ShuffleExchangeLike}
  import org.apache.spark.sql.execution.streaming._
- import org.apache.spark.sql.execution.streaming.sources.{MemorySink, TestForeachWriter}
- import org.apache.spark.sql.functions._
-@@ -1462,7 +1462,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
+ import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata}
+ import org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec
+@@ -1471,7 +1471,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
        CheckAnswer((1, 2), (2, 2), (3, 2)),
        Execute { qe =>
          val shuffleOpt = qe.lastExecution.executedPlan.collect {
@@ -3418,7 +3407,7 @@ index e33d4f1f6ab..ce0a21d1e9d 100644
  
          assert(shuffleOpt.nonEmpty, "No shuffle exchange found in the query plan")
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
-index 86c4e49f6f6..2e639e5f38d 100644
+index 89f65562295..166a11a9079 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
 @@ -22,7 +22,7 @@ import java.util
@@ -3591,7 +3580,7 @@ index 982d57fb287..6017f36c440 100644
    implicit val formats: DefaultFormats = new DefaultFormats {
      override def dateFormatter = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss")
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
-index 52abd248f3a..7a199931a08 100644
+index ff8cbc0ad6e..847fc3f15f4 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
 +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
 @@ -19,6 +19,7 @@ package org.apache.spark.sql.hive
@@ -3602,7 +3591,7 @@ index 52abd248f3a..7a199931a08 100644
  import org.apache.spark.sql.execution._
  import org.apache.spark.sql.execution.adaptive.{DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
  import org.apache.spark.sql.hive.execution.HiveTableScanExec
-@@ -35,6 +36,9 @@ abstract class DynamicPartitionPruningHiveScanSuiteBase
+@@ -36,6 +37,9 @@ abstract class DynamicPartitionPruningHiveScanSuiteBase
        case s: FileSourceScanExec => s.partitionFilters.collect {
          case d: DynamicPruningExpression => d.child
        }
@@ -3644,10 +3633,10 @@ index 4b27082e188..09f591dfed3 100644
          withUserDefinedFunction(udfInfo.funcName -> false) {
            val sparkClassLoader = Thread.currentThread().getContextClassLoader
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
-index cc7bb193731..06555d48da7 100644
+index f9c001a1a07..8f63820113e 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
 +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
-@@ -818,7 +818,8 @@ class InsertSuite extends QueryTest with TestHiveSingleton with BeforeAndAfter
+@@ -819,7 +819,8 @@ class InsertSuite extends QueryTest with TestHiveSingleton with BeforeAndAfter
      }
    }