From 6143e7a9973521844fb7e898a0f22a9c185972bc Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Wed, 5 Jun 2024 20:52:50 -0700
Subject: [PATCH 1/8] chore: Add UnboundColumn to carry datatype for unbound
 reference (#518)

* chore: Add UnboundColumn to carry datatype for unbound reference

* Update core/src/execution/datafusion/expressions/unbound.rs
---
 .../execution/datafusion/expressions/mod.rs   |   1 +
 .../datafusion/expressions/unbound.rs         | 110 ++++++++++++++++++
 core/src/execution/datafusion/planner.rs      |  11 +-
 3 files changed, 120 insertions(+), 2 deletions(-)
 create mode 100644 core/src/execution/datafusion/expressions/unbound.rs

diff --git a/core/src/execution/datafusion/expressions/mod.rs b/core/src/execution/datafusion/expressions/mod.rs
index 084fef2df..05230b4c2 100644
--- a/core/src/execution/datafusion/expressions/mod.rs
+++ b/core/src/execution/datafusion/expressions/mod.rs
@@ -36,5 +36,6 @@ pub mod strings;
 pub mod subquery;
 pub mod sum_decimal;
 pub mod temporal;
+pub mod unbound;
 mod utils;
 pub mod variance;
diff --git a/core/src/execution/datafusion/expressions/unbound.rs b/core/src/execution/datafusion/expressions/unbound.rs
new file mode 100644
index 000000000..5387b1012
--- /dev/null
+++ b/core/src/execution/datafusion/expressions/unbound.rs
@@ -0,0 +1,110 @@
+// 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.
+
+use crate::execution::datafusion::expressions::utils::down_cast_any_ref;
+use arrow_array::RecordBatch;
+use arrow_schema::{DataType, Schema};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_common::{internal_err, Result};
+use datafusion_physical_expr::PhysicalExpr;
+use std::{
+    any::Any,
+    hash::{Hash, Hasher},
+    sync::Arc,
+};
+
+/// This is similar to `UnKnownColumn` in DataFusion, but it has data type.
+/// This is only used when the column is not bound to a schema, for example, the
+/// inputs to aggregation functions in final aggregation. In the case, we cannot
+/// bind the aggregation functions to the input schema which is grouping columns
+/// and aggregate buffer attributes in Spark (DataFusion has different design).
+/// But when creating certain aggregation functions, we need to know its input
+/// data types. As `UnKnownColumn` doesn't have data type, we implement this
+/// `UnboundColumn` to carry the data type.
+#[derive(Debug, Hash, PartialEq, Eq, Clone)]
+pub struct UnboundColumn {
+    name: String,
+    datatype: DataType,
+}
+
+impl UnboundColumn {
+    /// Create a new unbound column expression
+    pub fn new(name: &str, datatype: DataType) -> Self {
+        Self {
+            name: name.to_owned(),
+            datatype,
+        }
+    }
+
+    /// Get the column name
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+}
+
+impl std::fmt::Display for UnboundColumn {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{}, datatype: {}", self.name, self.datatype)
+    }
+}
+
+impl PhysicalExpr for UnboundColumn {
+    /// Return a reference to Any that can be used for downcasting
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    /// Get the data type of this expression, given the schema of the input
+    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+        Ok(self.datatype.clone())
+    }
+
+    /// Decide whether this expression is nullable, given the schema of the input
+    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+        Ok(true)
+    }
+
+    /// Evaluate the expression
+    fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
+        internal_err!("UnboundColumn::evaluate() should not be called")
+    }
+
+    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+        vec![]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        _children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        Ok(self)
+    }
+
+    fn dyn_hash(&self, state: &mut dyn Hasher) {
+        let mut s = state;
+        self.hash(&mut s);
+    }
+}
+
+impl PartialEq<dyn Any> for UnboundColumn {
+    fn eq(&self, other: &dyn Any) -> bool {
+        down_cast_any_ref(other)
+            .downcast_ref::<Self>()
+            .map(|x| self == x)
+            .unwrap_or(false)
+    }
+}
diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
index a5bcf5654..7af5f6838 100644
--- a/core/src/execution/datafusion/planner.rs
+++ b/core/src/execution/datafusion/planner.rs
@@ -33,7 +33,7 @@ use datafusion::{
         expressions::{
             in_list, BinaryExpr, BitAnd, BitOr, BitXor, CaseExpr, CastExpr, Column, Count,
             FirstValue, InListExpr, IsNotNullExpr, IsNullExpr, LastValue,
-            Literal as DataFusionLiteral, Max, Min, NotExpr, Sum, UnKnownColumn,
+            Literal as DataFusionLiteral, Max, Min, NotExpr, Sum,
         },
         AggregateExpr, PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr,
     },
@@ -78,6 +78,7 @@ use crate::{
                 subquery::Subquery,
                 sum_decimal::SumDecimal,
                 temporal::{DateTruncExec, HourExec, MinuteExec, SecondExec, TimestampTruncExec},
+                unbound::UnboundColumn,
                 variance::Variance,
                 NormalizeNaNAndZero,
             },
@@ -239,7 +240,13 @@ impl PhysicalPlanner {
                 let field = input_schema.field(idx);
                 Ok(Arc::new(Column::new(field.name().as_str(), idx)))
             }
-            ExprStruct::Unbound(unbound) => Ok(Arc::new(UnKnownColumn::new(unbound.name.as_str()))),
+            ExprStruct::Unbound(unbound) => {
+                let data_type = to_arrow_datatype(unbound.datatype.as_ref().unwrap());
+                Ok(Arc::new(UnboundColumn::new(
+                    unbound.name.as_str(),
+                    data_type,
+                )))
+            }
             ExprStruct::IsNotNull(is_notnull) => {
                 let child = self.create_expr(is_notnull.child.as_ref().unwrap(), input_schema)?;
                 Ok(Arc::new(IsNotNullExpr::new(child)))

From 0fd1476cdd9e3892137d55a869fcfe906d74d9de Mon Sep 17 00:00:00 2001
From: KAZUYUKI TANIMURA <ktanimura@apple.com>
Date: Wed, 5 Jun 2024 23:50:22 -0700
Subject: [PATCH 2/8] chore: Remove 3.4.2.diff (#528)

---
 dev/diffs/3.4.2.diff | 2555 ------------------------------------------
 1 file changed, 2555 deletions(-)
 delete mode 100644 dev/diffs/3.4.2.diff

diff --git a/dev/diffs/3.4.2.diff b/dev/diffs/3.4.2.diff
deleted file mode 100644
index cd02970da..000000000
--- a/dev/diffs/3.4.2.diff
+++ /dev/null
@@ -1,2555 +0,0 @@
-diff --git a/pom.xml b/pom.xml
-index fab98342498..f2156d790d1 100644
---- a/pom.xml
-+++ b/pom.xml
-@@ -148,6 +148,8 @@
-     <chill.version>0.10.0</chill.version>
-     <ivy.version>2.5.1</ivy.version>
-     <oro.version>2.0.8</oro.version>
-+    <spark.version.short>3.4</spark.version.short>
-+    <comet.version>0.1.0-SNAPSHOT</comet.version>
-     <!--
-     If you changes codahale.metrics.version, you also need to change
-     the link to metrics.dropwizard.io in docs/monitoring.md.
-@@ -2766,6 +2768,25 @@
-         <artifactId>arpack</artifactId>
-         <version>${netlib.ludovic.dev.version}</version>
-       </dependency>
-+      <dependency>
-+        <groupId>org.apache.comet</groupId>
-+        <artifactId>comet-spark-spark${spark.version.short}_${scala.binary.version}</artifactId>
-+        <version>${comet.version}</version>
-+        <exclusions>
-+          <exclusion>
-+            <groupId>org.apache.spark</groupId>
-+            <artifactId>spark-sql_${scala.binary.version}</artifactId>
-+          </exclusion>
-+          <exclusion>
-+            <groupId>org.apache.spark</groupId>
-+            <artifactId>spark-core_${scala.binary.version}</artifactId>
-+          </exclusion>
-+          <exclusion>
-+            <groupId>org.apache.spark</groupId>
-+            <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
-+          </exclusion>
-+        </exclusions>
-+      </dependency>
-     </dependencies>
-   </dependencyManagement>
- 
-diff --git a/sql/core/pom.xml b/sql/core/pom.xml
-index 5b6cc8cb7af..5ce708adc38 100644
---- a/sql/core/pom.xml
-+++ b/sql/core/pom.xml
-@@ -77,6 +77,10 @@
-       <groupId>org.apache.spark</groupId>
-       <artifactId>spark-tags_${scala.binary.version}</artifactId>
-     </dependency>
-+    <dependency>
-+      <groupId>org.apache.comet</groupId>
-+      <artifactId>comet-spark-spark${spark.version.short}_${scala.binary.version}</artifactId>
-+    </dependency>
- 
-     <!--
-       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/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-index c595b50950b..6b60213e775 100644
---- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-+++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala
-@@ -102,7 +102,7 @@ class SparkSession private(
-       sc: SparkContext,
-       initialSessionOptions: java.util.HashMap[String, String]) = {
-     this(sc, None, None,
--      SparkSession.applyExtensions(
-+      SparkSession.applyExtensions(sc,
-         sc.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS).getOrElse(Seq.empty),
-         new SparkSessionExtensions), initialSessionOptions.asScala.toMap)
-   }
-@@ -1028,7 +1028,7 @@ object SparkSession extends Logging {
-         }
- 
-         loadExtensions(extensions)
--        applyExtensions(
-+        applyExtensions(sparkContext,
-           sparkContext.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS).getOrElse(Seq.empty),
-           extensions)
- 
-@@ -1282,14 +1282,24 @@ object SparkSession extends Logging {
-     }
-   }
- 
-+  private def loadCometExtension(sparkContext: SparkContext): Seq[String] = {
-+    if (sparkContext.getConf.getBoolean("spark.comet.enabled", false)) {
-+      Seq("org.apache.comet.CometSparkSessionExtensions")
-+    } else {
-+      Seq.empty
-+    }
-+  }
-+
-   /**
-    * Initialize extensions for given extension classnames. The classes will be applied to the
-    * extensions passed into this function.
-    */
-   private def applyExtensions(
-+      sparkContext: SparkContext,
-       extensionConfClassNames: Seq[String],
-       extensions: SparkSessionExtensions): SparkSessionExtensions = {
--    extensionConfClassNames.foreach { extensionConfClassName =>
-+    val extensionClassNames = extensionConfClassNames ++ loadCometExtension(sparkContext)
-+    extensionClassNames.foreach { extensionConfClassName =>
-       try {
-         val extensionConfClass = Utils.classForName(extensionConfClassName)
-         val extensionConf = extensionConfClass.getConstructor().newInstance()
-diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
-index db587dd9868..aac7295a53d 100644
---- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
-+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
-@@ -18,6 +18,7 @@
- package org.apache.spark.sql.execution
- 
- import org.apache.spark.annotation.DeveloperApi
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec}
- import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
- import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
-@@ -67,6 +68,7 @@ private[execution] object SparkPlanInfo {
-     // dump the file scan metadata (e.g file path) to event log
-     val metadata = plan match {
-       case fileScan: FileSourceScanExec => fileScan.metadata
-+      case cometScan: CometScanExec => cometScan.metadata
-       case _ => Map[String, String]()
-     }
-     new SparkPlanInfo(
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql b/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
-index 7aef901da4f..f3d6e18926d 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
-@@ -2,3 +2,4 @@
- 
- --SET spark.sql.adaptive.enabled=true
- --SET spark.sql.maxMetadataStringLength = 500
-+--SET spark.comet.enabled = false
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql b/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
-index eeb2180f7a5..afd1b5ec289 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
-@@ -1,5 +1,6 @@
- --SET spark.sql.cbo.enabled=true
- --SET spark.sql.maxMetadataStringLength = 500
-+--SET spark.comet.enabled = false
- 
- 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
---- 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 @@
- --SET spark.sql.codegen.wholeStage = true
- --SET spark.sql.adaptive.enabled = false
- --SET spark.sql.maxMetadataStringLength = 500
-+--SET spark.comet.enabled = false
- 
- -- Test tables
- CREATE table  explain_temp1 (key int, val int) USING PARQUET;
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part1.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part1.sql
-index 1152d77da0c..f77493f690b 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part1.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part1.sql
-@@ -7,6 +7,9 @@
- 
- -- avoid bit-exact output here because operations may not be bit-exact.
- -- SET extra_float_digits = 0;
-+-- Disable Comet exec due to floating point precision difference
-+--SET spark.comet.exec.enabled = false
-+
- 
- -- Test aggregate operator with codegen on and off.
- --CONFIG_DIM1 spark.sql.codegen.wholeStage=true
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
-index 41fd4de2a09..44cd244d3b0 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
-@@ -5,6 +5,9 @@
- -- AGGREGATES [Part 3]
- -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/aggregates.sql#L352-L605
- 
-+-- Disable Comet exec due to floating point precision difference
-+--SET spark.comet.exec.enabled = false
-+
- -- Test aggregate operator with codegen on and off.
- --CONFIG_DIM1 spark.sql.codegen.wholeStage=true
- --CONFIG_DIM1 spark.sql.codegen.wholeStage=false,spark.sql.codegen.factoryMode=CODEGEN_ONLY
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
-index fac23b4a26f..2b73732c33f 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
-@@ -1,6 +1,10 @@
- --
- -- Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
- --
-+
-+-- Disable Comet exec due to floating point precision difference
-+--SET spark.comet.exec.enabled = false
-+
- --
- -- INT8
- -- Test int8 64-bit integers.
-diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
-index 0efe0877e9b..423d3b3d76d 100644
---- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
-+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
-@@ -1,6 +1,10 @@
- --
- -- Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
- --
-+
-+-- Disable Comet exec due to floating point precision difference
-+--SET spark.comet.exec.enabled = false
-+
- --
- -- SELECT_HAVING
- -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql
-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 cf40e944c09..bdd5be4f462 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
-@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants
- import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan}
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.columnar._
--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
-@@ -516,7 +516,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils
-    */
-   private def verifyNumExchanges(df: DataFrame, expected: Int): Unit = {
-     assert(
--      collect(df.queryExecution.executedPlan) { case e: ShuffleExchangeExec => e }.size == expected)
-+      collect(df.queryExecution.executedPlan) {
-+        case _: ShuffleExchangeLike => 1 }.size == expected)
-   }
- 
-   test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
-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 ea5e47ede55..814b92d090f 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
-@@ -27,7 +27,7 @@ import org.apache.spark.SparkException
- import org.apache.spark.sql.execution.WholeStageCodegenExec
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
--import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
- import org.apache.spark.sql.expressions.Window
- import org.apache.spark.sql.functions._
- import org.apache.spark.sql.internal.SQLConf
-@@ -755,7 +755,7 @@ class DataFrameAggregateSuite extends QueryTest
-       assert(objHashAggPlans.nonEmpty)
- 
-       val exchangePlans = collect(aggPlan) {
--        case shuffle: ShuffleExchangeExec => shuffle
-+        case shuffle: ShuffleExchangeLike => shuffle
-       }
-       assert(exchangePlans.length == 1)
-     }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
-index 56e9520fdab..917932336df 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
-@@ -435,7 +435,9 @@ class DataFrameJoinSuite extends QueryTest
- 
-     withTempDatabase { dbName =>
-       withTable(table1Name, table2Name) {
--        withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-+        withSQLConf(
-+            SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
-+            "spark.comet.enabled" -> "false") {
-           spark.range(50).write.saveAsTable(s"$dbName.$table1Name")
-           spark.range(100).write.saveAsTable(s"$dbName.$table2Name")
- 
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
-index 9ddb4abe98b..1b9269acef1 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
-@@ -43,7 +43,7 @@ import org.apache.spark.sql.connector.FakeV2Provider
- import org.apache.spark.sql.execution.{FilterExec, LogicalRDD, QueryExecution, SortExec, WholeStageCodegenExec}
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.aggregate.HashAggregateExec
--import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
-+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeLike}
- import org.apache.spark.sql.expressions.{Aggregator, Window}
- import org.apache.spark.sql.functions._
- import org.apache.spark.sql.internal.SQLConf
-@@ -1981,7 +1981,7 @@ class DataFrameSuite extends QueryTest
-           fail("Should not have back to back Aggregates")
-         }
-         atFirstAgg = true
--      case e: ShuffleExchangeExec => atFirstAgg = false
-+      case e: ShuffleExchangeLike => atFirstAgg = false
-       case _ =>
-     }
-   }
-@@ -2291,7 +2291,7 @@ class DataFrameSuite extends QueryTest
-       checkAnswer(join, df)
-       assert(
-         collect(join.queryExecution.executedPlan) {
--          case e: ShuffleExchangeExec => true }.size === 1)
-+          case _: ShuffleExchangeLike => true }.size === 1)
-       assert(
-         collect(join.queryExecution.executedPlan) { case e: ReusedExchangeExec => true }.size === 1)
-       val broadcasted = broadcast(join)
-@@ -2299,7 +2299,7 @@ class DataFrameSuite extends QueryTest
-       checkAnswer(join2, df)
-       assert(
-         collect(join2.queryExecution.executedPlan) {
--          case e: ShuffleExchangeExec => true }.size == 1)
-+          case _: ShuffleExchangeLike => true }.size == 1)
-       assert(
-         collect(join2.queryExecution.executedPlan) {
-           case e: BroadcastExchangeExec => true }.size === 1)
-@@ -2862,7 +2862,7 @@ class DataFrameSuite extends QueryTest
- 
-     // Assert that no extra shuffle introduced by cogroup.
-     val exchanges = collect(df3.queryExecution.executedPlan) {
--      case h: ShuffleExchangeExec => h
-+      case h: ShuffleExchangeLike => h
-     }
-     assert(exchanges.size == 2)
-   }
-@@ -3311,7 +3311,8 @@ class DataFrameSuite extends QueryTest
-     assert(df2.isLocal)
-   }
- 
--  test("SPARK-35886: PromotePrecision should be subexpr replaced") {
-+  test("SPARK-35886: PromotePrecision should be subexpr replaced",
-+    IgnoreComet("TODO: fix Comet for this test")) {
-     withTable("tbl") {
-       sql(
-         """
-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 7dec558f8df..840dda15033 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
-@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi}
- import org.apache.spark.sql.catalyst.util.sideBySide
- import org.apache.spark.sql.execution.{LogicalRDD, RDDScanExec, SQLExecution}
- 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.expressions.UserDefinedFunction
- import org.apache.spark.sql.functions._
-@@ -2254,7 +2254,7 @@ class DatasetSuite extends QueryTest
- 
-     // Assert that no extra shuffle introduced by cogroup.
-     val exchanges = collect(df3.queryExecution.executedPlan) {
--      case h: ShuffleExchangeExec => h
-+      case h: ShuffleExchangeLike => h
-     }
-     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 f33432ddb6f..060f874ea72 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
- import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression}
- import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
- import org.apache.spark.sql.catalyst.plans.ExistenceJoin
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog}
- import org.apache.spark.sql.execution._
- import org.apache.spark.sql.execution.adaptive._
-@@ -262,6 +263,9 @@ abstract class DynamicPartitionPruningSuiteBase
-       case s: BatchScanExec => s.runtimeFilters.collect {
-         case d: DynamicPruningExpression => d.child
-       }
-+      case s: CometScanExec => s.partitionFilters.collect {
-+        case d: DynamicPruningExpression => d.child
-+      }
-       case _ => Nil
-     }
-   }
-@@ -1187,7 +1191,8 @@ abstract class DynamicPartitionPruningSuiteBase
-     }
-   }
- 
--  test("Make sure dynamic pruning works on uncorrelated queries") {
-+  test("Make sure dynamic pruning works on uncorrelated queries",
-+    IgnoreComet("TODO: Support SubqueryBroadcastExec in Comet: #242")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       val df = sql(
-         """
-@@ -1238,7 +1243,8 @@ abstract class DynamicPartitionPruningSuiteBase
-     }
-   }
- 
--  test("Plan broadcast pruning only when the broadcast can be reused") {
-+  test("Plan broadcast pruning only when the broadcast can be reused",
-+    IgnoreComet("TODO: Support SubqueryBroadcastExec in Comet: #242")) {
-     Given("dynamic pruning filter on the build side")
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       val df = sql(
-@@ -1485,7 +1491,7 @@ abstract class DynamicPartitionPruningSuiteBase
-   }
- 
-   test("SPARK-38148: Do not add dynamic partition pruning if there exists static partition " +
--    "pruning") {
-+    "pruning", IgnoreComet("TODO: Support SubqueryBroadcastExec in Comet: #242")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true") {
-       Seq(
-         "f.store_id = 1" -> false,
-@@ -1729,6 +1735,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")))
-+              case s: CometScanExec =>
-+                s.output.exists(_.exists(_.argString(maxFields = 100).contains("fid")))
-               case _ => false
-             }
-           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 a6b295578d6..91acca4306f 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
-@@ -463,7 +463,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
-     }
-   }
- 
--  test("Explain formatted output for scan operator for datasource V2") {
-+  test("Explain formatted output for scan operator for datasource V2",
-+      IgnoreComet("Comet explain output is different")) {
-     withTempDir { dir =>
-       Seq("parquet", "orc", "csv", "json").foreach { fmt =>
-         val basePath = dir.getCanonicalPath + "/" + fmt
-@@ -541,7 +542,9 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
-   }
- }
- 
--class ExplainSuiteAE extends ExplainSuiteHelper with EnableAdaptiveExecutionSuite {
-+// Ignored when Comet is enabled. Comet changes expected query plans.
-+class ExplainSuiteAE extends ExplainSuiteHelper with EnableAdaptiveExecutionSuite
-+    with IgnoreCometSuite {
-   import testImplicits._
- 
-   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 2796b1cf154..be7078b38f4 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.TestingUDT.{IntervalUDT, NullData, NullUDT}
- import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GreaterThan, Literal}
- import org.apache.spark.sql.catalyst.expressions.IntegralLiteralTestUtils.{negativeInt, positiveInt}
- import org.apache.spark.sql.catalyst.plans.logical.Filter
-+import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec, CometSortMergeJoinExec}
- import org.apache.spark.sql.execution.{FileSourceScanLike, SimpleMode}
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.datasources.FilePartition
-@@ -815,6 +816,7 @@ class FileBasedDataSourceSuite extends QueryTest
-             assert(bJoinExec.isEmpty)
-             val smJoinExec = collect(joinedDF.queryExecution.executedPlan) {
-               case smJoin: SortMergeJoinExec => smJoin
-+              case smJoin: CometSortMergeJoinExec => smJoin
-             }
-             assert(smJoinExec.nonEmpty)
-           }
-@@ -875,6 +877,7 @@ class FileBasedDataSourceSuite extends QueryTest
- 
-           val fileScan = df.queryExecution.executedPlan collectFirst {
-             case BatchScanExec(_, f: FileScan, _, _, _, _, _, _, _) => f
-+            case CometBatchScanExec(BatchScanExec(_, f: FileScan, _, _, _, _, _, _, _), _) => f
-           }
-           assert(fileScan.nonEmpty)
-           assert(fileScan.get.partitionFilters.nonEmpty)
-@@ -916,6 +919,7 @@ class FileBasedDataSourceSuite extends QueryTest
- 
-           val fileScan = df.queryExecution.executedPlan collectFirst {
-             case BatchScanExec(_, f: FileScan, _, _, _, _, _, _, _) => f
-+            case CometBatchScanExec(BatchScanExec(_, f: FileScan, _, _, _, _, _, _, _), _) => f
-           }
-           assert(fileScan.nonEmpty)
-           assert(fileScan.get.partitionFilters.isEmpty)
-@@ -1100,6 +1104,8 @@ class FileBasedDataSourceSuite extends QueryTest
-           val filters = df.queryExecution.executedPlan.collect {
-             case f: FileSourceScanLike => f.dataFilters
-             case b: BatchScanExec => b.scan.asInstanceOf[FileScan].dataFilters
-+            case b: CometScanExec => b.dataFilters
-+            case b: CometBatchScanExec => b.scan.asInstanceOf[FileScan].dataFilters
-           }.flatten
-           assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L))))
-         }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala
-new file mode 100644
-index 00000000000..4b31bea33de
---- /dev/null
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala
-@@ -0,0 +1,42 @@
-+/*
-+ * 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
-+
-+import org.scalactic.source.Position
-+import org.scalatest.Tag
-+
-+import org.apache.spark.sql.test.SQLTestUtils
-+
-+/**
-+ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`).
-+ */
-+case class IgnoreComet(reason: String) extends Tag("DisableComet")
-+
-+/**
-+ * Helper trait that disables Comet for all tests regardless of default config values.
-+ */
-+trait IgnoreCometSuite extends SQLTestUtils {
-+  override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
-+    (implicit pos: Position): Unit = {
-+    if (isCometEnabled) {
-+      ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun)
-+    } else {
-+      super.test(testName, testTags: _*)(testFun)
-+    }
-+  }
-+}
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
-index 1792b4c32eb..1616e6f39bd 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
-@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide
- import org.apache.spark.sql.catalyst.plans.PlanTest
- import org.apache.spark.sql.catalyst.plans.logical._
- import org.apache.spark.sql.catalyst.rules.RuleExecutor
-+import org.apache.spark.sql.comet.{CometHashJoinExec, CometSortMergeJoinExec}
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.joins._
- import org.apache.spark.sql.internal.SQLConf
-@@ -362,6 +363,7 @@ class JoinHintSuite extends PlanTest with SharedSparkSession with AdaptiveSparkP
-     val executedPlan = df.queryExecution.executedPlan
-     val shuffleHashJoins = collect(executedPlan) {
-       case s: ShuffledHashJoinExec => s
-+      case c: CometHashJoinExec => c.originalPlan.asInstanceOf[ShuffledHashJoinExec]
-     }
-     assert(shuffleHashJoins.size == 1)
-     assert(shuffleHashJoins.head.buildSide == buildSide)
-@@ -371,6 +373,7 @@ class JoinHintSuite extends PlanTest with SharedSparkSession with AdaptiveSparkP
-     val executedPlan = df.queryExecution.executedPlan
-     val shuffleMergeJoins = collect(executedPlan) {
-       case s: SortMergeJoinExec => s
-+      case c: CometSortMergeJoinExec => c
-     }
-     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 5125708be32..210ab4f3ce1 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,6 +29,7 @@ import org.apache.spark.sql.catalyst.TableIdentifier
- import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
- import org.apache.spark.sql.catalyst.expressions.{Ascending, GenericRow, SortOrder}
- import org.apache.spark.sql.catalyst.plans.logical.Filter
-+import org.apache.spark.sql.comet._
- import org.apache.spark.sql.execution.{BinaryExecNode, FilterExec, ProjectExec, SortExec, SparkPlan, WholeStageCodegenExec}
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
-@@ -739,7 +740,8 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
-     }
-   }
- 
--  test("test SortMergeJoin (with spill)") {
-+  test("test SortMergeJoin (with spill)",
-+      IgnoreComet("TODO: Comet SMJ doesn't support spill yet")) {
-     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") {
-@@ -1114,9 +1116,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
--      assert(collect(plan) { case _: ShuffledHashJoinExec => true }.size === 1)
-+      assert(collect(plan) {
-+        case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
-       // No extra shuffle before aggregate
--      assert(collect(plan) { case _: ShuffleExchangeExec => true }.size === 2)
-+      assert(collect(plan) {
-+        case _: ShuffleExchangeLike => true }.size === 2)
-     })
-   }
- 
-@@ -1133,10 +1137,11 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
-         .join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
-         .queryExecution
-         .executedPlan
--      assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 2)
-+      assert(collect(plan) {
-+        case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 2)
-       assert(collect(plan) { case _: BroadcastHashJoinExec => true }.size === 1)
-       // No extra sort before last sort merge join
--      assert(collect(plan) { case _: SortExec => true }.size === 3)
-+      assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 3)
-     })
- 
-     // Test shuffled hash join
-@@ -1146,10 +1151,13 @@ class JoinSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlan
-         .join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
-         .queryExecution
-         .executedPlan
--      assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 2)
--      assert(collect(plan) { case _: ShuffledHashJoinExec => true }.size === 1)
-+      assert(collect(plan) {
-+        case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 2)
-+      assert(collect(plan) {
-+        case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
-       // No extra sort before last sort merge join
--      assert(collect(plan) { case _: SortExec => true }.size === 3)
-+      assert(collect(plan) {
-+        case _: SortExec | _: CometSortExec => true }.size === 3)
-     })
-   }
- 
-@@ -1240,12 +1248,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) {
--        case _: SortMergeJoinExec => true }.size === 1)
-+        case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 1)
-       val smjResult = smjDF.collect()
- 
-       val shjDF = df1.join(df2.hint("SHUFFLE_HASH"), joinExprs, "full")
-       assert(collect(shjDF.queryExecution.executedPlan) {
--        case _: ShuffledHashJoinExec => true }.size === 1)
-+        case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
-       // Same result between shuffled hash join and sort merge join
-       checkAnswer(shjDF, smjResult)
-     }
-@@ -1340,7 +1348,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
--          assert(collect(plan) { case _: ShuffleExchangeExec => true }.size === 1)
-+          assert(collect(plan) {
-+            case _: ShuffleExchangeLike => true }.size === 1)
-       }
- 
-       def getJoinQuery(selectExpr: String, joinType: String): String = {
-@@ -1369,9 +1378,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)
--          assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 3)
-+          assert(collect(plan) {
-+            case _: SortMergeJoinExec => true
-+            case _: CometSortMergeJoinExec => true
-+          }.size === 3)
-           // No extra sort on left side before last sort merge join
--          assert(collect(plan) { case _: SortExec => true }.size === 5)
-+          assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 5)
-       }
- 
-       // Test output ordering is not preserved
-@@ -1380,9 +1392,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)
--          assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 3)
-+          assert(collect(plan) {
-+            case _: SortMergeJoinExec => true
-+            case _: CometSortMergeJoinExec => true
-+          }.size === 3)
-           // Have sort on left side before last sort merge join
--          assert(collect(plan) { case _: SortExec => true }.size === 6)
-+          assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 6)
-       }
- 
-       // Test singe partition
-@@ -1392,7 +1407,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
--      assert(collect(plan) { case _: ShuffleExchangeExec => true}.size == 1)
-+      assert(collect(plan) {
-+        case _: ShuffleExchangeLike => true}.size == 1)
-       checkAnswer(fullJoinDF, Row(100))
-     }
-   }
-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 b5b34922694..a72403780c4 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
-  * }}}
-  */
- // scalastyle:on line.size.limit
--trait PlanStabilitySuite extends DisableAdaptiveExecutionSuite {
-+trait PlanStabilitySuite extends DisableAdaptiveExecutionSuite with IgnoreCometSuite {
- 
-   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 525d97e4998..8a3e7457618 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
-@@ -1508,7 +1508,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
-     checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001")))
-   }
- 
--  test("external sorting updates peak execution memory") {
-+  test("external sorting updates peak execution memory",
-+    IgnoreComet("TODO: native CometSort does not update peak execution memory")) {
-     AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") {
-       sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
-     }
-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 3cfda19134a..7590b808def 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
-@@ -21,10 +21,11 @@ import scala.collection.mutable.ArrayBuffer
- 
- import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
- import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, LogicalPlan, Project, Sort, Union}
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.execution._
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
- import org.apache.spark.sql.execution.datasources.FileScanRDD
--import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
- import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
- import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.test.SharedSparkSession
-@@ -1543,6 +1544,12 @@ class SubquerySuite extends QueryTest
-             fs.inputRDDs().forall(
-               _.asInstanceOf[FileScanRDD].filePartitions.forall(
-                 _.files.forall(_.urlEncodedPath.contains("p=0"))))
-+        case WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(
-+        fs @ CometScanExec(_, _, _, partitionFilters, _, _, _, _, _, _)))) =>
-+          partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
-+            fs.inputRDDs().forall(
-+              _.asInstanceOf[FileScanRDD].filePartitions.forall(
-+                _.files.forall(_.urlEncodedPath.contains("p=0"))))
-         case _ => false
-       })
-     }
-@@ -2108,7 +2115,7 @@ class SubquerySuite extends QueryTest
- 
-       df.collect()
-       val exchanges = collect(df.queryExecution.executedPlan) {
--        case s: ShuffleExchangeExec => s
-+        case s: ShuffleExchangeLike => s
-       }
-       assert(exchanges.size === 1)
-     }
-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 02990a7a40d..bddf5e1ccc2 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
-@@ -24,6 +24,7 @@ import test.org.apache.spark.sql.connector._
- 
- import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row}
- import org.apache.spark.sql.catalyst.InternalRow
-+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}
-@@ -33,7 +34,7 @@ import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning,
- 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.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._
-@@ -268,13 +269,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) {
--            case e: ShuffleExchangeExec => e
-+            case e: ShuffleExchangeLike => e
-           }.isDefined)
- 
-           val groupByIPlusJ = df.groupBy($"i" + $"j").agg(count("*"))
-           checkAnswer(groupByIPlusJ, Seq(Row(5, 2), Row(6, 2), Row(8, 1), Row(9, 1)))
-           assert(collectFirst(groupByIPlusJ.queryExecution.executedPlan) {
--            case e: ShuffleExchangeExec => e
-+            case e: ShuffleExchangeLike => e
-           }.isDefined)
-         }
-       }
-@@ -334,10 +335,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
- 
-                 val (shuffleExpected, sortExpected) = groupByExpects
-                 assert(collectFirst(groupBy.queryExecution.executedPlan) {
--                  case e: ShuffleExchangeExec => e
-+                  case e: ShuffleExchangeLike => e
-                 }.isDefined === shuffleExpected)
-                 assert(collectFirst(groupBy.queryExecution.executedPlan) {
-                   case e: SortExec => e
-+                  case c: CometSortExec => c
-                 }.isDefined === sortExpected)
-               }
- 
-@@ -352,10 +354,11 @@ class DataSourceV2Suite extends QueryTest with SharedSparkSession with AdaptiveS
- 
-                 val (shuffleExpected, sortExpected) = windowFuncExpects
-                 assert(collectFirst(windowPartByColIOrderByColJ.queryExecution.executedPlan) {
--                  case e: ShuffleExchangeExec => e
-+                  case e: ShuffleExchangeLike => e
-                 }.isDefined === shuffleExpected)
-                 assert(collectFirst(windowPartByColIOrderByColJ.queryExecution.executedPlan) {
-                   case e: SortExec => e
-+                  case c: CometSortExec => c
-                 }.isDefined === sortExpected)
-               }
-             }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
-index cfc8b2cc845..c6fcfd7bd08 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
-@@ -21,6 +21,7 @@ import scala.collection.mutable.ArrayBuffer
- import org.apache.spark.SparkConf
- import org.apache.spark.sql.{AnalysisException, QueryTest}
- import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability}
- import org.apache.spark.sql.connector.read.ScanBuilder
- import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder}
-@@ -184,7 +185,11 @@ class FileDataSourceV2FallBackSuite extends QueryTest with SharedSparkSession {
-             val df = spark.read.format(format).load(path.getCanonicalPath)
-             checkAnswer(df, inputData.toDF())
-             assert(
--              df.queryExecution.executedPlan.exists(_.isInstanceOf[FileSourceScanExec]))
-+              df.queryExecution.executedPlan.exists {
-+                case _: FileSourceScanExec | _: CometScanExec => true
-+                case _ => false
-+              }
-+            )
-           }
-         } 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 cf76f6ca32c..f454128af06 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
-@@ -22,6 +22,7 @@ import org.apache.spark.sql.{DataFrame, Row}
- import org.apache.spark.sql.catalyst.InternalRow
- import org.apache.spark.sql.catalyst.expressions.{Literal, TransformExpression}
- import org.apache.spark.sql.catalyst.plans.physical
-+import org.apache.spark.sql.comet.CometSortMergeJoinExec
- import org.apache.spark.sql.connector.catalog.Identifier
- import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
- import org.apache.spark.sql.connector.catalog.functions._
-@@ -31,7 +32,7 @@ import org.apache.spark.sql.connector.expressions.Expressions._
- import org.apache.spark.sql.execution.SparkPlan
- import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
- import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation
--import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
- import org.apache.spark.sql.execution.joins.SortMergeJoinExec
- import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.internal.SQLConf._
-@@ -279,13 +280,14 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
-         Row("bbb", 20, 250.0), Row("bbb", 20, 350.0), Row("ccc", 30, 400.50)))
-   }
- 
--  private def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeExec] = {
-+  private def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = {
-     // here we skip collecting shuffle operators that are not associated with SMJ
-     collect(plan) {
-       case s: SortMergeJoinExec => s
-+      case c: CometSortMergeJoinExec => c.originalPlan
-     }.flatMap(smj =>
-       collect(smj) {
--        case s: ShuffleExchangeExec => s
-+        case s: ShuffleExchangeLike => s
-       })
-   }
- 
-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 c0ec8a58bd5..4e8bc6ed3c5 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
-@@ -27,7 +27,7 @@ import org.apache.hadoop.fs.permission.FsPermission
- import org.mockito.Mockito.{mock, spy, when}
- 
- import org.apache.spark._
--import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, QueryTest, Row, SaveMode}
-+import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, IgnoreComet, QueryTest, Row, SaveMode}
- import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
- import org.apache.spark.sql.catalyst.util.BadRecordException
- import org.apache.spark.sql.execution.datasources.jdbc.{DriverRegistry, JDBCOptions}
-@@ -248,7 +248,8 @@ class QueryExecutionErrorsSuite
-   }
- 
-   test("INCONSISTENT_BEHAVIOR_CROSS_VERSION: " +
--    "compatibility with Spark 2.4/3.2 in reading/writing dates") {
-+    "compatibility with Spark 2.4/3.2 in reading/writing dates",
-+    IgnoreComet("Comet doesn't completely support datetime rebase mode yet")) {
- 
-     // Fail to read ancient datetime values.
-     withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) {
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
-index 418ca3430bb..eb8267192f8 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
-@@ -23,7 +23,7 @@ import scala.util.Random
- import org.apache.hadoop.fs.Path
- 
- import org.apache.spark.SparkConf
--import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.{DataFrame, IgnoreComet, QueryTest}
- import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
- import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan
- import org.apache.spark.sql.internal.SQLConf
-@@ -195,7 +195,7 @@ class DataSourceV2ScanExecRedactionSuite extends DataSourceScanRedactionTest {
-     }
-   }
- 
--  test("FileScan description") {
-+  test("FileScan description", IgnoreComet("Comet doesn't use BatchScan")) {
-     Seq("json", "orc", "parquet").foreach { format =>
-       withTempPath { path =>
-         val dir = path.getCanonicalPath
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
-index 743ec41dbe7..9f30d6c8e04 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
-@@ -53,6 +53,10 @@ class LogicalPlanTagInSparkPlanSuite extends TPCDSQuerySuite with DisableAdaptiv
-     case ColumnarToRowExec(i: InputAdapter) => isScanPlanTree(i.child)
-     case p: ProjectExec => isScanPlanTree(p.child)
-     case f: FilterExec => isScanPlanTree(f.child)
-+    // Comet produces scan plan tree like:
-+    // ColumnarToRow
-+    //  +- ReusedExchange
-+    case _: ReusedExchangeExec => false
-     case _: LeafExecNode => true
-     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 4b3d3a4b805..56e1e0e6f16 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
-@@ -18,7 +18,7 @@
- package org.apache.spark.sql.execution
- 
- 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.catalyst.InternalRow
- import org.apache.spark.sql.catalyst.expressions._
- import org.apache.spark.sql.catalyst.plans._
-@@ -35,7 +35,9 @@ import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.test.SharedSparkSession
- import org.apache.spark.sql.types._
- 
--class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
-+// Ignore this suite when Comet is enabled. This suite tests the Spark planner and Comet planner
-+// comes out with too many difference. Simply ignoring this suite for now.
-+class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper with IgnoreCometSuite {
-   import testImplicits._
- 
-   setupTestData()
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
-index 9e9d717db3b..91a4f9a38d5 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
-@@ -18,6 +18,7 @@
- package org.apache.spark.sql.execution
- 
- import org.apache.spark.sql.{DataFrame, QueryTest, Row}
-+import org.apache.spark.sql.comet.CometProjectExec
- import org.apache.spark.sql.connector.SimpleWritableDataSource
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
- import org.apache.spark.sql.internal.SQLConf
-@@ -34,7 +35,10 @@ abstract class RemoveRedundantProjectsSuiteBase
-   private def assertProjectExecCount(df: DataFrame, expected: Int): Unit = {
-     withClue(df.queryExecution) {
-       val plan = df.queryExecution.executedPlan
--      val actual = collectWithSubqueries(plan) { case p: ProjectExec => p }.size
-+      val actual = collectWithSubqueries(plan) {
-+        case p: ProjectExec => p
-+        case p: CometProjectExec => p
-+      }.size
-       assert(actual == expected)
-     }
-   }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
-index 30ce940b032..0d3f6c6c934 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
-@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
- 
- import org.apache.spark.sql.{DataFrame, QueryTest}
- import org.apache.spark.sql.catalyst.plans.physical.{RangePartitioning, UnknownPartitioning}
-+import org.apache.spark.sql.comet.CometSortExec
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
- import org.apache.spark.sql.execution.joins.ShuffledJoin
- import org.apache.spark.sql.internal.SQLConf
-@@ -33,7 +34,7 @@ abstract class RemoveRedundantSortsSuiteBase
- 
-   private def checkNumSorts(df: DataFrame, count: Int): Unit = {
-     val plan = df.queryExecution.executedPlan
--    assert(collectWithSubqueries(plan) { case s: SortExec => s }.length == count)
-+    assert(collectWithSubqueries(plan) { case _: SortExec | _: CometSortExec => 1 }.length == count)
-   }
- 
-   private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-index 47679ed7865..9ffbaecb98e 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
-@@ -18,6 +18,7 @@
- package org.apache.spark.sql.execution
- 
- import org.apache.spark.sql.{DataFrame, QueryTest}
-+import org.apache.spark.sql.comet.CometHashAggregateExec
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
- import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
- import org.apache.spark.sql.internal.SQLConf
-@@ -31,7 +32,7 @@ abstract class ReplaceHashWithSortAggSuiteBase
-   private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
-     val plan = df.queryExecution.executedPlan
-     assert(collectWithSubqueries(plan) {
--      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
-+      case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
-     }.length == hashAggCount)
-     assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
-   }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
-index ac710c32296..e163c1a6a76 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
-@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
- 
- import org.apache.spark.sql.{Dataset, QueryTest, Row, SaveMode}
- import org.apache.spark.sql.catalyst.expressions.codegen.{ByteCodeStats, CodeAndComment, CodeGenerator}
-+import org.apache.spark.sql.comet.CometSortMergeJoinExec
- import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecutionSuite
- import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
- import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
-@@ -224,6 +225,7 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-       assert(twoJoinsDF.queryExecution.executedPlan.collect {
-         case WholeStageCodegenExec(_ : ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
-         case WholeStageCodegenExec(_ : SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
-+        case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
-       }.size === 2)
-       checkAnswer(twoJoinsDF,
-         Seq(Row(0, 0, 0), Row(1, 1, null), Row(2, 2, 2), Row(3, 3, null), Row(4, 4, null),
-@@ -258,6 +260,7 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-       .join(df1.hint("SHUFFLE_MERGE"), $"k3" === $"k1", "right_outer")
-     assert(twoJoinsDF.queryExecution.executedPlan.collect {
-       case WholeStageCodegenExec(_ : SortMergeJoinExec) => true
-+      case _: CometSortMergeJoinExec => true
-     }.size === 2)
-     checkAnswer(twoJoinsDF,
-       Seq(Row(0, 0, 0), Row(1, 1, 1), Row(2, 2, 2), Row(3, 3, 3), Row(4, null, 4), Row(5, null, 5),
-@@ -280,8 +283,7 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-     val twoJoinsDF = df3.join(df2.hint("SHUFFLE_MERGE"), $"k3" === $"k2", "left_semi")
-       .join(df1.hint("SHUFFLE_MERGE"), $"k3" === $"k1", "left_semi")
-     assert(twoJoinsDF.queryExecution.executedPlan.collect {
--      case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) |
--           WholeStageCodegenExec(_ : SortMergeJoinExec) => true
-+      case _: SortMergeJoinExec => true
-     }.size === 2)
-     checkAnswer(twoJoinsDF, Seq(Row(0), Row(1), Row(2), Row(3)))
-   }
-@@ -302,8 +304,7 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-     val twoJoinsDF = df1.join(df2.hint("SHUFFLE_MERGE"), $"k1" === $"k2", "left_anti")
-       .join(df3.hint("SHUFFLE_MERGE"), $"k1" === $"k3", "left_anti")
-     assert(twoJoinsDF.queryExecution.executedPlan.collect {
--      case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) |
--           WholeStageCodegenExec(_ : SortMergeJoinExec) => true
-+      case _: SortMergeJoinExec => true
-     }.size === 2)
-     checkAnswer(twoJoinsDF, Seq(Row(6), Row(7), Row(8), Row(9)))
-   }
-@@ -436,7 +437,9 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-     val plan = df.queryExecution.executedPlan
-     assert(plan.exists(p =>
-       p.isInstanceOf[WholeStageCodegenExec] &&
--        p.asInstanceOf[WholeStageCodegenExec].child.isInstanceOf[SortExec]))
-+        p.asInstanceOf[WholeStageCodegenExec].collect {
-+          case _: SortExec => true
-+        }.nonEmpty))
-     assert(df.collect() === Array(Row(1), Row(2), Row(3)))
-   }
- 
-@@ -616,7 +619,9 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
-           .write.mode(SaveMode.Overwrite).parquet(path)
- 
-         withSQLConf(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "255",
--            SQLConf.WHOLESTAGE_SPLIT_CONSUME_FUNC_BY_OPERATOR.key -> "true") {
-+            SQLConf.WHOLESTAGE_SPLIT_CONSUME_FUNC_BY_OPERATOR.key -> "true",
-+            // Disable Comet native execution because this checks wholestage codegen.
-+            "spark.comet.exec.enabled" -> "false") {
-           val projection = Seq.tabulate(columnNum)(i => s"c$i + c$i as newC$i")
-           val df = spark.read.parquet(path).selectExpr(projection: _*)
- 
-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 593bd7bb4ba..7ad55e3ab20 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
-@@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._
- 
- import org.apache.spark.SparkException
- import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
--import org.apache.spark.sql.{Dataset, QueryTest, Row, SparkSession, Strategy}
-+import org.apache.spark.sql.{Dataset, IgnoreComet, QueryTest, Row, SparkSession, Strategy}
- import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
- import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
-+import org.apache.spark.sql.comet._
-+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
- import org.apache.spark.sql.execution.{CollectLimitExec, LocalTableScanExec, PartialReducerPartitionSpec, QueryExecution, ReusedSubqueryExec, ShuffledRowRDD, SortExec, SparkPlan, SparkPlanInfo, UnionExec}
- import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
- import org.apache.spark.sql.execution.command.DataWritingCommandExec
-@@ -104,6 +106,7 @@ class AdaptiveQueryExecSuite
-   private def findTopLevelBroadcastHashJoin(plan: SparkPlan): Seq[BroadcastHashJoinExec] = {
-     collect(plan) {
-       case j: BroadcastHashJoinExec => j
-+      case j: CometBroadcastHashJoinExec => j.originalPlan.asInstanceOf[BroadcastHashJoinExec]
-     }
-   }
- 
-@@ -116,30 +119,38 @@ class AdaptiveQueryExecSuite
-   private def findTopLevelSortMergeJoin(plan: SparkPlan): Seq[SortMergeJoinExec] = {
-     collect(plan) {
-       case j: SortMergeJoinExec => j
-+      case j: CometSortMergeJoinExec =>
-+        assert(j.originalPlan.isInstanceOf[SortMergeJoinExec])
-+        j.originalPlan.asInstanceOf[SortMergeJoinExec]
-     }
-   }
- 
-   private def findTopLevelShuffledHashJoin(plan: SparkPlan): Seq[ShuffledHashJoinExec] = {
-     collect(plan) {
-       case j: ShuffledHashJoinExec => j
-+      case j: CometHashJoinExec => j.originalPlan.asInstanceOf[ShuffledHashJoinExec]
-     }
-   }
- 
-   private def findTopLevelBaseJoin(plan: SparkPlan): Seq[BaseJoinExec] = {
-     collect(plan) {
-       case j: BaseJoinExec => j
-+      case c: CometHashJoinExec => c.originalPlan.asInstanceOf[BaseJoinExec]
-+      case c: CometSortMergeJoinExec => c.originalPlan.asInstanceOf[BaseJoinExec]
-     }
-   }
- 
-   private def findTopLevelSort(plan: SparkPlan): Seq[SortExec] = {
-     collect(plan) {
-       case s: SortExec => s
-+      case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
-     }
-   }
- 
-   private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
-     collect(plan) {
-       case agg: BaseAggregateExec => agg
-+      case agg: CometHashAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
-     }
-   }
- 
-@@ -176,6 +187,7 @@ class AdaptiveQueryExecSuite
-       val parts = rdd.partitions
-       assert(parts.forall(rdd.preferredLocations(_).nonEmpty))
-     }
-+
-     assert(numShuffles === (numLocalReads.length + numShufflesWithoutLocalRead))
-   }
- 
-@@ -184,7 +196,7 @@ class AdaptiveQueryExecSuite
-     val plan = df.queryExecution.executedPlan
-     assert(plan.isInstanceOf[AdaptiveSparkPlanExec])
-     val shuffle = plan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan.collect {
--      case s: ShuffleExchangeExec => s
-+      case s: ShuffleExchangeLike => s
-     }
-     assert(shuffle.size == 1)
-     assert(shuffle(0).outputPartitioning.numPartitions == numPartition)
-@@ -200,7 +212,8 @@ class AdaptiveQueryExecSuite
-       assert(smj.size == 1)
-       val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
-       assert(bhj.size == 1)
--      checkNumLocalShuffleReads(adaptivePlan)
-+      // Comet shuffle changes shuffle metrics
-+      // checkNumLocalShuffleReads(adaptivePlan)
-     }
-   }
- 
-@@ -227,7 +240,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Reuse the parallelism of coalesced shuffle in local shuffle read") {
-+  test("Reuse the parallelism of coalesced shuffle in local shuffle read",
-+      IgnoreComet("Comet shuffle changes shuffle partition size")) {
-     withSQLConf(
-       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
-@@ -259,7 +273,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Reuse the default parallelism in local shuffle read") {
-+  test("Reuse the default parallelism in local shuffle read",
-+      IgnoreComet("Comet shuffle changes shuffle partition size")) {
-     withSQLConf(
-       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
-@@ -273,7 +288,8 @@ class AdaptiveQueryExecSuite
-       val localReads = collect(adaptivePlan) {
-         case read: AQEShuffleReadExec if read.isLocalRead => read
-       }
--      assert(localReads.length == 2)
-+      // Comet shuffle changes shuffle metrics
-+      assert(localReads.length == 1)
-       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
-@@ -322,7 +338,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Scalar subquery") {
-+  test("Scalar subquery", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -337,7 +353,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Scalar subquery in later stages") {
-+  test("Scalar subquery in later stages", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -353,7 +369,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("multiple joins") {
-+  test("multiple joins", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -398,7 +414,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("multiple joins with aggregate") {
-+  test("multiple joins with aggregate", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -443,7 +459,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("multiple joins with aggregate 2") {
-+  test("multiple joins with aggregate 2", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
-@@ -508,7 +524,7 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Exchange reuse with subqueries") {
-+  test("Exchange reuse with subqueries", IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
-@@ -539,7 +555,9 @@ class AdaptiveQueryExecSuite
-       assert(smj.size == 1)
-       val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
-       assert(bhj.size == 1)
--      checkNumLocalShuffleReads(adaptivePlan)
-+      // Comet shuffle changes shuffle metrics,
-+      // so we can't check the number of local shuffle reads.
-+      // checkNumLocalShuffleReads(adaptivePlan)
-       // Even with local shuffle read, the query stage reuse can also work.
-       val ex = findReusedExchange(adaptivePlan)
-       assert(ex.nonEmpty)
-@@ -560,7 +578,9 @@ class AdaptiveQueryExecSuite
-       assert(smj.size == 1)
-       val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
-       assert(bhj.size == 1)
--      checkNumLocalShuffleReads(adaptivePlan)
-+      // Comet shuffle changes shuffle metrics,
-+      // so we can't check the number of local shuffle reads.
-+      // checkNumLocalShuffleReads(adaptivePlan)
-       // Even with local shuffle read, the query stage reuse can also work.
-       val ex = findReusedExchange(adaptivePlan)
-       assert(ex.isEmpty)
-@@ -569,7 +589,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("Broadcast exchange reuse across subqueries") {
-+  test("Broadcast exchange reuse across subqueries",
-+      IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-         SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000",
-@@ -664,7 +685,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.
--      checkNumLocalShuffleReads(adaptivePlan, 2)
-+      // Comet shuffle changes shuffle metrics
-+      // checkNumLocalShuffleReads(adaptivePlan, 2)
-     }
-   }
- 
-@@ -786,7 +808,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-29544: adaptive skew join with different join types") {
-+  test("SPARK-29544: adaptive skew join with different join types",
-+      IgnoreComet("Comet shuffle has different partition metrics")) {
-     Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint =>
-       def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") {
-         findTopLevelSortMergeJoin(plan)
-@@ -1004,7 +1027,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("metrics of the shuffle read") {
-+  test("metrics of the shuffle read",
-+      IgnoreComet("Comet shuffle changes the metrics")) {
-     withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
-       val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
-         "SELECT key FROM testData GROUP BY key")
-@@ -1599,7 +1623,7 @@ class AdaptiveQueryExecSuite
-         val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
-           "SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id")
-         assert(collect(adaptivePlan) {
--          case s: ShuffleExchangeExec => s
-+          case s: ShuffleExchangeLike => s
-         }.length == 1)
-       }
-     }
-@@ -1679,7 +1703,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-33551: Do not use AQE shuffle read for repartition") {
-+  test("SPARK-33551: Do not use AQE shuffle read for repartition",
-+      IgnoreComet("Comet shuffle changes partition size")) {
-     def hasRepartitionShuffle(plan: SparkPlan): Boolean = {
-       find(plan) {
-         case s: ShuffleExchangeLike =>
-@@ -1864,6 +1889,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
-+        case c: CometShuffleExchangeExec
-+          if c.originalPlan.shuffleOrigin == origin &&
-+            c.originalPlan.numPartitions == 2 => c
-       }.size == 1)
-       ds.collect()
-       val plan = ds.queryExecution.executedPlan
-@@ -1872,6 +1900,9 @@ class AdaptiveQueryExecSuite
-       }.isEmpty)
-       assert(collect(plan) {
-         case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
-+        case c: CometShuffleExchangeExec
-+          if c.originalPlan.shuffleOrigin == origin &&
-+            c.originalPlan.numPartitions == 2 => c
-       }.size == 1)
-       checkAnswer(ds, testData)
-     }
-@@ -2028,7 +2059,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-35264: Support AQE side shuffled hash join formula") {
-+  test("SPARK-35264: Support AQE side shuffled hash join formula",
-+      IgnoreComet("Comet shuffle changes the partition size")) {
-     withTempView("t1", "t2") {
-       def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
-         Seq("100", "100000").foreach { size =>
-@@ -2114,7 +2146,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-35725: Support optimize skewed partitions in RebalancePartitions") {
-+  test("SPARK-35725: Support optimize skewed partitions in RebalancePartitions",
-+      IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withTempView("v") {
-       withSQLConf(
-         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-@@ -2213,7 +2246,7 @@ class AdaptiveQueryExecSuite
-               runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
-                 s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
-             val shuffles1 = collect(adaptive1) {
--              case s: ShuffleExchangeExec => s
-+              case s: ShuffleExchangeLike => s
-             }
-             assert(shuffles1.size == 3)
-             // shuffles1.head is the top-level shuffle under the Aggregate operator
-@@ -2226,7 +2259,7 @@ class AdaptiveQueryExecSuite
-               runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
-                 s"JOIN skewData2 ON key1 = key2")
-             val shuffles2 = collect(adaptive2) {
--              case s: ShuffleExchangeExec => s
-+              case s: ShuffleExchangeLike => s
-             }
-             if (hasRequiredDistribution) {
-               assert(shuffles2.size == 3)
-@@ -2260,7 +2293,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-35794: Allow custom plugin for cost evaluator") {
-+  test("SPARK-35794: Allow custom plugin for cost evaluator",
-+      IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     CostEvaluator.instantiate(
-       classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
-     intercept[IllegalArgumentException] {
-@@ -2404,6 +2438,7 @@ class AdaptiveQueryExecSuite
-           val (_, adaptive) = runAdaptiveAndVerifyResult(query)
-           assert(adaptive.collect {
-             case sort: SortExec => sort
-+            case sort: CometSortExec => sort
-           }.size == 1)
-           val read = collect(adaptive) {
-             case read: AQEShuffleReadExec => read
-@@ -2421,7 +2456,8 @@ class AdaptiveQueryExecSuite
-     }
-   }
- 
--  test("SPARK-37357: Add small partition factor for rebalance partitions") {
-+  test("SPARK-37357: Add small partition factor for rebalance partitions",
-+      IgnoreComet("Comet shuffle changes shuffle metrics")) {
-     withTempView("v") {
-       withSQLConf(
-         SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
-@@ -2533,7 +2569,7 @@ class AdaptiveQueryExecSuite
-           runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
-             "JOIN skewData3 ON value2 = value3")
-         val shuffles1 = collect(adaptive1) {
--          case s: ShuffleExchangeExec => s
-+          case s: ShuffleExchangeLike => s
-         }
-         assert(shuffles1.size == 4)
-         val smj1 = findTopLevelSortMergeJoin(adaptive1)
-@@ -2544,7 +2580,7 @@ class AdaptiveQueryExecSuite
-           runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
-             "JOIN skewData3 ON value1 = value3")
-         val shuffles2 = collect(adaptive2) {
--          case s: ShuffleExchangeExec => s
-+          case s: ShuffleExchangeLike => s
-         }
-         assert(shuffles2.size == 4)
-         val smj2 = findTopLevelSortMergeJoin(adaptive2)
-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 bd9c79e5b96..ab7584e768e 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
-@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.SchemaPruningTest
- import org.apache.spark.sql.catalyst.expressions.Concat
- import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
- import org.apache.spark.sql.catalyst.plans.logical.Expand
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.execution.FileSourceScanExec
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.functions._
-@@ -867,6 +868,7 @@ abstract class SchemaPruningSuite
-     val fileSourceScanSchemata =
-       collect(df.queryExecution.executedPlan) {
-         case scan: FileSourceScanExec => scan.requiredSchema
-+        case scan: CometScanExec => scan.requiredSchema
-       }
-     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 ce43edb79c1..c414b19eda7 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,7 +17,7 @@
- 
- package org.apache.spark.sql.execution.datasources
- 
--import org.apache.spark.sql.{QueryTest, Row}
-+import org.apache.spark.sql.{IgnoreComet, QueryTest, Row}
- import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, NullsFirst, SortOrder}
- import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Sort}
- import org.apache.spark.sql.execution.{QueryExecution, SortExec}
-@@ -305,7 +305,8 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
-     }
-   }
- 
--  test("v1 write with AQE changing SMJ to BHJ") {
-+  test("v1 write with AQE changing SMJ to BHJ",
-+      IgnoreComet("TODO: Comet SMJ to BHJ by AQE")) {
-     withPlannedWrite { enabled =>
-       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 1d2e467c94c..3ea82cd1a3f 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}
- import org.mockito.Mockito.{mock, when}
- 
- import org.apache.spark.SparkException
--import org.apache.spark.sql.{DataFrame, QueryTest, Row}
-+import org.apache.spark.sql.{DataFrame, IgnoreCometSuite, QueryTest, Row}
- import org.apache.spark.sql.catalyst.encoders.RowEncoder
- import org.apache.spark.sql.execution.datasources.PartitionedFile
- import org.apache.spark.sql.functions.col
-@@ -38,7 +38,9 @@ import org.apache.spark.sql.test.SharedSparkSession
- import org.apache.spark.sql.types._
- import org.apache.spark.util.Utils
- 
--class BinaryFileFormatSuite extends QueryTest with SharedSparkSession {
-+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
-+// Since it isn't related to Comet, we disable it for now.
-+class BinaryFileFormatSuite extends QueryTest with SharedSparkSession with IgnoreCometSuite {
-   import BinaryFileFormat._
- 
-   private var testDir: String = _
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
-index 07e2849ce6f..3e73645b638 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
-@@ -28,7 +28,7 @@ import org.apache.parquet.hadoop.ParquetOutputFormat
- 
- import org.apache.spark.TestUtils
- import org.apache.spark.memory.MemoryMode
--import org.apache.spark.sql.Row
-+import org.apache.spark.sql.{IgnoreComet, Row}
- import org.apache.spark.sql.catalyst.util.DateTimeUtils
- import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.test.SharedSparkSession
-@@ -201,7 +201,8 @@ class ParquetEncodingSuite extends ParquetCompatibilityTest with SharedSparkSess
-     }
-   }
- 
--  test("parquet v2 pages - rle encoding for boolean value columns") {
-+  test("parquet v2 pages - rle encoding for boolean value columns",
-+      IgnoreComet("Comet doesn't support RLE encoding yet")) {
-     val extraOptions = Map[String, String](
-       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 9adcb43c838..c6872c7b24b 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
-@@ -1025,7 +1025,11 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
-           // When a filter is pushed to Parquet, Parquet can apply it to every row.
-           // So, we can check the number of rows returned from the Parquet
-           // to make sure our filter pushdown work.
--          assert(stripSparkFilter(df).count == 1)
-+          // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
-+          // on Spark to apply the data filters after columnar batches are returned
-+          if (!isCometEnabled) {
-+            assert(stripSparkFilter(df).count == 1)
-+          }
-         }
-       }
-     }
-@@ -1510,7 +1514,11 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
-           // than the total length but should not be a single record.
-           // Note that, if record level filtering is enabled, it should be a single record.
-           // If no filter is pushed down to Parquet, it should be the total length of data.
--          assert(actual > 1 && actual < data.length)
-+          // Only enable Comet test iff it's scan only, since with native execution
-+          // `stripSparkFilter` can't remove the native filter
-+          if (!isCometEnabled || isCometScanOnly) {
-+            assert(actual > 1 && actual < data.length)
-+          }
-         }
-       }
-     }
-@@ -1537,7 +1545,11 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
-         // than the total length but should not be a single record.
-         // Note that, if record level filtering is enabled, it should be a single record.
-         // If no filter is pushed down to Parquet, it should be the total length of data.
--        assert(actual > 1 && actual < data.length)
-+        // Only enable Comet test iff it's scan only, since with native execution
-+        // `stripSparkFilter` can't remove the native filter
-+        if (!isCometEnabled || isCometScanOnly) {
-+          assert(actual > 1 && actual < data.length)
-+        }
-       }
-     }
-   }
-@@ -1673,7 +1685,8 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
-     }
-   }
- 
--  test("SPARK-17091: Convert IN predicate to Parquet filter push-down") {
-+  test("SPARK-17091: Convert IN predicate to Parquet filter push-down",
-+      IgnoreComet("IN predicate is not yet supported in Comet, see issue #36")) {
-     val schema = StructType(Seq(
-       StructField("a", IntegerType, nullable = false)
-     ))
-@@ -1914,7 +1927,8 @@ abstract class ParquetFilterSuite extends QueryTest with ParquetTest with Shared
-     }
-   }
- 
--  test("Support Parquet column index") {
-+  test("Support Parquet column index",
-+      IgnoreComet("Comet doesn't support Parquet column index yet")) {
-     // block 1:
-     //                      null count  min                                       max
-     // page-0                         0  0                                         99
-@@ -2206,7 +2220,11 @@ class ParquetV1FilterSuite extends ParquetFilterSuite {
-           assert(pushedParquetFilters.exists(_.getClass === filterClass),
-             s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
- 
--          checker(stripSparkFilter(query), expected)
-+          // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
-+          // on Spark to apply the data filters after columnar batches are returned
-+          if (!isCometEnabled) {
-+            checker(stripSparkFilter(query), expected)
-+          }
-         } else {
-           assert(selectedFilters.isEmpty, "There is filter pushed down")
-         }
-@@ -2266,7 +2284,11 @@ class ParquetV2FilterSuite extends ParquetFilterSuite {
-           assert(pushedParquetFilters.exists(_.getClass === filterClass),
-             s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
- 
--          checker(stripSparkFilter(query), expected)
-+          // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
-+          // on Spark to apply the data filters after columnar batches are returned
-+          if (!isCometEnabled) {
-+            checker(stripSparkFilter(query), expected)
-+          }
- 
-         case _ =>
-           throw new AnalysisException("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 8670d95c65e..b624c3811dd 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
-@@ -1335,7 +1335,8 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSparkSession
-     }
-   }
- 
--  test("SPARK-40128 read DELTA_LENGTH_BYTE_ARRAY encoded strings") {
-+  test("SPARK-40128 read DELTA_LENGTH_BYTE_ARRAY encoded strings",
-+      IgnoreComet("Comet doesn't support DELTA encoding yet")) {
-     withAllParquetReaders {
-       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 2e7b26126d2..f7368eb026e 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
-@@ -1029,7 +1029,8 @@ abstract class ParquetQuerySuite extends QueryTest with ParquetTest with SharedS
-         checkAnswer(readParquet(schema, path), df)
-       }
- 
--      withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
-+      withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false",
-+          "spark.comet.enabled" -> "false") {
-         val schema1 = "a DECIMAL(3, 2), b DECIMAL(18, 3), c DECIMAL(37, 3)"
-         checkAnswer(readParquet(schema1, path), df)
-         val schema2 = "a DECIMAL(3, 0), b DECIMAL(18, 1), c DECIMAL(37, 1)"
-@@ -1051,7 +1052,8 @@ abstract class ParquetQuerySuite extends QueryTest with ParquetTest with SharedS
-       val df = sql(s"SELECT 1 a, 123456 b, ${Int.MaxValue.toLong * 10} c, CAST('1.2' AS BINARY) d")
-       df.write.parquet(path.toString)
- 
--      withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
-+      withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false",
-+          "spark.comet.enabled" -> "false") {
-         checkAnswer(readParquet("a DECIMAL(3, 2)", path), sql("SELECT 1.00"))
-         checkAnswer(readParquet("b DECIMAL(3, 2)", path), Row(null))
-         checkAnswer(readParquet("b DECIMAL(11, 1)", path), sql("SELECT 123456.0"))
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
-index 240bb4e6dcb..8287ffa03ca 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
-@@ -21,7 +21,7 @@ import java.nio.file.{Files, Paths, StandardCopyOption}
- import java.sql.{Date, Timestamp}
- 
- import org.apache.spark.{SPARK_VERSION_SHORT, SparkConf, SparkException, SparkUpgradeException}
--import org.apache.spark.sql.{QueryTest, Row, SPARK_LEGACY_DATETIME_METADATA_KEY, SPARK_LEGACY_INT96_METADATA_KEY, SPARK_TIMEZONE_METADATA_KEY}
-+import org.apache.spark.sql.{IgnoreCometSuite, QueryTest, Row, SPARK_LEGACY_DATETIME_METADATA_KEY, SPARK_LEGACY_INT96_METADATA_KEY, SPARK_TIMEZONE_METADATA_KEY}
- import org.apache.spark.sql.catalyst.util.DateTimeTestUtils
- import org.apache.spark.sql.internal.SQLConf
- import org.apache.spark.sql.internal.SQLConf.{LegacyBehaviorPolicy, ParquetOutputTimestampType}
-@@ -30,9 +30,11 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType.{INT96,
- import org.apache.spark.sql.test.SharedSparkSession
- import org.apache.spark.tags.SlowSQLTest
- 
-+// Comet is disabled for this suite because it doesn't support datetime rebase mode
- abstract class ParquetRebaseDatetimeSuite
-   extends QueryTest
-   with ParquetTest
-+  with IgnoreCometSuite
-   with SharedSparkSession {
- 
-   import testImplicits._
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
-index 351c6d698fc..36492fe936d 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
-@@ -26,6 +26,7 @@ import org.apache.parquet.hadoop.{ParquetFileReader, ParquetOutputFormat}
- import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE
- 
- import org.apache.spark.sql.QueryTest
-+import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec}
- import org.apache.spark.sql.execution.FileSourceScanExec
- import org.apache.spark.sql.execution.datasources.FileFormat
- import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
-@@ -230,6 +231,12 @@ class ParquetRowIndexSuite extends QueryTest with SharedSparkSession {
-             case f: FileSourceScanExec =>
-               numPartitions += f.inputRDD.partitions.length
-               numOutputRows += f.metrics("numOutputRows").value
-+            case b: CometScanExec =>
-+              numPartitions += b.inputRDD.partitions.length
-+              numOutputRows += b.metrics("numOutputRows").value
-+            case b: CometBatchScanExec =>
-+              numPartitions += b.inputRDD.partitions.length
-+              numOutputRows += b.metrics("numOutputRows").value
-             case _ =>
-           }
-           assert(numPartitions > 0)
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
-index 5c0b7def039..151184bc98c 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
-@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.datasources.parquet
- import org.apache.spark.SparkConf
- import org.apache.spark.sql.DataFrame
- import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
-+import org.apache.spark.sql.comet.CometBatchScanExec
- import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
- import org.apache.spark.sql.execution.datasources.SchemaPruningSuite
- import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
-@@ -56,6 +57,7 @@ class ParquetV2SchemaPruningSuite extends ParquetSchemaPruningSuite {
-     val fileSourceScanSchemata =
-       collect(df.queryExecution.executedPlan) {
-         case scan: BatchScanExec => scan.scan.asInstanceOf[ParquetScan].readDataSchema
-+        case scan: CometBatchScanExec => scan.scan.asInstanceOf[ParquetScan].readDataSchema
-       }
-     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 bf5c51b89bb..ca22370ca3b 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,6 +27,7 @@ import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
- import org.apache.parquet.schema.Type._
- 
- import org.apache.spark.SparkException
-+import org.apache.spark.sql.IgnoreComet
- import org.apache.spark.sql.catalyst.ScalaReflection
- import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
- import org.apache.spark.sql.functions.desc
-@@ -1016,7 +1017,8 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
-     e
-   }
- 
--  test("schema mismatch failure error message for parquet reader") {
-+  test("schema mismatch failure error message for parquet reader",
-+      IgnoreComet("Comet doesn't work with vectorizedReaderEnabled = false")) {
-     withTempPath { dir =>
-       val e = testSchemaMismatch(dir.getCanonicalPath, vectorizedReaderEnabled = false)
-       val expectedMessage = "Encountered error while reading file"
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
-index 3a0bd35cb70..b28f06a757f 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
-@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.debug
- import java.io.ByteArrayOutputStream
- 
- import org.apache.spark.rdd.RDD
-+import org.apache.spark.sql.IgnoreComet
- import org.apache.spark.sql.catalyst.InternalRow
- import org.apache.spark.sql.catalyst.expressions.Attribute
- import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
-@@ -124,7 +125,8 @@ class DebuggingSuite extends DebuggingSuiteBase with DisableAdaptiveExecutionSui
-          | id LongType: {}""".stripMargin))
-   }
- 
--  test("SPARK-28537: DebugExec cannot debug columnar related queries") {
-+  test("SPARK-28537: DebugExec cannot debug columnar related queries",
-+      IgnoreComet("Comet does not use FileScan")) {
-     withTempPath { workDir =>
-       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 26e61c6b58d..cb09d7e116a 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
-@@ -45,8 +45,10 @@ import org.apache.spark.sql.util.QueryExecutionListener
- import org.apache.spark.util.{AccumulatorContext, JsonProtocol}
- 
- // Disable AQE because metric info is different with AQE on/off
-+// This test suite runs tests against the metrics of physical operators.
-+// Disabling it for Comet because the metrics are different with Comet enabled.
- class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
--  with DisableAdaptiveExecutionSuite {
-+  with DisableAdaptiveExecutionSuite with IgnoreCometSuite {
-   import testImplicits._
- 
-   /**
-@@ -737,7 +739,8 @@ class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
-     }
-   }
- 
--  test("SPARK-26327: FileSourceScanExec metrics") {
-+  test("SPARK-26327: FileSourceScanExec metrics",
-+      IgnoreComet("Spark uses row-based Parquet reader while Comet is vectorized")) {
-     withTable("testDataForScan") {
-       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
---- 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 @@
- package org.apache.spark.sql.execution.python
- 
- import org.apache.spark.sql.catalyst.plans.logical.{ArrowEvalPython, BatchEvalPython, Limit, LocalLimit}
-+import org.apache.spark.sql.comet._
- import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan, SparkPlanTest}
- import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
- import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
-@@ -108,6 +109,7 @@ class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSparkSession {
- 
-           val scanNodes = query.queryExecution.executedPlan.collect {
-             case scan: FileSourceScanExec => scan
-+            case scan: CometScanExec => scan
-           }
-           assert(scanNodes.length == 1)
-           assert(scanNodes.head.output.map(_.name) == Seq("a"))
-@@ -120,11 +122,16 @@ class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSparkSession {
- 
-           val scanNodes = query.queryExecution.executedPlan.collect {
-             case scan: FileSourceScanExec => scan
-+            case scan: CometScanExec => scan
-           }
-           assert(scanNodes.length == 1)
-           // $"a" is not null and $"a" > 1
--          assert(scanNodes.head.dataFilters.length == 2)
--          assert(scanNodes.head.dataFilters.flatMap(_.references.map(_.name)).distinct == Seq("a"))
-+          val dataFilters = scanNodes.head match {
-+            case scan: FileSourceScanExec => scan.dataFilters
-+            case scan: CometScanExec => scan.dataFilters
-+          }
-+          assert(dataFilters.length == 2)
-+          assert(dataFilters.flatMap(_.references.map(_.name)).distinct == Seq("a"))
-         }
-       }
-     }
-@@ -145,6 +152,7 @@ class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSparkSession {
- 
-           val scanNodes = query.queryExecution.executedPlan.collect {
-             case scan: BatchScanExec => scan
-+            case scan: CometBatchScanExec => scan
-           }
-           assert(scanNodes.length == 1)
-           assert(scanNodes.head.output.map(_.name) == Seq("a"))
-@@ -157,6 +165,7 @@ class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSparkSession {
- 
-           val scanNodes = query.queryExecution.executedPlan.collect {
-             case scan: BatchScanExec => scan
-+            case scan: CometBatchScanExec => scan
-           }
-           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 d083cac48ff..3c11bcde807 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,
- import org.apache.spark.sql.streaming.util.StreamManualClock
- import org.apache.spark.util.Utils
- 
-+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
-+// Since it isn't related to Comet, we disable it for now.
- class AsyncProgressTrackingMicroBatchExecutionSuite
--  extends StreamTest with BeforeAndAfter with Matchers {
-+  extends StreamTest with BeforeAndAfter with Matchers with IgnoreCometSuite {
- 
-   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 266bb343526..a426d8396be 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
-@@ -24,10 +24,11 @@ import org.apache.spark.sql.catalyst.catalog.BucketSpec
- import org.apache.spark.sql.catalyst.expressions
- import org.apache.spark.sql.catalyst.expressions._
- import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
--import org.apache.spark.sql.execution.{FileSourceScanExec, SortExec, SparkPlan}
-+import org.apache.spark.sql.comet._
-+import org.apache.spark.sql.execution.{ColumnarToRowExec, FileSourceScanExec, SortExec, SparkPlan}
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
- import org.apache.spark.sql.execution.datasources.BucketingUtils
--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._
- import org.apache.spark.sql.internal.SQLConf
-@@ -101,12 +102,20 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti
-     }
-   }
- 
--  private def getFileScan(plan: SparkPlan): FileSourceScanExec = {
--    val fileScan = collect(plan) { case f: FileSourceScanExec => f }
-+  private def getFileScan(plan: SparkPlan): SparkPlan = {
-+    val fileScan = collect(plan) {
-+      case f: FileSourceScanExec => f
-+      case f: CometScanExec => f
-+    }
-     assert(fileScan.nonEmpty, plan)
-     fileScan.head
-   }
- 
-+  private def getBucketScan(plan: SparkPlan): Boolean = getFileScan(plan) match {
-+    case fs: FileSourceScanExec => fs.bucketedScan
-+    case bs: CometScanExec => bs.bucketedScan
-+  }
-+
-   // To verify if the bucket pruning works, this function checks two conditions:
-   //   1) Check if the pruned buckets (before filtering) are empty.
-   //   2) Verify the final result is the same as the expected one
-@@ -155,7 +164,8 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti
-           val planWithoutBucketedScan = bucketedDataFrame.filter(filterCondition)
-             .queryExecution.executedPlan
-           val fileScan = getFileScan(planWithoutBucketedScan)
--          assert(!fileScan.bucketedScan, s"except no bucketed scan but found\n$fileScan")
-+          val bucketedScan = getBucketScan(planWithoutBucketedScan)
-+          assert(!bucketedScan, s"except no bucketed scan but found\n$fileScan")
- 
-           val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType
-           val rowsWithInvalidBuckets = fileScan.execute().filter(row => {
-@@ -451,28 +461,44 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti
-         val joinOperator = if (joined.sqlContext.conf.adaptiveExecutionEnabled) {
-           val executedPlan =
-             joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
--          assert(executedPlan.isInstanceOf[SortMergeJoinExec])
--          executedPlan.asInstanceOf[SortMergeJoinExec]
-+          executedPlan match {
-+            case s: SortMergeJoinExec => s
-+            case b: CometSortMergeJoinExec =>
-+              b.originalPlan match {
-+                case s: SortMergeJoinExec => s
-+                case o => fail(s"expected SortMergeJoinExec, but found\n$o")
-+              }
-+            case o => fail(s"expected SortMergeJoinExec, but found\n$o")
-+          }
-         } else {
-           val executedPlan = joined.queryExecution.executedPlan
--          assert(executedPlan.isInstanceOf[SortMergeJoinExec])
--          executedPlan.asInstanceOf[SortMergeJoinExec]
-+          executedPlan match {
-+            case s: SortMergeJoinExec => s
-+            case ColumnarToRowExec(child) =>
-+              child.asInstanceOf[CometSortMergeJoinExec].originalPlan match {
-+                case s: SortMergeJoinExec => s
-+                case o => fail(s"expected SortMergeJoinExec, but found\n$o")
-+              }
-+            case o => fail(s"expected SortMergeJoinExec, but found\n$o")
-+          }
-         }
- 
-         // check existence of shuffle
-         assert(
--          joinOperator.left.exists(_.isInstanceOf[ShuffleExchangeExec]) == shuffleLeft,
-+          joinOperator.left.exists(op => op.isInstanceOf[ShuffleExchangeLike]) == shuffleLeft,
-           s"expected shuffle in plan to be $shuffleLeft but found\n${joinOperator.left}")
-         assert(
--          joinOperator.right.exists(_.isInstanceOf[ShuffleExchangeExec]) == shuffleRight,
-+          joinOperator.right.exists(op => op.isInstanceOf[ShuffleExchangeLike]) == shuffleRight,
-           s"expected shuffle in plan to be $shuffleRight but found\n${joinOperator.right}")
- 
-         // check existence of sort
-         assert(
--          joinOperator.left.exists(_.isInstanceOf[SortExec]) == sortLeft,
-+          joinOperator.left.exists(op => op.isInstanceOf[SortExec] || op.isInstanceOf[CometExec] &&
-+            op.asInstanceOf[CometExec].originalPlan.isInstanceOf[SortExec]) == sortLeft,
-           s"expected sort in the left child to be $sortLeft but found\n${joinOperator.left}")
-         assert(
--          joinOperator.right.exists(_.isInstanceOf[SortExec]) == sortRight,
-+          joinOperator.right.exists(op => op.isInstanceOf[SortExec] || op.isInstanceOf[CometExec] &&
-+            op.asInstanceOf[CometExec].originalPlan.isInstanceOf[SortExec]) == sortRight,
-           s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}")
- 
-         // check the output partitioning
-@@ -835,11 +861,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti
-       df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table")
- 
-       val scanDF = spark.table("bucketed_table").select("j")
--      assert(!getFileScan(scanDF.queryExecution.executedPlan).bucketedScan)
-+      assert(!getBucketScan(scanDF.queryExecution.executedPlan))
-       checkAnswer(scanDF, df1.select("j"))
- 
-       val aggDF = spark.table("bucketed_table").groupBy("j").agg(max("k"))
--      assert(!getFileScan(aggDF.queryExecution.executedPlan).bucketedScan)
-+      assert(!getBucketScan(aggDF.queryExecution.executedPlan))
-       checkAnswer(aggDF, df1.groupBy("j").agg(max("k")))
-     }
-   }
-@@ -1026,15 +1052,23 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti
-             expectedNumShuffles: Int,
-             expectedCoalescedNumBuckets: Option[Int]): Unit = {
-           val plan = sql(query).queryExecution.executedPlan
--          val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
-+          val shuffles = plan.collect {
-+            case s: ShuffleExchangeLike => s
-+          }
-           assert(shuffles.length == expectedNumShuffles)
- 
-           val scans = plan.collect {
-             case f: FileSourceScanExec if f.optionalNumCoalescedBuckets.isDefined => f
-+            case b: CometScanExec if b.optionalNumCoalescedBuckets.isDefined => b
-           }
-           if (expectedCoalescedNumBuckets.isDefined) {
-             assert(scans.length == 1)
--            assert(scans.head.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
-+            scans.head match {
-+              case f: FileSourceScanExec =>
-+                assert(f.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
-+              case b: CometScanExec =>
-+                assert(b.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
-+            }
-           } else {
-             assert(scans.isEmpty)
-           }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
-index b5f6d2f9f68..277784a92af 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
-@@ -20,7 +20,7 @@ package org.apache.spark.sql.sources
- import java.io.File
- 
- import org.apache.spark.SparkException
--import org.apache.spark.sql.AnalysisException
-+import org.apache.spark.sql.{AnalysisException, IgnoreCometSuite}
- import org.apache.spark.sql.catalyst.TableIdentifier
- import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogTableType}
- import org.apache.spark.sql.catalyst.parser.ParseException
-@@ -28,7 +28,10 @@ import org.apache.spark.sql.internal.SQLConf.BUCKETING_MAX_BUCKETS
- import org.apache.spark.sql.test.SharedSparkSession
- import org.apache.spark.util.Utils
- 
--class CreateTableAsSelectSuite extends DataSourceTest with SharedSparkSession {
-+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
-+// Since it isn't related to Comet, we disable it for now.
-+class CreateTableAsSelectSuite extends DataSourceTest with SharedSparkSession
-+    with IgnoreCometSuite {
-   import testImplicits._
- 
-   protected override lazy val sql = spark.sql _
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
-index 1f55742cd67..42377f7cf26 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
-@@ -20,6 +20,7 @@ package org.apache.spark.sql.sources
- import org.apache.spark.sql.QueryTest
- import org.apache.spark.sql.catalyst.expressions.AttributeReference
- import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
-+import org.apache.spark.sql.comet.CometScanExec
- import org.apache.spark.sql.execution.FileSourceScanExec
- import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
- import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-@@ -71,7 +72,10 @@ abstract class DisableUnnecessaryBucketedScanSuite
- 
-     def checkNumBucketedScan(query: String, expectedNumBucketedScan: Int): Unit = {
-       val plan = sql(query).queryExecution.executedPlan
--      val bucketedScan = collect(plan) { case s: FileSourceScanExec if s.bucketedScan => s }
-+      val bucketedScan = collect(plan) {
-+        case s: FileSourceScanExec if s.bucketedScan => s
-+        case s: CometScanExec if s.bucketedScan => s
-+      }
-       assert(bucketedScan.length == expectedNumBucketedScan)
-     }
- 
-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 75f440caefc..36b1146bc3a 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
- import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
- import org.apache.spark.sql.{AnalysisException, DataFrame}
- import org.apache.spark.sql.catalyst.util.stringToFile
-+import org.apache.spark.sql.comet.CometBatchScanExec
- import org.apache.spark.sql.execution.DataSourceScanExec
- import org.apache.spark.sql.execution.datasources._
- import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, FileScan, FileTable}
-@@ -748,6 +749,8 @@ class FileStreamSinkV2Suite extends FileStreamSinkSuite {
-       val fileScan = df.queryExecution.executedPlan.collect {
-         case batch: BatchScanExec if batch.scan.isInstanceOf[FileScan] =>
-           batch.scan.asInstanceOf[FileScan]
-+        case batch: CometBatchScanExec if batch.scan.isInstanceOf[FileScan] =>
-+          batch.scan.asInstanceOf[FileScan]
-       }.headOption.getOrElse {
-         fail(s"No FileScan in query\n${df.queryExecution}")
-       }
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateDistributionSuite.scala
-index b597a244710..b2e8be41065 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateDistributionSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateDistributionSuite.scala
-@@ -21,6 +21,7 @@ import java.io.File
- 
- import org.apache.commons.io.FileUtils
- 
-+import org.apache.spark.sql.IgnoreComet
- import org.apache.spark.sql.catalyst.streaming.InternalOutputModes.Update
- import org.apache.spark.sql.execution.streaming.{FlatMapGroupsWithStateExec, MemoryStream}
- import org.apache.spark.sql.internal.SQLConf
-@@ -91,7 +92,7 @@ class FlatMapGroupsWithStateDistributionSuite extends StreamTest
-   }
- 
-   test("SPARK-38204: flatMapGroupsWithState should require StatefulOpClusteredDistribution " +
--    "from children - without initial state") {
-+    "from children - without initial state", IgnoreComet("TODO: fix Comet for this test")) {
-     // function will return -1 on timeout and returns count of the state otherwise
-     val stateFunc =
-       (key: (String, String), values: Iterator[(String, String, Long)],
-@@ -243,7 +244,8 @@ class FlatMapGroupsWithStateDistributionSuite extends StreamTest
-   }
- 
-   test("SPARK-38204: flatMapGroupsWithState should require ClusteredDistribution " +
--    "from children if the query starts from checkpoint in 3.2.x - without initial state") {
-+    "from children if the query starts from checkpoint in 3.2.x - without initial state",
-+    IgnoreComet("TODO: fix Comet for this test")) {
-     // function will return -1 on timeout and returns count of the state otherwise
-     val stateFunc =
-       (key: (String, String), values: Iterator[(String, String, Long)],
-@@ -335,7 +337,8 @@ class FlatMapGroupsWithStateDistributionSuite extends StreamTest
-   }
- 
-   test("SPARK-38204: flatMapGroupsWithState should require ClusteredDistribution " +
--    "from children if the query starts from checkpoint in prior to 3.2") {
-+    "from children if the query starts from checkpoint in prior to 3.2",
-+    IgnoreComet("TODO: fix Comet for this test")) {
-     // function will return -1 on timeout and returns count of the state otherwise
-     val stateFunc =
-       (key: (String, String), values: Iterator[(String, String, Long)],
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala
-index 6aa7d0945c7..38523536154 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala
-@@ -25,7 +25,7 @@ import org.scalatest.exceptions.TestFailedException
- 
- import org.apache.spark.SparkException
- import org.apache.spark.api.java.function.FlatMapGroupsWithStateFunction
--import org.apache.spark.sql.{DataFrame, Encoder}
-+import org.apache.spark.sql.{DataFrame, Encoder, IgnoreCometSuite}
- import org.apache.spark.sql.catalyst.InternalRow
- import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow}
- import org.apache.spark.sql.catalyst.plans.logical.FlatMapGroupsWithState
-@@ -46,8 +46,9 @@ case class RunningCount(count: Long)
- 
- case class Result(key: Long, count: Int)
- 
-+// TODO: fix Comet to enable this suite
- @SlowSQLTest
--class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest {
-+class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest with IgnoreCometSuite {
- 
-   import testImplicits._
- 
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateWithInitialStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateWithInitialStateSuite.scala
-index 2a2a83d35e1..e3b7b290b3e 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateWithInitialStateSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateWithInitialStateSuite.scala
-@@ -18,7 +18,7 @@
- package org.apache.spark.sql.streaming
- 
- import org.apache.spark.SparkException
--import org.apache.spark.sql.{AnalysisException, Dataset, KeyValueGroupedDataset}
-+import org.apache.spark.sql.{AnalysisException, Dataset, IgnoreComet, KeyValueGroupedDataset}
- import org.apache.spark.sql.catalyst.streaming.InternalOutputModes.Update
- import org.apache.spark.sql.execution.streaming.MemoryStream
- import org.apache.spark.sql.execution.streaming.state.FlatMapGroupsWithStateExecHelper
-@@ -253,7 +253,8 @@ class FlatMapGroupsWithStateWithInitialStateSuite extends StateStoreMetricsTest
-     assert(e.message.contains(expectedError))
-   }
- 
--  test("flatMapGroupsWithState - initial state - initial state has flatMapGroupsWithState") {
-+  test("flatMapGroupsWithState - initial state - initial state has flatMapGroupsWithState",
-+    IgnoreComet("TODO: fix Comet for this test")) {
-     val initialStateDS = Seq(("keyInStateAndData", new RunningCount(1))).toDS()
-     val initialState: KeyValueGroupedDataset[String, RunningCount] =
-       initialStateDS.groupByKey(_._1).mapValues(_._2)
-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 ef5b8a769fe..84fe1bfabc9 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
-@@ -37,6 +37,7 @@ import org.apache.spark.sql._
- import org.apache.spark.sql.catalyst.plans.logical.{Range, RepartitionByExpression}
- import org.apache.spark.sql.catalyst.streaming.{InternalOutputModes, StreamingRelationV2}
- import org.apache.spark.sql.catalyst.util.DateTimeUtils
-+import org.apache.spark.sql.comet.CometLocalLimitExec
- import org.apache.spark.sql.execution.{LocalLimitExec, SimpleMode, SparkPlan}
- import org.apache.spark.sql.execution.command.ExplainCommand
- import org.apache.spark.sql.execution.streaming._
-@@ -1103,11 +1104,12 @@ class StreamSuite extends StreamTest {
-       val localLimits = execPlan.collect {
-         case l: LocalLimitExec => l
-         case l: StreamingLocalLimitExec => l
-+        case l: CometLocalLimitExec => l
-       }
- 
-       require(
-         localLimits.size == 1,
--        s"Cant verify local limit optimization with this plan:\n$execPlan")
-+        s"Cant verify local limit optimization ${localLimits.size} with this plan:\n$execPlan")
- 
-       if (expectStreamingLimit) {
-         assert(
-@@ -1115,7 +1117,8 @@ class StreamSuite extends StreamTest {
-           s"Local limit was not StreamingLocalLimitExec:\n$execPlan")
-       } else {
-         assert(
--          localLimits.head.isInstanceOf[LocalLimitExec],
-+          localLimits.head.isInstanceOf[LocalLimitExec] ||
-+            localLimits.head.isInstanceOf[CometLocalLimitExec],
-           s"Local limit was not LocalLimitExec:\n$execPlan")
-       }
-     }
-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
---- 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
- 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
-@@ -67,6 +68,7 @@ class StreamingAggregationDistributionSuite extends StreamTest
-         // verify aggregations in between, except partial aggregation
-         val allAggregateExecs = query.lastExecution.executedPlan.collect {
-           case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
-         }
- 
-         val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
-@@ -201,6 +203,7 @@ class StreamingAggregationDistributionSuite extends StreamTest
-         // verify aggregations in between, except partial aggregation
-         val allAggregateExecs = executedPlan.collect {
-           case a: BaseAggregateExec => a
-+          case c: CometHashAggregateExec => c.originalPlan
-         }
- 
-         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 4d92e270539..33f1c2eb75e 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
-@@ -31,7 +31,7 @@ import org.apache.spark.scheduler.ExecutorCacheTaskLocation
- 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.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._
-@@ -619,14 +619,28 @@ class StreamingInnerJoinSuite extends StreamingJoinSuite {
- 
-         val numPartitions = spark.sqlContext.conf.getConf(SQLConf.SHUFFLE_PARTITIONS)
- 
--        assert(query.lastExecution.executedPlan.collect {
--          case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _,
--            ShuffleExchangeExec(opA: HashPartitioning, _, _),
--            ShuffleExchangeExec(opB: HashPartitioning, _, _))
--              if partitionExpressionsColumns(opA.expressions) === Seq("a", "b")
--                && partitionExpressionsColumns(opB.expressions) === Seq("a", "b")
--                && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j
--        }.size == 1)
-+        val join = query.lastExecution.executedPlan.collect {
-+          case j: StreamingSymmetricHashJoinExec => j
-+        }.head
-+        val opA = join.left.collect {
-+          case s: ShuffleExchangeLike
-+            if s.outputPartitioning.isInstanceOf[HashPartitioning] &&
-+              partitionExpressionsColumns(
-+                s.outputPartitioning
-+                  .asInstanceOf[HashPartitioning].expressions) === Seq("a", "b") =>
-+            s.outputPartitioning
-+              .asInstanceOf[HashPartitioning]
-+        }.head
-+        val opB = join.right.collect {
-+          case s: ShuffleExchangeLike
-+            if s.outputPartitioning.isInstanceOf[HashPartitioning] &&
-+              partitionExpressionsColumns(
-+                s.outputPartitioning
-+                  .asInstanceOf[HashPartitioning].expressions) === Seq("a", "b") =>
-+            s.outputPartitioning
-+              .asInstanceOf[HashPartitioning]
-+        }.head
-+        assert(opA.numPartitions == numPartitions && opB.numPartitions == numPartitions)
-       })
-   }
- 
-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 abe606ad9c1..2d930b64cca 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
- 
- import org.scalatest.BeforeAndAfter
- 
--import org.apache.spark.sql.{AnalysisException, Row, SaveMode}
-+import org.apache.spark.sql.{AnalysisException, IgnoreComet, Row, SaveMode}
- import org.apache.spark.sql.catalyst.TableIdentifier
- import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
- import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType}
-@@ -327,7 +327,8 @@ class DataStreamTableAPISuite extends StreamTest with BeforeAndAfter {
-     }
-   }
- 
--  test("explain with table on DSv1 data source") {
-+  test("explain with table on DSv1 data source",
-+      IgnoreComet("Comet explain output is different")) {
-     val tblSourceName = "tbl_src"
-     val tblTargetName = "tbl_target"
-     val tblSourceQualified = s"default.$tblSourceName"
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-index dd55fcfe42c..293e9dc2986 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
-@@ -41,6 +41,7 @@ import org.apache.spark.sql.catalyst.plans.PlanTest
- import org.apache.spark.sql.catalyst.plans.PlanTestBase
- import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
- import org.apache.spark.sql.catalyst.util._
-+import org.apache.spark.sql.comet._
- import org.apache.spark.sql.execution.FilterExec
- import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecution
- import org.apache.spark.sql.execution.datasources.DataSourceUtils
-@@ -126,7 +127,11 @@ private[sql] trait SQLTestUtils extends SparkFunSuite with SQLTestUtilsBase with
-         }
-       }
-     } else {
--      super.test(testName, testTags: _*)(testFun)
-+      if (isCometEnabled && testTags.exists(_.isInstanceOf[IgnoreComet])) {
-+        ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun)
-+      } else {
-+        super.test(testName, testTags: _*)(testFun)
-+      }
-     }
-   }
- 
-@@ -242,6 +247,32 @@ private[sql] trait SQLTestUtilsBase
-     protected override def _sqlContext: SQLContext = self.spark.sqlContext
-   }
- 
-+  /**
-+   * Whether Comet extension is enabled
-+   */
-+  protected def isCometEnabled: Boolean = {
-+    val v = System.getenv("ENABLE_COMET")
-+    v != null && v.toBoolean
-+  }
-+
-+  /**
-+   * Whether to enable ansi mode This is only effective when
-+   * [[isCometEnabled]] returns true.
-+   */
-+  protected def enableCometAnsiMode: Boolean = {
-+    val v = System.getenv("ENABLE_COMET_ANSI_MODE")
-+    v != null && v.toBoolean
-+  }
-+
-+  /**
-+   * Whether Spark should only apply Comet scan optimization. This is only effective when
-+   * [[isCometEnabled]] returns true.
-+   */
-+  protected def isCometScanOnly: Boolean = {
-+    val v = System.getenv("ENABLE_COMET_SCAN_ONLY")
-+    v != null && v.toBoolean
-+  }
-+
-   protected override def withSQLConf(pairs: (String, String)*)(f: => Unit): Unit = {
-     SparkSession.setActiveSession(spark)
-     super.withSQLConf(pairs: _*)(f)
-@@ -434,6 +465,8 @@ private[sql] trait SQLTestUtilsBase
-     val schema = df.schema
-     val withoutFilters = df.queryExecution.executedPlan.transform {
-       case FilterExec(_, child) => child
-+      case CometFilterExec(_, _, _, child, _) => child
-+      case CometProjectExec(_, _, _, _, CometFilterExec(_, _, _, child, _), _) => child
-     }
- 
-     spark.internalCreateDataFrame(withoutFilters.execute(), schema)
-diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-index ed2e309fa07..e071fc44960 100644
---- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
-@@ -74,6 +74,28 @@ trait SharedSparkSessionBase
-       // this rule may potentially block testing of other optimization rules such as
-       // ConstantPropagation etc.
-       .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
-+    // Enable Comet if `ENABLE_COMET` environment variable is set
-+    if (isCometEnabled) {
-+      conf
-+        .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
-+        .set("spark.comet.enabled", "true")
-+
-+      if (!isCometScanOnly) {
-+        conf
-+          .set("spark.comet.exec.enabled", "true")
-+          .set("spark.comet.exec.all.enabled", "true")
-+          .set("spark.shuffle.manager",
-+            "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
-+          .set("spark.comet.exec.shuffle.enabled", "true")
-+          .set("spark.comet.memoryOverhead", "10g")
-+      }
-+
-+      if (enableCometAnsiMode) {
-+        conf
-+          .set("spark.sql.ansi.enabled", "true")
-+          .set("spark.comet.ansi.enabled", "true")
-+      }
-+    }
-     conf.set(
-       StaticSQLConf.WAREHOUSE_PATH,
-       conf.get(StaticSQLConf.WAREHOUSE_PATH) + "/" + getClass.getCanonicalName)
-diff --git a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
-index 1510e8957f9..7618419d8ff 100644
---- a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
-+++ b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
-@@ -43,7 +43,7 @@ class SqlResourceWithActualMetricsSuite
-   import testImplicits._
- 
-   // Exclude nodes which may not have the metrics
--  val excludedNodes = List("WholeStageCodegen", "Project", "SerializeFromObject")
-+  val excludedNodes = List("WholeStageCodegen", "Project", "SerializeFromObject", "RowToColumnar")
- 
-   implicit val formats = 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
---- 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
- 
- import org.apache.spark.sql._
- import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression}
-+import org.apache.spark.sql.comet._
- 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
-       case s: FileSourceScanExec => s.partitionFilters.collect {
-         case d: DynamicPruningExpression => d.child
-       }
-+      case s: CometScanExec => s.partitionFilters.collect {
-+        case d: DynamicPruningExpression => d.child
-+      }
-       case h: HiveTableScanExec => h.partitionPruningPred.collect {
-         case d: DynamicPruningExpression => d.child
-       }
-diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
-index 1966e1e64fd..cde97a0aafe 100644
---- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
-+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
-@@ -656,7 +656,8 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
-         Row(3, 4, 4, 3, null) :: Nil)
-   }
- 
--  test("single distinct multiple columns set") {
-+  test("single distinct multiple columns set",
-+    IgnoreComet("TODO: fix Comet for this test")) {
-     checkAnswer(
-       spark.sql(
-         """
-diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-index 07361cfdce9..25b0dc3ef7e 100644
---- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
-@@ -55,25 +55,53 @@ object TestHive
-     new SparkContext(
-       System.getProperty("spark.sql.test.master", "local[1]"),
-       "TestSQLContext",
--      new SparkConf()
--        .set("spark.sql.test", "")
--        .set(SQLConf.CODEGEN_FALLBACK.key, "false")
--        .set(SQLConf.CODEGEN_FACTORY_MODE.key, CodegenObjectFactoryMode.CODEGEN_ONLY.toString)
--        .set(HiveUtils.HIVE_METASTORE_BARRIER_PREFIXES.key,
--          "org.apache.spark.sql.hive.execution.PairSerDe")
--        .set(WAREHOUSE_PATH.key, TestHiveContext.makeWarehouseDir().toURI.getPath)
--        // SPARK-8910
--        .set(UI_ENABLED, false)
--        .set(config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK, true)
--        // Hive changed the default of hive.metastore.disallow.incompatible.col.type.changes
--        // from false to true. For details, see the JIRA HIVE-12320 and HIVE-17764.
--        .set("spark.hadoop.hive.metastore.disallow.incompatible.col.type.changes", "false")
--        // Disable ConvertToLocalRelation for better test coverage. Test cases built on
--        // LocalRelation will exercise the optimization rules better by disabling it as
--        // this rule may potentially block testing of other optimization rules such as
--        // ConstantPropagation etc.
--        .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)))
-+      {
-+        val conf = new SparkConf()
-+          .set("spark.sql.test", "")
-+          .set(SQLConf.CODEGEN_FALLBACK.key, "false")
-+          .set(SQLConf.CODEGEN_FACTORY_MODE.key, CodegenObjectFactoryMode.CODEGEN_ONLY.toString)
-+          .set(HiveUtils.HIVE_METASTORE_BARRIER_PREFIXES.key,
-+            "org.apache.spark.sql.hive.execution.PairSerDe")
-+          .set(WAREHOUSE_PATH.key, TestHiveContext.makeWarehouseDir().toURI.getPath)
-+          // SPARK-8910
-+          .set(UI_ENABLED, false)
-+          .set(config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK, true)
-+          // Hive changed the default of hive.metastore.disallow.incompatible.col.type.changes
-+          // from false to true. For details, see the JIRA HIVE-12320 and HIVE-17764.
-+          .set("spark.hadoop.hive.metastore.disallow.incompatible.col.type.changes", "false")
-+          // Disable ConvertToLocalRelation for better test coverage. Test cases built on
-+          // LocalRelation will exercise the optimization rules better by disabling it as
-+          // this rule may potentially block testing of other optimization rules such as
-+          // ConstantPropagation etc.
-+          .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
-+
-+        val v = System.getenv("ENABLE_COMET")
-+        if (v != null && v.toBoolean) {
-+          conf
-+            .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
-+            .set("spark.comet.enabled", "true")
-+
-+          val v = System.getenv("ENABLE_COMET_SCAN_ONLY")
-+          if (v == null || !v.toBoolean) {
-+            conf
-+              .set("spark.comet.exec.enabled", "true")
-+              .set("spark.comet.exec.all.enabled", "true")
-+              .set("spark.shuffle.manager",
-+                "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
-+              .set("spark.comet.exec.shuffle.enabled", "true")
-+          }
-+
-+          val a = System.getenv("ENABLE_COMET_ANSI_MODE")
-+          if (a != null && a.toBoolean) {
-+            conf
-+              .set("spark.sql.ansi.enabled", "true")
-+              .set("spark.comet.ansi.enabled", "true")
-+          }
-+        }
- 
-+        conf
-+      }
-+    ))
- 
- case class TestHiveVersion(hiveClient: HiveClient)
-   extends TestHiveContext(TestHive.sparkContext, hiveClient)

From c1cdf46baaec3d75e1f853662588a4ab4ce8cb5a Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Thu, 6 Jun 2024 00:04:14 -0700
Subject: [PATCH 3/8] fix: Input batch to ShuffleRepartitioner.insert_batch
 should not be larger than configured batch size (#523)

* fix: Input batch to ShuffleRepartitioner.insert_batch should not be larger than configured batch size

* Add test

* For review
---
 .../execution/datafusion/shuffle_writer.rs    | 64 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 3 deletions(-)

diff --git a/core/src/execution/datafusion/shuffle_writer.rs b/core/src/execution/datafusion/shuffle_writer.rs
index 967340979..99ac885b5 100644
--- a/core/src/execution/datafusion/shuffle_writer.rs
+++ b/core/src/execution/datafusion/shuffle_writer.rs
@@ -575,6 +575,8 @@ struct ShuffleRepartitioner {
     hashes_buf: Vec<u32>,
     /// Partition ids for each row in the current batch
     partition_ids: Vec<u64>,
+    /// The configured batch size
+    batch_size: usize,
 }
 
 struct ShuffleRepartitionerMetrics {
@@ -642,17 +644,41 @@ impl ShuffleRepartitioner {
             reservation,
             hashes_buf,
             partition_ids,
+            batch_size,
         }
     }
 
+    /// Shuffles rows in input batch into corresponding partition buffer.
+    /// This function will slice input batch according to configured batch size and then
+    /// shuffle rows into corresponding partition buffer.
+    async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> {
+        let mut start = 0;
+        while start < batch.num_rows() {
+            let end = (start + self.batch_size).min(batch.num_rows());
+            let batch = batch.slice(start, end - start);
+            self.partitioning_batch(batch).await?;
+            start = end;
+        }
+        Ok(())
+    }
+
     /// Shuffles rows in input batch into corresponding partition buffer.
     /// This function first calculates hashes for rows and then takes rows in same
     /// partition as a record batch which is appended into partition buffer.
-    async fn insert_batch(&mut self, input: RecordBatch) -> Result<()> {
+    /// This should not be called directly. Use `insert_batch` instead.
+    async fn partitioning_batch(&mut self, input: RecordBatch) -> Result<()> {
         if input.num_rows() == 0 {
             // skip empty batch
             return Ok(());
         }
+
+        if input.num_rows() > self.batch_size {
+            return Err(DataFusionError::Internal(
+                "Input batch size exceeds configured batch size. Call `insert_batch` instead."
+                    .to_string(),
+            ));
+        }
+
         let _timer = self.metrics.baseline.elapsed_compute().timer();
 
         // NOTE: in shuffle writer exec, the output_rows metrics represents the
@@ -951,8 +977,7 @@ async fn external_shuffle(
     );
 
     while let Some(batch) = input.next().await {
-        let batch = batch?;
-        repartitioner.insert_batch(batch).await?;
+        repartitioner.insert_batch(batch?).await?;
     }
     repartitioner.shuffle_write().await
 }
@@ -1387,6 +1412,11 @@ impl RecordBatchStream for EmptyStream {
 #[cfg(test)]
 mod test {
     use super::*;
+    use datafusion::physical_plan::common::collect;
+    use datafusion::physical_plan::memory::MemoryExec;
+    use datafusion::prelude::SessionContext;
+    use datafusion_physical_expr::expressions::Column;
+    use tokio::runtime::Runtime;
 
     #[test]
     fn test_slot_size() {
@@ -1415,4 +1445,32 @@ mod test {
                 assert_eq!(slot_size, *expected);
             })
     }
+
+    #[test]
+    fn test_insert_larger_batch() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
+        let mut b = StringBuilder::new();
+        for i in 0..10000 {
+            b.append_value(format!("{i}"));
+        }
+        let array = b.finish();
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
+
+        let mut batches = Vec::new();
+        batches.push(batch.clone());
+
+        let partitions = &[batches];
+        let exec = ShuffleWriterExec::try_new(
+            Arc::new(MemoryExec::try_new(partitions, batch.schema(), None).unwrap()),
+            Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16),
+            "/tmp/data.out".to_string(),
+            "/tmp/index.out".to_string(),
+        )
+        .unwrap();
+        let ctx = SessionContext::new();
+        let task_ctx = ctx.task_ctx();
+        let stream = exec.execute(0, task_ctx).unwrap();
+        let rt = Runtime::new().unwrap();
+        rt.block_on(collect(stream)).unwrap();
+    }
 }

From c3e27c0405cd2ce3888dc1943323113735c7859e Mon Sep 17 00:00:00 2001
From: Emil Ejbyfeldt <eejbyfeldt@liveintent.com>
Date: Fri, 7 Jun 2024 07:11:29 +0200
Subject: [PATCH 4/8] Fix overflow in date_parser (#529)

The date_parser was introduced in #383 and is mostly a direct port of
code in Spark. Since the code uses the JVM it has defined integer
overflow as wrapping. The proposed fixed is to use std::num::Wrapping to
get the same wrapping behavior in rust.  The overflown value will still
be disgarded in a later check that uses `current_segment_digits` so
allowing the overflow does not lead to correctness issues.

This resolves one of the overflows discussed in #481
---
 .../execution/datafusion/expressions/cast.rs   | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/core/src/execution/datafusion/expressions/cast.rs b/core/src/execution/datafusion/expressions/cast.rs
index 7e8a96f28..a39587b6e 100644
--- a/core/src/execution/datafusion/expressions/cast.rs
+++ b/core/src/execution/datafusion/expressions/cast.rs
@@ -19,6 +19,7 @@ use std::{
     any::Any,
     fmt::{Debug, Display, Formatter},
     hash::{Hash, Hasher},
+    num::Wrapping,
     sync::Arc,
 };
 
@@ -1570,7 +1571,7 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> CometResult<Option<i32>>
     let mut date_segments = [1, 1, 1];
     let mut sign = 1;
     let mut current_segment = 0;
-    let mut current_segment_value = 0;
+    let mut current_segment_value = Wrapping(0);
     let mut current_segment_digits = 0;
     let bytes = date_str.as_bytes();
 
@@ -1597,16 +1598,16 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> CometResult<Option<i32>>
                 return return_result(date_str, eval_mode);
             }
             //if valid update corresponding segment with the current segment value.
-            date_segments[current_segment as usize] = current_segment_value;
-            current_segment_value = 0;
+            date_segments[current_segment as usize] = current_segment_value.0;
+            current_segment_value = Wrapping(0);
             current_segment_digits = 0;
             current_segment += 1;
         } else if !b.is_ascii_digit() {
             return return_result(date_str, eval_mode);
         } else {
             //increment value of current segment by the next digit
-            let parsed_value = (b - b'0') as i32;
-            current_segment_value = current_segment_value * 10 + parsed_value;
+            let parsed_value = Wrapping((b - b'0') as i32);
+            current_segment_value = current_segment_value * Wrapping(10) + parsed_value;
             current_segment_digits += 1;
         }
         j += 1;
@@ -1622,7 +1623,7 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> CometResult<Option<i32>>
         return return_result(date_str, eval_mode);
     }
 
-    date_segments[current_segment as usize] = current_segment_value;
+    date_segments[current_segment as usize] = current_segment_value.0;
 
     match NaiveDate::from_ymd_opt(
         sign * date_segments[0],
@@ -1836,6 +1837,8 @@ mod tests {
             Some(" 202 "),
             Some("\n 2020-\r8 "),
             Some("2020-01-01T"),
+            // Overflows i32
+            Some("-4607172990231812908"),
         ]));
 
         for eval_mode in &[EvalMode::Legacy, EvalMode::Try] {
@@ -1857,7 +1860,8 @@ mod tests {
                     None,
                     None,
                     None,
-                    Some(18262)
+                    Some(18262),
+                    None
                 ]
             );
         }

From c6d387c79507a1e4a71ecae5a26f7fc437d4f1d2 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 6 Jun 2024 23:12:36 -0600
Subject: [PATCH 5/8] feat: Add fuzz testing for arithmetic expressions (#519)

* Add fuzz tests for aritmetic expressions

* add unary math

* add bit-wise expressions

* bug fix
---
 fuzz-testing/README.md                        |  1 -
 .../scala/org/apache/comet/fuzz/Meta.scala    |  4 +++
 .../org/apache/comet/fuzz/QueryGen.scala      | 31 ++++++++++++++++++-
 .../org/apache/comet/fuzz/QueryRunner.scala   |  6 ++--
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/fuzz-testing/README.md b/fuzz-testing/README.md
index 076ff6aea..0f5f4f606 100644
--- a/fuzz-testing/README.md
+++ b/fuzz-testing/README.md
@@ -32,7 +32,6 @@ Planned areas of improvement:
 
 - ANSI mode
 - Support for all data types, expressions, and operators supported by Comet
-- Unary and binary arithmetic expressions
 - IF and CASE WHEN expressions
 - Complex (nested) expressions
 - Literal scalar values in queries
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala
index 13ebbf9ed..bbba5a443 100644
--- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala
+++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala
@@ -106,4 +106,8 @@ object Meta {
     Function("stddev_samp", 1),
     Function("corr", 2))
 
+  val unaryArithmeticOps: Seq[String] = Seq("+", "-")
+
+  val binaryArithmeticOps: Seq[String] = Seq("+", "-", "*", "/", "%", "&", "|", "^")
+
 }
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala
index 7584e76ce..e75726d7b 100644
--- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala
+++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala
@@ -42,11 +42,13 @@ object QueryGen {
     val uniqueQueries = mutable.HashSet[String]()
 
     for (_ <- 0 until numQueries) {
-      val sql = r.nextInt().abs % 4 match {
+      val sql = r.nextInt().abs % 6 match {
         case 0 => generateJoin(r, spark, numFiles)
         case 1 => generateAggregate(r, spark, numFiles)
         case 2 => generateScalar(r, spark, numFiles)
         case 3 => generateCast(r, spark, numFiles)
+        case 4 => generateUnaryArithmetic(r, spark, numFiles)
+        case 5 => generateBinaryArithmetic(r, spark, numFiles)
       }
       if (!uniqueQueries.contains(sql)) {
         uniqueQueries += sql
@@ -92,6 +94,33 @@ object QueryGen {
       s"ORDER BY ${args.mkString(", ")};"
   }
 
+  private def generateUnaryArithmetic(r: Random, spark: SparkSession, numFiles: Int): String = {
+    val tableName = s"test${r.nextInt(numFiles)}"
+    val table = spark.table(tableName)
+
+    val op = Utils.randomChoice(Meta.unaryArithmeticOps, r)
+    val a = Utils.randomChoice(table.columns, r)
+
+    // Example SELECT a, -a FROM test0
+    s"SELECT $a, $op$a " +
+      s"FROM $tableName " +
+      s"ORDER BY $a;"
+  }
+
+  private def generateBinaryArithmetic(r: Random, spark: SparkSession, numFiles: Int): String = {
+    val tableName = s"test${r.nextInt(numFiles)}"
+    val table = spark.table(tableName)
+
+    val op = Utils.randomChoice(Meta.binaryArithmeticOps, r)
+    val a = Utils.randomChoice(table.columns, r)
+    val b = Utils.randomChoice(table.columns, r)
+
+    // Example SELECT a, b, a+b FROM test0
+    s"SELECT $a, $b, $a $op $b " +
+      s"FROM $tableName " +
+      s"ORDER BY $a, $b;"
+  }
+
   private def generateCast(r: Random, spark: SparkSession, numFiles: Int): String = {
     val tableName = s"test${r.nextInt(numFiles)}"
     val table = spark.table(tableName)
diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala
index b2ceae9d0..f928c93a2 100644
--- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala
+++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryRunner.scala
@@ -19,7 +19,7 @@
 
 package org.apache.comet.fuzz
 
-import java.io.{BufferedWriter, FileWriter, PrintWriter}
+import java.io.{BufferedWriter, FileWriter, PrintWriter, StringWriter}
 
 import scala.io.Source
 
@@ -111,9 +111,11 @@ object QueryRunner {
                 showSQL(w, sql)
                 w.write(s"[ERROR] Query failed in Comet: ${e.getMessage}:\n")
                 w.write("```\n")
-                val p = new PrintWriter(w)
+                val sw = new StringWriter()
+                val p = new PrintWriter(sw)
                 e.printStackTrace(p)
                 p.close()
+                w.write(s"${sw.toString}\n")
                 w.write("```\n")
             }
 

From fd596ed985643be5be65fb71806b358219a146d9 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Fri, 7 Jun 2024 12:07:01 -0700
Subject: [PATCH 6/8] build: Switch back to official DataFusion repo and
 arrow-rs after Arrow Java 16 is released (#403)

* build: Switch back to released version of DataFusion and arrow-rs

* Exclude all arrow dependencies from Spark

* Revert "build: Switch back to released version of DataFusion and arrow-rs"

This reverts commit 29c89bfb25ddf4757ab17f951d3ccf17e55422da.

* Test

* Test

* Test arrow-rs fix

* Fix

* Use DataFusion repo

* Fix

* Fix

* Use 39.0.0-rc1
---
 core/Cargo.lock                               | 240 ++++++++++++------
 core/Cargo.toml                               |  25 +-
 .../datafusion/expressions/bitwise_not.rs     |   4 +-
 .../expressions/bloom_filter_might_contain.rs |   4 +-
 .../execution/datafusion/expressions/cast.rs  |   4 +-
 .../datafusion/expressions/checkoverflow.rs   |   4 +-
 .../datafusion/expressions/if_expr.rs         |  12 +-
 .../datafusion/expressions/negative.rs        |  19 +-
 .../datafusion/expressions/normalize_nan.rs   |   2 +-
 .../datafusion/expressions/scalar_funcs.rs    |  37 +--
 .../datafusion/expressions/strings.rs         |  12 +-
 .../datafusion/expressions/subquery.rs        |   2 +-
 .../datafusion/expressions/temporal.rs        |  20 +-
 .../datafusion/expressions/unbound.rs         |   2 +-
 .../execution/datafusion/operators/expand.rs  |   4 +-
 core/src/execution/datafusion/planner.rs      |  66 ++---
 .../execution/datafusion/shuffle_writer.rs    |   4 +-
 core/src/execution/kernels/hash.rs            |  16 ++
 core/src/execution/operators/copy.rs          |   4 +-
 core/src/execution/operators/scan.rs          |   2 +-
 dev/ensure-jars-have-correct-contents.sh      |   8 +
 pom.xml                                       |  27 +-
 22 files changed, 301 insertions(+), 217 deletions(-)

diff --git a/core/Cargo.lock b/core/Cargo.lock
index 6c8d54ab7..59c46c92b 100644
--- a/core/Cargo.lock
+++ b/core/Cargo.lock
@@ -114,8 +114,9 @@ checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711"
 
 [[package]]
 name = "arrow"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ae9728f104939be6d8d9b368a354b4929b0569160ea1641f0721b55a861ce38"
 dependencies = [
  "arrow-arith",
  "arrow-array",
@@ -134,48 +135,52 @@ dependencies = [
 
 [[package]]
 name = "arrow-arith"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7029a5b3efbeafbf4a12d12dc16b8f9e9bff20a410b8c25c5d28acc089e1043"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "chrono",
- "half 2.1.0",
+ "half 2.4.1",
  "num",
 ]
 
 [[package]]
 name = "arrow-array"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d33238427c60271710695f17742f45b1a5dc5bcfc5c15331c25ddfe7abf70d97"
 dependencies = [
  "ahash",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "chrono",
- "chrono-tz",
- "half 2.1.0",
+ "chrono-tz 0.9.0",
+ "half 2.4.1",
  "hashbrown",
  "num",
 ]
 
 [[package]]
 name = "arrow-buffer"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe9b95e825ae838efaf77e366c00d3fc8cca78134c9db497d6bda425f2e7b7c1"
 dependencies = [
  "bytes",
- "half 2.1.0",
+ "half 2.4.1",
  "num",
 ]
 
 [[package]]
 name = "arrow-cast"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "87cf8385a9d5b5fcde771661dd07652b79b9139fea66193eda6a88664400ccab"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -186,7 +191,7 @@ dependencies = [
  "base64",
  "chrono",
  "comfy-table",
- "half 2.1.0",
+ "half 2.4.1",
  "lexical-core",
  "num",
  "ryu",
@@ -194,8 +199,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-csv"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cea5068bef430a86690059665e40034625ec323ffa4dd21972048eebb0127adc"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -212,19 +218,21 @@ dependencies = [
 
 [[package]]
 name = "arrow-data"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb29be98f987bcf217b070512bb7afba2f65180858bca462edf4a39d84a23e10"
 dependencies = [
  "arrow-buffer",
  "arrow-schema",
- "half 2.1.0",
+ "half 2.4.1",
  "num",
 ]
 
 [[package]]
 name = "arrow-ipc"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffc68f6523970aa6f7ce1dc9a33a7d9284cfb9af77d4ad3e617dbe5d79cc6ec8"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -237,8 +245,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-json"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2041380f94bd6437ab648e6c2085a045e45a0c44f91a1b9a4fe3fed3d379bfb1"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -246,7 +255,7 @@ dependencies = [
  "arrow-data",
  "arrow-schema",
  "chrono",
- "half 2.1.0",
+ "half 2.4.1",
  "indexmap",
  "lexical-core",
  "num",
@@ -256,44 +265,48 @@ dependencies = [
 
 [[package]]
 name = "arrow-ord"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcb56ed1547004e12203652f12fe12e824161ff9d1e5cf2a7dc4ff02ba94f413"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "arrow-select",
- "half 2.1.0",
+ "half 2.4.1",
  "num",
 ]
 
 [[package]]
 name = "arrow-row"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "575b42f1fc588f2da6977b94a5ca565459f5ab07b60545e17243fb9a7ed6d43e"
 dependencies = [
  "ahash",
  "arrow-array",
  "arrow-buffer",
  "arrow-data",
  "arrow-schema",
- "half 2.1.0",
+ "half 2.4.1",
  "hashbrown",
 ]
 
 [[package]]
 name = "arrow-schema"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32aae6a60458a2389c0da89c9de0b7932427776127da1a738e2efc21d32f3393"
 dependencies = [
  "bitflags 2.5.0",
 ]
 
 [[package]]
 name = "arrow-select"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "de36abaef8767b4220d7b4a8c2fe5ffc78b47db81b03d77e2136091c3ba39102"
 dependencies = [
  "ahash",
  "arrow-array",
@@ -305,8 +318,9 @@ dependencies = [
 
 [[package]]
 name = "arrow-string"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e435ada8409bcafc910bc3e0077f532a4daa20e99060a496685c0e3e53cc2597"
 dependencies = [
  "arrow-array",
  "arrow-buffer",
@@ -508,7 +522,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d59ae0466b83e838b81a54256c39d5d7c20b9d7daa10510a242d9b75abd5936e"
 dependencies = [
  "chrono",
- "chrono-tz-build",
+ "chrono-tz-build 0.2.1",
+ "phf",
+]
+
+[[package]]
+name = "chrono-tz"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93698b29de5e97ad0ae26447b344c482a7284c737d9ddc5f9e52b74a336671bb"
+dependencies = [
+ "chrono",
+ "chrono-tz-build 0.3.0",
  "phf",
 ]
 
@@ -523,6 +548,17 @@ dependencies = [
  "phf_codegen",
 ]
 
+[[package]]
+name = "chrono-tz-build"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0c088aee841df9c3041febbb73934cfc39708749bf96dc827e3359cd39ef11b1"
+dependencies = [
+ "parse-zoneinfo",
+ "phf",
+ "phf_codegen",
+]
+
 [[package]]
 name = "ciborium"
 version = "0.2.1"
@@ -592,6 +628,7 @@ dependencies = [
  "ahash",
  "arrow",
  "arrow-array",
+ "arrow-buffer",
  "arrow-data",
  "arrow-schema",
  "arrow-string",
@@ -601,16 +638,18 @@ dependencies = [
  "byteorder",
  "bytes",
  "chrono",
- "chrono-tz",
+ "chrono-tz 0.8.6",
  "crc32fast",
  "criterion",
  "datafusion",
  "datafusion-common",
+ "datafusion-expr",
  "datafusion-functions",
  "datafusion-physical-expr",
+ "datafusion-physical-expr-common",
  "flate2",
  "futures",
- "half 2.1.0",
+ "half 2.4.1",
  "hashbrown",
  "itertools 0.11.0",
  "jni",
@@ -826,8 +865,8 @@ dependencies = [
 
 [[package]]
 name = "datafusion"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "ahash",
  "arrow",
@@ -843,13 +882,15 @@ dependencies = [
  "datafusion-execution",
  "datafusion-expr",
  "datafusion-functions",
+ "datafusion-functions-aggregate",
  "datafusion-optimizer",
  "datafusion-physical-expr",
+ "datafusion-physical-expr-common",
  "datafusion-physical-plan",
  "datafusion-sql",
  "futures",
  "glob",
- "half 2.1.0",
+ "half 2.4.1",
  "hashbrown",
  "indexmap",
  "itertools 0.12.1",
@@ -857,6 +898,7 @@ dependencies = [
  "num_cpus",
  "object_store",
  "parking_lot",
+ "paste",
  "pin-project-lite",
  "rand",
  "sqlparser",
@@ -868,8 +910,8 @@ dependencies = [
 
 [[package]]
 name = "datafusion-common"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "ahash",
  "arrow",
@@ -877,7 +919,8 @@ dependencies = [
  "arrow-buffer",
  "arrow-schema",
  "chrono",
- "half 2.1.0",
+ "half 2.4.1",
+ "hashbrown",
  "instant",
  "libc",
  "num_cpus",
@@ -887,16 +930,16 @@ dependencies = [
 
 [[package]]
 name = "datafusion-common-runtime"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "tokio",
 ]
 
 [[package]]
 name = "datafusion-execution"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "arrow",
  "chrono",
@@ -915,15 +958,17 @@ dependencies = [
 
 [[package]]
 name = "datafusion-expr"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "ahash",
  "arrow",
  "arrow-array",
+ "arrow-buffer",
  "chrono",
  "datafusion-common",
  "paste",
+ "serde_json",
  "sqlparser",
  "strum",
  "strum_macros",
@@ -931,8 +976,8 @@ dependencies = [
 
 [[package]]
 name = "datafusion-functions"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "arrow",
  "base64",
@@ -943,20 +988,39 @@ dependencies = [
  "datafusion-execution",
  "datafusion-expr",
  "datafusion-physical-expr",
+ "hashbrown",
  "hex",
  "itertools 0.12.1",
  "log",
  "md-5",
+ "rand",
  "regex",
  "sha2",
  "unicode-segmentation",
  "uuid",
 ]
 
+[[package]]
+name = "datafusion-functions-aggregate"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
+dependencies = [
+ "ahash",
+ "arrow",
+ "arrow-schema",
+ "datafusion-common",
+ "datafusion-execution",
+ "datafusion-expr",
+ "datafusion-physical-expr-common",
+ "log",
+ "paste",
+ "sqlparser",
+]
+
 [[package]]
 name = "datafusion-optimizer"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "arrow",
  "async-trait",
@@ -965,6 +1029,7 @@ dependencies = [
  "datafusion-expr",
  "datafusion-physical-expr",
  "hashbrown",
+ "indexmap",
  "itertools 0.12.1",
  "log",
  "regex-syntax",
@@ -972,8 +1037,8 @@ dependencies = [
 
 [[package]]
 name = "datafusion-physical-expr"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "ahash",
  "arrow",
@@ -983,36 +1048,44 @@ dependencies = [
  "arrow-schema",
  "arrow-string",
  "base64",
- "blake2",
- "blake3",
  "chrono",
  "datafusion-common",
  "datafusion-execution",
  "datafusion-expr",
- "half 2.1.0",
+ "datafusion-functions-aggregate",
+ "datafusion-physical-expr-common",
+ "half 2.4.1",
  "hashbrown",
  "hex",
  "indexmap",
  "itertools 0.12.1",
  "log",
- "md-5",
  "paste",
  "petgraph",
- "rand",
  "regex",
- "sha2",
- "unicode-segmentation",
+]
+
+[[package]]
+name = "datafusion-physical-expr-common"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
+dependencies = [
+ "arrow",
+ "datafusion-common",
+ "datafusion-expr",
+ "rand",
 ]
 
 [[package]]
 name = "datafusion-physical-plan"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "ahash",
  "arrow",
  "arrow-array",
  "arrow-buffer",
+ "arrow-ord",
  "arrow-schema",
  "async-trait",
  "chrono",
@@ -1020,9 +1093,11 @@ dependencies = [
  "datafusion-common-runtime",
  "datafusion-execution",
  "datafusion-expr",
+ "datafusion-functions-aggregate",
  "datafusion-physical-expr",
+ "datafusion-physical-expr-common",
  "futures",
- "half 2.1.0",
+ "half 2.4.1",
  "hashbrown",
  "indexmap",
  "itertools 0.12.1",
@@ -1036,8 +1111,8 @@ dependencies = [
 
 [[package]]
 name = "datafusion-sql"
-version = "36.0.0"
-source = "git+https://github.com/viirya/arrow-datafusion.git?rev=57b3be4#57b3be4297a47aa45094c16e37ddf0141d723bf0"
+version = "39.0.0"
+source = "git+https://github.com/apache/arrow-datafusion.git?tag=39.0.0-rc1#6a4a280e3cf70fe5f1a1cfe7c2de13e4c39f89bb"
 dependencies = [
  "arrow",
  "arrow-array",
@@ -1045,6 +1120,7 @@ dependencies = [
  "datafusion-common",
  "datafusion-expr",
  "log",
+ "regex",
  "sqlparser",
  "strum",
 ]
@@ -1140,9 +1216,9 @@ checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
 
 [[package]]
 name = "flatbuffers"
-version = "23.5.26"
+version = "24.3.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4dac53e22462d78c16d64a1cd22371b54cc3fe94aa15e7886a2fa6e5d1ab8640"
+checksum = "8add37afff2d4ffa83bc748a70b4b1370984f6980768554182424ef71447c35f"
 dependencies = [
  "bitflags 1.3.2",
  "rustc_version",
@@ -1303,10 +1379,11 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403"
 
 [[package]]
 name = "half"
-version = "2.1.0"
+version = "2.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ad6a9459c9c30b177b925162351f97e7d967c7ea8bab3b8352805327daf45554"
+checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888"
 dependencies = [
+ "cfg-if",
  "crunchy",
  "num-traits",
 ]
@@ -1905,9 +1982,9 @@ dependencies = [
 
 [[package]]
 name = "object_store"
-version = "0.9.1"
+version = "0.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b8718f8b65fdf67a45108d1548347d4af7d71fb81ce727bbf9e3b2535e079db3"
+checksum = "fbebfd32c213ba1907fa7a9c9138015a8de2b43e30c5aa45b18f7deb46786ad6"
 dependencies = [
  "async-trait",
  "bytes",
@@ -1979,13 +2056,14 @@ dependencies = [
 
 [[package]]
 name = "parquet"
-version = "51.0.0"
-source = "git+https://github.com/viirya/arrow-rs.git?rev=3f1ae0c#3f1ae0c836b0769c88220d2180ef008b7a59158c"
+version = "52.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29c3b5322cc1bbf67f11c079c42be41a55949099b78732f7dba9e15edde40eab"
 dependencies = [
  "ahash",
  "bytes",
  "chrono",
- "half 2.1.0",
+ "half 2.4.1",
  "hashbrown",
  "num",
  "num-bigint",
@@ -2535,9 +2613,9 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 
 [[package]]
 name = "sqlparser"
-version = "0.44.0"
+version = "0.47.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aaf9c7ff146298ffda83a200f8d5084f08dcee1edfc135fcc1d646a45d50ffd6"
+checksum = "295e9930cd7a97e58ca2a070541a3ca502b17f5d1fa7157376d0fabd85324f25"
 dependencies = [
  "log",
  "sqlparser_derive",
diff --git a/core/Cargo.toml b/core/Cargo.toml
index 4584dffce..564c450cb 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -29,13 +29,14 @@ include = [
 
 [dependencies]
 parquet-format = "4.0.0" # This must be kept in sync with that from parquet crate
-arrow = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c", features = ["prettyprint", "ffi", "chrono-tz"] }
-arrow-array = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c" }
-arrow-data = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c" }
-arrow-schema = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c" }
-arrow-string = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c" }
-parquet = { git = "https://github.com/viirya/arrow-rs.git", rev = "3f1ae0c", default-features = false, features = ["experimental"] }
-half = { version = "~2.1", default-features = false }
+arrow = { version = "52.0.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow-array = { version = "52.0.0" }
+arrow-buffer = { version = "52.0.0" }
+arrow-data = { version = "52.0.0" }
+arrow-schema = { version = "52.0.0" }
+arrow-string = { version = "52.0.0" }
+parquet = { version = "52.0.0", default-features = false, features = ["experimental"] }
+half = { version = "2.4.1", default-features = false }
 futures = "0.3.28"
 mimalloc = { version = "*", default-features = false, optional = true }
 tokio = { version = "1", features = ["rt-multi-thread"] }
@@ -66,10 +67,12 @@ itertools = "0.11.0"
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
 chrono-tz = { version = "0.8" }
 paste = "1.0.14"
-datafusion-common = { git = "https://github.com/viirya/arrow-datafusion.git", rev = "57b3be4" }
-datafusion = { default-features = false, git = "https://github.com/viirya/arrow-datafusion.git", rev = "57b3be4", features = ["unicode_expressions", "crypto_expressions"] }
-datafusion-functions = { git = "https://github.com/viirya/arrow-datafusion.git", rev = "57b3be4", features = ["crypto_expressions"]}
-datafusion-physical-expr = { git = "https://github.com/viirya/arrow-datafusion.git", rev = "57b3be4", default-features = false, features = ["unicode_expressions"] }
+datafusion-common = { git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1" }
+datafusion = { default-features = false, git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1", features = ["unicode_expressions", "crypto_expressions"] }
+datafusion-functions = { git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1", features = ["crypto_expressions"] }
+datafusion-expr = { git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1", default-features = false }
+datafusion-physical-expr-common = { git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1", default-features = false }
+datafusion-physical-expr = { git = "https://github.com/apache/arrow-datafusion.git", tag = "39.0.0-rc1", default-features = false }
 unicode-segmentation = "^1.10.1"
 once_cell = "1.18.0"
 regex = "1.9.6"
diff --git a/core/src/execution/datafusion/expressions/bitwise_not.rs b/core/src/execution/datafusion/expressions/bitwise_not.rs
index f9f8ee392..06ead2670 100644
--- a/core/src/execution/datafusion/expressions/bitwise_not.rs
+++ b/core/src/execution/datafusion/expressions/bitwise_not.rs
@@ -105,8 +105,8 @@ impl PhysicalExpr for BitwiseNotExpr {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.arg.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.arg]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/bloom_filter_might_contain.rs b/core/src/execution/datafusion/expressions/bloom_filter_might_contain.rs
index 6a4d07b89..b922119f8 100644
--- a/core/src/execution/datafusion/expressions/bloom_filter_might_contain.rs
+++ b/core/src/execution/datafusion/expressions/bloom_filter_might_contain.rs
@@ -129,8 +129,8 @@ impl PhysicalExpr for BloomFilterMightContain {
             })
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.bloom_filter_expr.clone(), self.value_expr.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.bloom_filter_expr, &self.value_expr]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/cast.rs b/core/src/execution/datafusion/expressions/cast.rs
index a39587b6e..045626465 100644
--- a/core/src/execution/datafusion/expressions/cast.rs
+++ b/core/src/execution/datafusion/expressions/cast.rs
@@ -1292,8 +1292,8 @@ impl PhysicalExpr for Cast {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/checkoverflow.rs b/core/src/execution/datafusion/expressions/checkoverflow.rs
index 1e4b5f333..044b366e3 100644
--- a/core/src/execution/datafusion/expressions/checkoverflow.rs
+++ b/core/src/execution/datafusion/expressions/checkoverflow.rs
@@ -165,8 +165,8 @@ impl PhysicalExpr for CheckOverflow {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/if_expr.rs b/core/src/execution/datafusion/expressions/if_expr.rs
index 6f2ed6a54..fa235cc66 100644
--- a/core/src/execution/datafusion/expressions/if_expr.rs
+++ b/core/src/execution/datafusion/expressions/if_expr.rs
@@ -110,12 +110,8 @@ impl PhysicalExpr for IfExpr {
         Ok(ColumnarValue::Array(current_value))
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![
-            self.if_expr.clone(),
-            self.true_expr.clone(),
-            self.false_expr.clone(),
-        ]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.if_expr, &self.true_expr, &self.false_expr]
     }
 
     fn with_new_children(
@@ -225,8 +221,8 @@ mod tests {
         let true_expr = lit(123i32);
         let false_expr = lit(999i32);
 
-        let expr = if_fn(if_expr, true_expr, false_expr);
-        let children = expr.unwrap().children();
+        let expr = if_fn(if_expr, true_expr, false_expr).unwrap();
+        let children = expr.children();
         assert_eq!(children.len(), 3);
         assert_eq!(children[0].to_string(), "true");
         assert_eq!(children[1].to_string(), "123");
diff --git a/core/src/execution/datafusion/expressions/negative.rs b/core/src/execution/datafusion/expressions/negative.rs
index e7aa2ac64..a85cde89e 100644
--- a/core/src/execution/datafusion/expressions/negative.rs
+++ b/core/src/execution/datafusion/expressions/negative.rs
@@ -18,15 +18,15 @@
 use crate::errors::CometError;
 use arrow::{compute::kernels::numeric::neg_wrapping, datatypes::IntervalDayTimeType};
 use arrow_array::RecordBatch;
+use arrow_buffer::IntervalDayTime;
 use arrow_schema::{DataType, Schema};
 use datafusion::{
     logical_expr::{interval_arithmetic::Interval, ColumnarValue},
     physical_expr::PhysicalExpr,
 };
 use datafusion_common::{Result, ScalarValue};
-use datafusion_physical_expr::{
-    aggregate::utils::down_cast_any_ref, sort_properties::SortProperties,
-};
+use datafusion_expr::sort_properties::ExprProperties;
+use datafusion_physical_expr::aggregate::utils::down_cast_any_ref;
 use std::{
     any::Any,
     hash::{Hash, Hasher},
@@ -63,7 +63,7 @@ macro_rules! check_overflow {
         for i in 0..typed_array.len() {
             if typed_array.value(i) == $min_val {
                 if $type_name == "byte" || $type_name == "short" {
-                    let value = typed_array.value(i).to_string() + " caused";
+                    let value = format!("{:?} caused", typed_array.value(i));
                     return Err(arithmetic_overflow_error(value.as_str()).into());
                 }
                 return Err(arithmetic_overflow_error($type_name).into());
@@ -135,7 +135,7 @@ impl PhysicalExpr for NegativeExpr {
                             arrow::datatypes::IntervalUnit::DayTime => check_overflow!(
                                 array,
                                 arrow::array::IntervalDayTimeArray,
-                                i64::MIN,
+                                IntervalDayTime::MIN,
                                 "interval"
                             ),
                             arrow::datatypes::IntervalUnit::MonthDayNano => {
@@ -195,8 +195,8 @@ impl PhysicalExpr for NegativeExpr {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.arg.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.arg]
     }
 
     fn with_new_children(
@@ -255,8 +255,9 @@ impl PhysicalExpr for NegativeExpr {
     }
 
     /// The ordering of a [`NegativeExpr`] is simply the reverse of its child.
-    fn get_ordering(&self, children: &[SortProperties]) -> SortProperties {
-        -children[0]
+    fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> {
+        let properties = children[0].clone().with_order(children[0].sort_properties);
+        Ok(properties)
     }
 }
 
diff --git a/core/src/execution/datafusion/expressions/normalize_nan.rs b/core/src/execution/datafusion/expressions/normalize_nan.rs
index 111a34d5d..3bd5feea5 100644
--- a/core/src/execution/datafusion/expressions/normalize_nan.rs
+++ b/core/src/execution/datafusion/expressions/normalize_nan.rs
@@ -77,7 +77,7 @@ impl PhysicalExpr for NormalizeNaNAndZero {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
         self.child.children()
     }
 
diff --git a/core/src/execution/datafusion/expressions/scalar_funcs.rs b/core/src/execution/datafusion/expressions/scalar_funcs.rs
index dc333e8be..5f98ce3f6 100644
--- a/core/src/execution/datafusion/expressions/scalar_funcs.rs
+++ b/core/src/execution/datafusion/expressions/scalar_funcs.rs
@@ -19,7 +19,6 @@ use std::{
     any::Any,
     cmp::min,
     fmt::{Debug, Write},
-    str::FromStr,
     sync::Arc,
 };
 
@@ -35,17 +34,15 @@ use arrow_array::{Array, ArrowNativeTypeOp, Decimal128Array, StringArray};
 use arrow_schema::DataType;
 use datafusion::{
     execution::FunctionRegistry,
-    logical_expr::{
-        BuiltinScalarFunction, ScalarFunctionDefinition, ScalarFunctionImplementation,
-        ScalarUDFImpl, Signature, Volatility,
-    },
+    functions::math::round::round,
+    logical_expr::{ScalarFunctionImplementation, ScalarUDFImpl, Signature, Volatility},
     physical_plan::ColumnarValue,
 };
 use datafusion_common::{
     cast::{as_binary_array, as_generic_string_array},
     exec_err, internal_err, DataFusionError, Result as DataFusionResult, ScalarValue,
 };
-use datafusion_physical_expr::{math_expressions, udf::ScalarUDF};
+use datafusion_expr::ScalarUDF;
 use num::{
     integer::{div_ceil, div_floor},
     BigInt, Signed, ToPrimitive,
@@ -66,9 +63,7 @@ macro_rules! make_comet_scalar_udf {
             $data_type.clone(),
             Arc::new(move |args| $func(args, &$data_type)),
         );
-        Ok(ScalarFunctionDefinition::UDF(Arc::new(
-            ScalarUDF::new_from_impl(scalar_func),
-        )))
+        Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
     }};
     ($name:expr, $func:expr, without $data_type:ident) => {{
         let scalar_func = CometScalarFunction::new(
@@ -77,9 +72,7 @@ macro_rules! make_comet_scalar_udf {
             $data_type,
             $func,
         );
-        Ok(ScalarFunctionDefinition::UDF(Arc::new(
-            ScalarUDF::new_from_impl(scalar_func),
-        )))
+        Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
     }};
 }
 
@@ -88,7 +81,7 @@ pub fn create_comet_physical_fun(
     fun_name: &str,
     data_type: DataType,
     registry: &dyn FunctionRegistry,
-) -> Result<ScalarFunctionDefinition, DataFusionError> {
+) -> Result<Arc<ScalarUDF>, DataFusionError> {
     let sha2_functions = ["sha224", "sha256", "sha384", "sha512"];
     match fun_name {
         "ceil" => {
@@ -140,13 +133,11 @@ pub fn create_comet_physical_fun(
             let spark_func_name = "spark".to_owned() + sha;
             make_comet_scalar_udf!(spark_func_name, wrapped_func, without data_type)
         }
-        _ => {
-            if let Ok(fun) = BuiltinScalarFunction::from_str(fun_name) {
-                Ok(ScalarFunctionDefinition::BuiltIn(fun))
-            } else {
-                Ok(ScalarFunctionDefinition::UDF(registry.udf(fun_name)?))
-            }
-        }
+        _ => registry.udf(fun_name).map_err(|e| {
+            DataFusionError::Execution(format!(
+                "Function {fun_name} not found in the registry: {e}",
+            ))
+        }),
     }
 }
 
@@ -509,9 +500,7 @@ fn spark_round(
                 make_decimal_array(array, precision, scale, &f)
             }
             DataType::Float32 | DataType::Float64 => {
-                Ok(ColumnarValue::Array(math_expressions::round(&[
-                    array.clone()
-                ])?))
+                Ok(ColumnarValue::Array(round(&[array.clone()])?))
             }
             dt => exec_err!("Not supported datatype for ROUND: {dt}"),
         },
@@ -534,7 +523,7 @@ fn spark_round(
                 make_decimal_scalar(a, precision, scale, &f)
             }
             ScalarValue::Float32(_) | ScalarValue::Float64(_) => Ok(ColumnarValue::Scalar(
-                ScalarValue::try_from_array(&math_expressions::round(&[a.to_array()?])?, 0)?,
+                ScalarValue::try_from_array(&round(&[a.to_array()?])?, 0)?,
             )),
             dt => exec_err!("Not supported datatype for ROUND: {dt}"),
         },
diff --git a/core/src/execution/datafusion/expressions/strings.rs b/core/src/execution/datafusion/expressions/strings.rs
index ee9a22212..cbbd4cfa4 100644
--- a/core/src/execution/datafusion/expressions/strings.rs
+++ b/core/src/execution/datafusion/expressions/strings.rs
@@ -111,8 +111,8 @@ macro_rules! make_predicate_function {
                 Ok(ColumnarValue::Array(Arc::new(array)))
             }
 
-            fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-                vec![self.left.clone(), self.right.clone()]
+            fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+                vec![&self.left, &self.right]
             }
 
             fn with_new_children(
@@ -221,8 +221,8 @@ impl PhysicalExpr for SubstringExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
@@ -286,8 +286,8 @@ impl PhysicalExpr for StringSpaceExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/subquery.rs b/core/src/execution/datafusion/expressions/subquery.rs
index bf37cb895..9b1be2df6 100644
--- a/core/src/execution/datafusion/expressions/subquery.rs
+++ b/core/src/execution/datafusion/expressions/subquery.rs
@@ -199,7 +199,7 @@ impl PhysicalExpr for Subquery {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
         vec![]
     }
 
diff --git a/core/src/execution/datafusion/expressions/temporal.rs b/core/src/execution/datafusion/expressions/temporal.rs
index 4ae3c2605..22b4aee8a 100644
--- a/core/src/execution/datafusion/expressions/temporal.rs
+++ b/core/src/execution/datafusion/expressions/temporal.rs
@@ -111,8 +111,8 @@ impl PhysicalExpr for HourExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
@@ -205,8 +205,8 @@ impl PhysicalExpr for MinuteExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
@@ -299,8 +299,8 @@ impl PhysicalExpr for SecondExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
@@ -386,8 +386,8 @@ impl PhysicalExpr for DateTruncExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
@@ -511,8 +511,8 @@ impl PhysicalExpr for TimestampTruncExec {
         }
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/expressions/unbound.rs b/core/src/execution/datafusion/expressions/unbound.rs
index 5387b1012..95f9912c9 100644
--- a/core/src/execution/datafusion/expressions/unbound.rs
+++ b/core/src/execution/datafusion/expressions/unbound.rs
@@ -83,7 +83,7 @@ impl PhysicalExpr for UnboundColumn {
         internal_err!("UnboundColumn::evaluate() should not be called")
     }
 
-    fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
         vec![]
     }
 
diff --git a/core/src/execution/datafusion/operators/expand.rs b/core/src/execution/datafusion/operators/expand.rs
index ca3fdc1aa..5285dfb46 100644
--- a/core/src/execution/datafusion/operators/expand.rs
+++ b/core/src/execution/datafusion/operators/expand.rs
@@ -96,8 +96,8 @@ impl ExecutionPlan for CometExpandExec {
         self.schema.clone()
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
-        vec![self.child.clone()]
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.child]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
index 7af5f6838..7d504f878 100644
--- a/core/src/execution/datafusion/planner.rs
+++ b/core/src/execution/datafusion/planner.rs
@@ -17,7 +17,7 @@
 
 //! Converts Spark physical plan to DataFusion physical plan
 
-use std::{collections::HashMap, str::FromStr, sync::Arc};
+use std::{collections::HashMap, sync::Arc};
 
 use arrow_schema::{DataType, Field, Schema, TimeUnit};
 use datafusion::{
@@ -25,9 +25,7 @@ use datafusion::{
     common::DataFusionError,
     execution::FunctionRegistry,
     functions::math,
-    logical_expr::{
-        BuiltinScalarFunction, Operator as DataFusionOperator, ScalarFunctionDefinition,
-    },
+    logical_expr::Operator as DataFusionOperator,
     physical_expr::{
         execution_props::ExecutionProps,
         expressions::{
@@ -52,6 +50,7 @@ use datafusion_common::{
     tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter},
     JoinType as DFJoinType, ScalarValue,
 };
+use datafusion_physical_expr_common::aggregate::create_aggregate_expr;
 use itertools::Itertools;
 use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
@@ -499,10 +498,7 @@ impl PhysicalPlanner {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema.clone())?;
                 let return_type = child.data_type(&input_schema)?;
                 let args = vec![child];
-                let scalar_def = ScalarFunctionDefinition::UDF(math::abs());
-
-                let expr =
-                    ScalarFunctionExpr::new("abs", scalar_def, args, return_type, None, false);
+                let expr = ScalarFunctionExpr::new("abs", math::abs(), args, return_type);
                 Ok(Arc::new(expr))
             }
             ExprStruct::CaseWhen(case_when) => {
@@ -690,8 +686,6 @@ impl PhysicalPlanner {
                     fun_expr,
                     vec![left, right],
                     data_type,
-                    None,
-                    false,
                 )))
             }
             _ => Ok(Arc::new(BinaryExpr::new(left, op, right))),
@@ -1209,26 +1203,18 @@ impl PhysicalPlanner {
                 }
             }
             AggExprStruct::First(expr) => {
-                let child = self.create_expr(expr.child.as_ref().unwrap(), schema)?;
-                let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
-                Ok(Arc::new(FirstValue::new(
-                    child,
-                    "first",
-                    datatype,
-                    vec![],
-                    vec![],
-                )))
+                let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;
+                let func = datafusion_expr::AggregateUDF::new_from_impl(FirstValue::new());
+
+                create_aggregate_expr(&func, &[child], &[], &[], &schema, "first", false, false)
+                    .map_err(|e| e.into())
             }
             AggExprStruct::Last(expr) => {
-                let child = self.create_expr(expr.child.as_ref().unwrap(), schema)?;
-                let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
-                Ok(Arc::new(LastValue::new(
-                    child,
-                    "last",
-                    datatype,
-                    vec![],
-                    vec![],
-                )))
+                let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;
+                let func = datafusion_expr::AggregateUDF::new_from_impl(LastValue::new());
+
+                create_aggregate_expr(&func, &[child], &[], &[], &schema, "last", false, false)
+                    .map_err(|e| e.into())
             }
             AggExprStruct::BitAndAgg(expr) => {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), schema)?;
@@ -1373,21 +1359,11 @@ impl PhysicalPlanner {
 
         let data_type = match expr.return_type.as_ref().map(to_arrow_datatype) {
             Some(t) => t,
-            None => {
-                // If no data type is provided from Spark, we'll use DF's return type from the
-                // scalar function
-                // Note this assumes the `fun_name` is a defined function in DF. Otherwise, it'll
-                // throw error.
-
-                if let Ok(fun) = BuiltinScalarFunction::from_str(fun_name) {
-                    fun.return_type(&input_expr_types)?
-                } else {
-                    self.session_ctx
-                        .udf(fun_name)?
-                        .inner()
-                        .return_type(&input_expr_types)?
-                }
-            }
+            None => self
+                .session_ctx
+                .udf(fun_name)?
+                .inner()
+                .return_type(&input_expr_types)?,
         };
 
         let fun_expr =
@@ -1398,8 +1374,6 @@ impl PhysicalPlanner {
             fun_expr,
             args.to_vec(),
             data_type,
-            None,
-            args.is_empty(),
         ));
 
         Ok(scalar_expr)
@@ -1444,7 +1418,7 @@ fn expr_to_columns(
     let mut left_field_indices: Vec<usize> = vec![];
     let mut right_field_indices: Vec<usize> = vec![];
 
-    expr.apply(&mut |expr| {
+    expr.apply(&mut |expr: &Arc<dyn PhysicalExpr>| {
         Ok({
             if let Some(column) = expr.as_any().downcast_ref::<Column>() {
                 if column.index() > left_field_len + right_field_len {
diff --git a/core/src/execution/datafusion/shuffle_writer.rs b/core/src/execution/datafusion/shuffle_writer.rs
index 99ac885b5..5afc9a53e 100644
--- a/core/src/execution/datafusion/shuffle_writer.rs
+++ b/core/src/execution/datafusion/shuffle_writer.rs
@@ -104,8 +104,8 @@ impl ExecutionPlan for ShuffleWriterExec {
         self.input.schema()
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
-        vec![self.input.clone()]
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/kernels/hash.rs b/core/src/execution/kernels/hash.rs
index de30f74cd..b39fd6224 100644
--- a/core/src/execution/kernels/hash.rs
+++ b/core/src/execution/kernels/hash.rs
@@ -22,6 +22,7 @@ use arrow_array::{
     downcast_dictionary_array, downcast_primitive_array, Array, ArrayAccessor, ArrayRef,
     ArrowPrimitiveType, PrimitiveArray,
 };
+use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
 use std::fmt::Debug;
 
 pub fn hash(src: &[ArrayRef], dst: &mut [u64]) {
@@ -169,3 +170,18 @@ impl Hashable for f64 {
         state.hash_one(u64::from_ne_bytes(self.to_ne_bytes()))
     }
 }
+
+impl Hashable for IntervalDayTime {
+    fn create_hash(&self, state: &RandomState) -> u64 {
+        state.hash_one(self.days);
+        state.hash_one(self.milliseconds)
+    }
+}
+
+impl Hashable for IntervalMonthDayNano {
+    fn create_hash(&self, state: &RandomState) -> u64 {
+        state.hash_one(self.months);
+        state.hash_one(self.days);
+        state.hash_one(self.nanoseconds)
+    }
+}
diff --git a/core/src/execution/operators/copy.rs b/core/src/execution/operators/copy.rs
index 96c244935..d011b3cb2 100644
--- a/core/src/execution/operators/copy.rs
+++ b/core/src/execution/operators/copy.rs
@@ -93,8 +93,8 @@ impl ExecutionPlan for CopyExec {
         self.schema.clone()
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
-        vec![self.input.clone()]
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
     }
 
     fn with_new_children(
diff --git a/core/src/execution/operators/scan.rs b/core/src/execution/operators/scan.rs
index 99c7c8391..bd518eda1 100644
--- a/core/src/execution/operators/scan.rs
+++ b/core/src/execution/operators/scan.rs
@@ -248,7 +248,7 @@ impl ExecutionPlan for ScanExec {
         scan_schema(input_batch, &self.data_types)
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         vec![]
     }
 
diff --git a/dev/ensure-jars-have-correct-contents.sh b/dev/ensure-jars-have-correct-contents.sh
index 12f555b8e..23d0be323 100755
--- a/dev/ensure-jars-have-correct-contents.sh
+++ b/dev/ensure-jars-have-correct-contents.sh
@@ -60,6 +60,14 @@ allowed_expr+="|^x86_64/libarrow_cdata_jni.so$"
 allowed_expr+="|^x86_64/libarrow_cdata_jni.dylib$"
 allowed_expr+="|^x86_64/arrow_cdata_jni.dll$"
 allowed_expr+="|^aarch_64/libarrow_cdata_jni.dylib$"
+
+allowed_expr+="|^arrow_cdata_jni/"
+allowed_expr+="|^arrow_cdata_jni/x86_64/"
+allowed_expr+="|^arrow_cdata_jni/aarch_64/"
+allowed_expr+="|^arrow_cdata_jni/x86_64/libarrow_cdata_jni.so$"
+allowed_expr+="|^arrow_cdata_jni/x86_64/libarrow_cdata_jni.dylib$"
+allowed_expr+="|^arrow_cdata_jni/x86_64/arrow_cdata_jni.dll$"
+allowed_expr+="|^arrow_cdata_jni/aarch_64/libarrow_cdata_jni.dylib$"
 # Two classes in Arrow C module: StructVectorLoader and StructVectorUnloader, are not
 # under org/apache/arrow/c, so we'll need to treat them specially.
 allowed_expr+="|^org/apache/arrow/$"
diff --git a/pom.xml b/pom.xml
index 34d949e14..fb42303c9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -53,7 +53,7 @@ under the License.
     <protobuf.version>3.19.6</protobuf.version>
     <parquet.version>1.13.1</parquet.version>
     <parquet.maven.scope>provided</parquet.maven.scope>
-    <arrow.version>14.0.2</arrow.version>
+    <arrow.version>16.0.0</arrow.version>
     <codehaus.jackson.version>1.9.13</codehaus.jackson.version>
     <spotless.version>2.43.0</spotless.version>
     <jacoco.version>0.8.11</jacoco.version>
@@ -120,6 +120,11 @@ under the License.
             <groupId>commons-logging</groupId>
             <artifactId>commons-logging</artifactId>
           </exclusion>
+
+          <exclusion>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>*</artifactId>
+          </exclusion>
         </exclusions>
       </dependency>
       <dependency>
@@ -138,7 +143,7 @@ under the License.
           <!-- Comet uses arrow-memory-unsafe -->
           <exclusion>
             <groupId>org.apache.arrow</groupId>
-            <artifactId>arrow-memory-netty</artifactId>
+            <artifactId>*</artifactId>
           </exclusion>
         </exclusions>
       </dependency>
@@ -258,7 +263,7 @@ under the License.
           <!-- We are using arrow-memory-unsafe -->
           <exclusion>
             <groupId>org.apache.arrow</groupId>
-            <artifactId>arrow-memory-netty</artifactId>
+            <artifactId>*</artifactId>
           </exclusion>
 
           <!-- We're using "org.slf4j:jcl-over-slf4j" -->
@@ -284,7 +289,7 @@ under the License.
           <!-- We are using arrow-memory-unsafe -->
           <exclusion>
             <groupId>org.apache.arrow</groupId>
-            <artifactId>arrow-memory-netty</artifactId>
+            <artifactId>*</artifactId>
           </exclusion>
 
           <!-- We're using "org.slf4j:jcl-over-slf4j" -->
@@ -324,6 +329,11 @@ under the License.
             <groupId>commons-logging</groupId>
             <artifactId>commons-logging</artifactId>
           </exclusion>
+
+          <exclusion>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>*</artifactId>
+          </exclusion>
         </exclusions>
       </dependency>
       <dependency>
@@ -340,6 +350,10 @@ under the License.
             <groupId>commons-logging</groupId>
             <artifactId>commons-logging</artifactId>
           </exclusion>
+          <exclusion>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>*</artifactId>
+          </exclusion>
         </exclusions>
       </dependency>
       <dependency>
@@ -400,6 +414,11 @@ under the License.
             <groupId>commons-logging</groupId>
             <artifactId>commons-logging</artifactId>
           </exclusion>
+
+          <exclusion>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>*</artifactId>
+          </exclusion>
         </exclusions>
       </dependency>
 

From f75aeefab58dc7e14cb70742b9d7bb656b727dbd Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Fri, 7 Jun 2024 15:15:23 -0600
Subject: [PATCH 7/8] docs: Improve user documentation for supported operators
 and expressions (#520)

* Improve documentation about supported operators and expressions

* Improve documentation about supported operators and expressions

* more notes

* Add more supported expressions

* rename protobuf Negative to UnaryMinus for consistency

* format

* remove duplicate ASF header

* SMJ not disabled by default

* Update docs/source/user-guide/operators.md

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>

* Update docs/source/user-guide/operators.md

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>

* remove RLike

---------

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>
---
 core/src/execution/datafusion/planner.rs      |   2 +-
 core/src/execution/proto/expr.proto           |   4 +-
 docs/source/user-guide/expressions.md         | 267 +++++++++++-------
 docs/source/user-guide/operators.md           |  27 +-
 .../apache/comet/serde/QueryPlanSerde.scala   |   4 +-
 5 files changed, 192 insertions(+), 112 deletions(-)

diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
index 7d504f878..6c7ea0de4 100644
--- a/core/src/execution/datafusion/planner.rs
+++ b/core/src/execution/datafusion/planner.rs
@@ -566,7 +566,7 @@ impl PhysicalPlanner {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
                 Ok(Arc::new(NotExpr::new(child)))
             }
-            ExprStruct::Negative(expr) => {
+            ExprStruct::UnaryMinus(expr) => {
                 let child: Arc<dyn PhysicalExpr> =
                     self.create_expr(expr.child.as_ref().unwrap(), input_schema.clone())?;
                 let result = negative::create_negate_expr(child, expr.fail_on_error);
diff --git a/core/src/execution/proto/expr.proto b/core/src/execution/proto/expr.proto
index 9c6049013..5192bbd4c 100644
--- a/core/src/execution/proto/expr.proto
+++ b/core/src/execution/proto/expr.proto
@@ -65,7 +65,7 @@ message Expr {
     CaseWhen caseWhen = 38;
     In in = 39;
     Not not = 40;
-    Negative negative = 41;
+    UnaryMinus unary_minus = 41;
     BitwiseShiftRight bitwiseShiftRight = 42;
     BitwiseShiftLeft bitwiseShiftLeft = 43;
     IfExpr if = 44;
@@ -452,7 +452,7 @@ message Not {
   Expr child = 1;
 }
 
-message Negative {
+message UnaryMinus {
   Expr child = 1;
   bool fail_on_error = 2;
 }
diff --git a/docs/source/user-guide/expressions.md b/docs/source/user-guide/expressions.md
index 775ebedb6..14b6f18d0 100644
--- a/docs/source/user-guide/expressions.md
+++ b/docs/source/user-guide/expressions.md
@@ -19,99 +19,174 @@
 
 # Supported Spark Expressions
 
-The following Spark expressions are currently available:
-
-- Literals
-- Arithmetic Operators
-  - UnaryMinus
-  - Add/Minus/Multiply/Divide/Remainder
-- Conditional functions
-  - Case When
-  - If
-- Cast
-- Coalesce
-- BloomFilterMightContain
-- Boolean functions
-  - And
-  - Or
-  - Not
-  - EqualTo
-  - EqualNullSafe
-  - GreaterThan
-  - GreaterThanOrEqual
-  - LessThan
-  - LessThanOrEqual
-  - IsNull
-  - IsNotNull
-  - In
-- String functions
-  - Substring
-  - Coalesce
-  - StringSpace
-  - Like
-  - Contains
-  - Startswith
-  - Endswith
-  - Ascii
-  - Bit_length
-  - Octet_length
-  - Upper
-  - Lower
-  - Chr
-  - Initcap
-  - Trim/Btrim/Ltrim/Rtrim
-  - Concat_ws
-  - Repeat
-  - Length
-  - Reverse
-  - Instr
-  - Replace
-  - Translate
-- Bitwise functions
-  - Shiftright/Shiftleft
-- Date/Time functions
-  - Year/Hour/Minute/Second
-- Hash functions
-  - Md5
-  - Sha2
-  - Hash
-  - Xxhash64
-- Math functions
-  - Abs
-  - Acos
-  - Asin
-  - Atan
-  - Atan2
-  - Cos
-  - Exp
-  - Ln
-  - Log10
-  - Log2
-  - Pow
-  - Round
-  - Signum
-  - Sin
-  - Sqrt
-  - Tan
-  - Ceil
-  - Floor
-- Aggregate functions
-  - Count
-  - Sum
-  - Max
-  - Min
-  - Avg
-  - First
-  - Last
-  - BitAnd
-  - BitOr
-  - BitXor
-  - BoolAnd
-  - BoolOr
-  - CovPopulation
-  - CovSample
-  - VariancePop
-  - VarianceSamp
-  - StddevPop
-  - StddevSamp
-  - Corr
+The following Spark expressions are currently available. Any known compatibility issues are noted in the following tables.
+
+## Literal Values
+
+| Expression                             | Notes |
+| -------------------------------------- | ----- |
+| Literal values of supported data types |       |
+
+## Unary Arithmetic
+
+| Expression       | Notes |
+| ---------------- | ----- |
+| UnaryMinus (`-`) |       |
+
+## Binary Arithmeticx
+
+| Expression      | Notes                                               |
+| --------------- | --------------------------------------------------- |
+| Add (`+`)       |                                                     |
+| Subtract (`-`)  |                                                     |
+| Multiply (`*`)  |                                                     |
+| Divide (`/`)    |                                                     |
+| Remainder (`%`) | Comet produces `NaN` instead of `NULL` for `% -0.0` |
+
+## Conditional Expressions
+
+| Expression | Notes |
+| ---------- | ----- |
+| CaseWhen   |       |
+| If         |       |
+
+## Comparison
+
+| Expression                | Notes |
+| ------------------------- | ----- |
+| EqualTo (`=`)             |       |
+| EqualNullSafe (`<=>`)     |       |
+| GreaterThan (`>`)         |       |
+| GreaterThanOrEqual (`>=`) |       |
+| LessThan (`<`)            |       |
+| LessThanOrEqual (`<=`)    |       |
+| IsNull (`IS NULL`)        |       |
+| IsNotNull (`IS NOT NULL`) |       |
+| In (`IN`)                 |       |
+
+## String Functions
+
+| Expression      | Notes                                                                                                       |
+| --------------- | ----------------------------------------------------------------------------------------------------------- |
+| Ascii           |                                                                                                             |
+| BitLength       |                                                                                                             |
+| Chr             |                                                                                                             |
+| ConcatWs        |                                                                                                             |
+| Contains        |                                                                                                             |
+| EndsWith        |                                                                                                             |
+| InitCap         |                                                                                                             |
+| Instr           |                                                                                                             |
+| Length          |                                                                                                             |
+| Like            |                                                                                                             |
+| Lower           |                                                                                                             |
+| OctetLength     |                                                                                                             |
+| Repeat          | Negative argument for number of times to repeat causes exception                                            |
+| Replace         |                                                                                                             |
+| Reverse         |                                                                                                             |
+| StartsWith      |                                                                                                             |
+| StringSpace     |                                                                                                             |
+| StringTrim      |                                                                                                             |
+| StringTrimBoth  |                                                                                                             |
+| StringTrimLeft  |                                                                                                             |
+| StringTrimRight |                                                                                                             |
+| Substring       |                                                                                                             |
+| Translate       |                                                                                                             |
+| Upper           |                                                                                                             |
+
+## Date/Time Functions
+
+| Expression     | Notes                    |
+| -------------- | ------------------------ |
+| DatePart       | Only `year` is supported |
+| Extract        | Only `year` is supported |
+| Hour           |                          |
+| Minute         |                          |
+| Second         |                          |
+| TruncDate      |                          |
+| TruncTimestamp |                          |
+| Year           |                          |
+
+## Math Expressions
+
+| Expression | Notes                                                               |
+| ---------- | ------------------------------------------------------------------- |
+| Abs        |                                                                     |
+| Acos       |                                                                     |
+| Asin       |                                                                     |
+| Atan       |                                                                     |
+| Atan2      |                                                                     |
+| Ceil       |                                                                     |
+| Cos        |                                                                     |
+| Exp        |                                                                     |
+| Floor      |                                                                     |
+| Log        | log(0) will produce `-Infinity` unlike Spark which returns `null`   |
+| Log2       | log2(0) will produce `-Infinity` unlike Spark which returns `null`  |
+| Log10      | log10(0) will produce `-Infinity` unlike Spark which returns `null` |
+| Pow        |                                                                     |
+| Round      |                                                                     |
+| Signum     | Signum does not differentiate between `0.0` and `-0.0`              |
+| Sin        |                                                                     |
+| Sqrt       |                                                                     |
+| Tan        |                                                                     |
+
+## Hashing Functions
+
+| Expression | Notes |
+| ---------- | ----- |
+| Md5        |       |
+| Hash       |       |
+| Sha2       |       |
+| XxHash64   |       |
+
+## Boolean Expressions
+
+| Expression | Notes |
+| ---------- | ----- |
+| And        |       |
+| Or         |       |
+| Not        |       |
+
+## Bitwise Expressions
+
+| Expression           | Notes |
+| -------------------- | ----- |
+| ShiftLeft (`<<`)     |       |
+| ShiftRight (`>>`)    |       |
+| BitAnd (`&`)         |       |
+| BitOr (`\|`)         |       |
+| BitXor (`^`)         |       |
+| BitwiseNot (`~`)     |       |
+| BoolAnd (`bool_and`) |       |
+| BoolOr (`bool_or`)   |       |
+
+## Aggregate Expressions
+
+| Expression    | Notes |
+| ------------- | ----- |
+| Avg           |       |
+| BitAndAgg     |       |
+| BitOrAgg      |       |
+| BitXorAgg     |       |
+| Corr          |       |
+| Count         |       |
+| CovPopulation |       |
+| CovSample     |       |
+| First         |       |
+| Last          |       |
+| Max           |       |
+| Min           |       |
+| StddevPop     |       |
+| StddevSamp    |       |
+| Sum           |       |
+| VariancePop   |       |
+| VarianceSamp  |       |
+
+## Other
+
+| Expression              | Notes                                                                           |
+| ----------------------- | ------------------------------------------------------------------------------- |
+| Cast                    | See compatibility guide for list of supported cast expressions and known issues |
+| BloomFilterMightContain |                                                                                 |
+| ScalarSubquery          |                                                                                 |
+| Coalesce                |                                                                                 |
+| NormalizeNaNAndZero     |                                                                                 |
diff --git a/docs/source/user-guide/operators.md b/docs/source/user-guide/operators.md
index ec82e9f69..e3a3ac522 100644
--- a/docs/source/user-guide/operators.md
+++ b/docs/source/user-guide/operators.md
@@ -19,15 +19,20 @@
 
 # Supported Spark Operators
 
-The following Spark operators are currently available:
+The following Spark operators are currently replaced with native versions. Query stages that contain any operators
+not supported by Comet will fall back to regular Spark execution.
 
-- FileSourceScanExec/BatchScanExec for Parquet
-- Projection
-- Filter
-- Sort
-- Hash Aggregate
-- Limit
-- Sort-merge Join
-- Hash Join
-- Shuffle
-- Expand
+| Operator                                     | Notes |
+| -------------------------------------------- | ----- |
+| FileSourceScanExec/BatchScanExec for Parquet |       |
+| Projection                                   |       |
+| Filter                                       |       |
+| Sort                                         |       |
+| Hash Aggregate                               |       |
+| Limit                                        |       |
+| Sort-merge Join                              |       |
+| Hash Join                                    |       |
+| BroadcastHashJoinExec                        |       |
+| Shuffle                                      |       |
+| Expand                                       |       |
+| Union                                        |       |
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 8d81b57c4..448c4ff0f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -1987,13 +1987,13 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         case UnaryMinus(child, failOnError) =>
           val childExpr = exprToProtoInternal(child, inputs)
           if (childExpr.isDefined) {
-            val builder = ExprOuterClass.Negative.newBuilder()
+            val builder = ExprOuterClass.UnaryMinus.newBuilder()
             builder.setChild(childExpr.get)
             builder.setFailOnError(failOnError)
             Some(
               ExprOuterClass.Expr
                 .newBuilder()
-                .setNegative(builder)
+                .setUnaryMinus(builder)
                 .build())
           } else {
             withInfo(expr, child)

From 311e13e3ec5fe40e10e0f2671e317e1142272af4 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Fri, 7 Jun 2024 15:15:47 -0600
Subject: [PATCH 8/8] chore: Add CometEvalMode enum to replace string literals
 (#539)

* Add CometEvalMode enum

* address feedback
---
 .../scala/org/apache/comet/GenerateDocs.scala |  6 +-
 .../apache/comet/expressions/CometCast.scala  |  4 +-
 .../comet/expressions/CometEvalMode.scala     | 42 ++++++++++++++
 .../apache/comet/serde/QueryPlanSerde.scala   | 57 +++++++------------
 .../apache/comet/shims/CometExprShim.scala    |  6 +-
 .../apache/comet/shims/CometExprShim.scala    |  5 +-
 .../apache/comet/shims/CometExprShim.scala    | 15 ++++-
 .../apache/comet/shims/CometExprShim.scala    | 13 +++++
 .../org/apache/comet/CometCastSuite.scala     |  4 +-
 9 files changed, 107 insertions(+), 45 deletions(-)
 create mode 100644 spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala

diff --git a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
index a2d5e2515..fb86389fe 100644
--- a/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
+++ b/spark/src/main/scala/org/apache/comet/GenerateDocs.scala
@@ -25,7 +25,7 @@ import scala.io.Source
 
 import org.apache.spark.sql.catalyst.expressions.Cast
 
-import org.apache.comet.expressions.{CometCast, Compatible, Incompatible}
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible, Incompatible}
 
 /**
  * Utility for generating markdown documentation from the configs.
@@ -72,7 +72,7 @@ object GenerateDocs {
             if (Cast.canCast(fromType, toType) && fromType != toType) {
               val fromTypeName = fromType.typeName.replace("(10,2)", "")
               val toTypeName = toType.typeName.replace("(10,2)", "")
-              CometCast.isSupported(fromType, toType, None, "LEGACY") match {
+              CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) match {
                 case Compatible(notes) =>
                   val notesStr = notes.getOrElse("").trim
                   w.write(s"| $fromTypeName | $toTypeName | $notesStr |\n".getBytes)
@@ -89,7 +89,7 @@ object GenerateDocs {
             if (Cast.canCast(fromType, toType) && fromType != toType) {
               val fromTypeName = fromType.typeName.replace("(10,2)", "")
               val toTypeName = toType.typeName.replace("(10,2)", "")
-              CometCast.isSupported(fromType, toType, None, "LEGACY") match {
+              CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) match {
                 case Incompatible(notes) =>
                   val notesStr = notes.getOrElse("").trim
                   w.write(s"| $fromTypeName | $toTypeName  | $notesStr |\n".getBytes)
diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
index 11c5a53cc..811c61d46 100644
--- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
+++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
@@ -55,7 +55,7 @@ object CometCast {
       fromType: DataType,
       toType: DataType,
       timeZoneId: Option[String],
-      evalMode: String): SupportLevel = {
+      evalMode: CometEvalMode.Value): SupportLevel = {
 
     if (fromType == toType) {
       return Compatible()
@@ -102,7 +102,7 @@ object CometCast {
   private def canCastFromString(
       toType: DataType,
       timeZoneId: Option[String],
-      evalMode: String): SupportLevel = {
+      evalMode: CometEvalMode.Value): SupportLevel = {
     toType match {
       case DataTypes.BooleanType =>
         Compatible()
diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala
new file mode 100644
index 000000000..59e9c89a6
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala
@@ -0,0 +1,42 @@
+/*
+ * 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.expressions
+
+/**
+ * We cannot reference Spark's EvalMode directly because the package is different between Spark
+ * versions, so we copy it here.
+ *
+ * Expression evaluation modes.
+ *   - LEGACY: the default evaluation mode, which is compliant to Hive SQL.
+ *   - ANSI: a evaluation mode which is compliant to ANSI SQL standard.
+ *   - TRY: a evaluation mode for `try_*` functions. It is identical to ANSI evaluation mode
+ *     except for returning null result on errors.
+ */
+object CometEvalMode extends Enumeration {
+  val LEGACY, ANSI, TRY = Value
+
+  def fromBoolean(ansiEnabled: Boolean): Value = if (ansiEnabled) {
+    ANSI
+  } else {
+    LEGACY
+  }
+
+  def fromString(str: String): CometEvalMode.Value = CometEvalMode.withName(str)
+}
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 448c4ff0f..ed3f2fae6 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -19,8 +19,6 @@
 
 package org.apache.comet.serde
 
-import java.util.Locale
-
 import scala.collection.JavaConverters._
 
 import org.apache.spark.internal.Logging
@@ -45,7 +43,7 @@ import org.apache.spark.unsafe.types.UTF8String
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.{isCometOperatorEnabled, isCometScan, isSpark32, isSpark34Plus, withInfo}
-import org.apache.comet.expressions.{CometCast, Compatible, Incompatible, Unsupported}
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible, Incompatible, Unsupported}
 import org.apache.comet.serde.ExprOuterClass.{AggExpr, DataType => ProtoDataType, Expr, ScalarFunc}
 import org.apache.comet.serde.ExprOuterClass.DataType.{DataTypeInfo, DecimalInfo, ListInfo, MapInfo, StructInfo}
 import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, JoinType, Operator}
@@ -578,6 +576,15 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     }
   }
 
+  def evalModeToProto(evalMode: CometEvalMode.Value): ExprOuterClass.EvalMode = {
+    evalMode match {
+      case CometEvalMode.LEGACY => ExprOuterClass.EvalMode.LEGACY
+      case CometEvalMode.TRY => ExprOuterClass.EvalMode.TRY
+      case CometEvalMode.ANSI => ExprOuterClass.EvalMode.ANSI
+      case _ => throw new IllegalStateException(s"Invalid evalMode $evalMode")
+    }
+  }
+
   /**
    * Convert a Spark expression to protobuf.
    *
@@ -590,18 +597,6 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
    * @return
    *   The protobuf representation of the expression, or None if the expression is not supported
    */
-
-  def stringToEvalMode(evalModeStr: String): ExprOuterClass.EvalMode =
-    evalModeStr.toUpperCase(Locale.ROOT) match {
-      case "LEGACY" => ExprOuterClass.EvalMode.LEGACY
-      case "TRY" => ExprOuterClass.EvalMode.TRY
-      case "ANSI" => ExprOuterClass.EvalMode.ANSI
-      case invalid =>
-        throw new IllegalArgumentException(
-          s"Invalid eval mode '$invalid' "
-        ) // Assuming we want to catch errors strictly
-    }
-
   def exprToProto(
       expr: Expression,
       input: Seq[Attribute],
@@ -610,15 +605,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         timeZoneId: Option[String],
         dt: DataType,
         childExpr: Option[Expr],
-        evalMode: String): Option[Expr] = {
+        evalMode: CometEvalMode.Value): Option[Expr] = {
       val dataType = serializeDataType(dt)
-      val evalModeEnum = stringToEvalMode(evalMode) // Convert string to enum
 
       if (childExpr.isDefined && dataType.isDefined) {
         val castBuilder = ExprOuterClass.Cast.newBuilder()
         castBuilder.setChild(childExpr.get)
         castBuilder.setDatatype(dataType.get)
-        castBuilder.setEvalMode(evalModeEnum) // Set the enum in protobuf
+        castBuilder.setEvalMode(evalModeToProto(evalMode))
 
         val timeZone = timeZoneId.getOrElse("UTC")
         castBuilder.setTimezone(timeZone)
@@ -646,26 +640,26 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           inputs: Seq[Attribute],
           dt: DataType,
           timeZoneId: Option[String],
-          actualEvalModeStr: String): Option[Expr] = {
+          evalMode: CometEvalMode.Value): Option[Expr] = {
 
         val childExpr = exprToProtoInternal(child, inputs)
         if (childExpr.isDefined) {
           val castSupport =
-            CometCast.isSupported(child.dataType, dt, timeZoneId, actualEvalModeStr)
+            CometCast.isSupported(child.dataType, dt, timeZoneId, evalMode)
 
           def getIncompatMessage(reason: Option[String]): String =
             "Comet does not guarantee correct results for cast " +
               s"from ${child.dataType} to $dt " +
-              s"with timezone $timeZoneId and evalMode $actualEvalModeStr" +
+              s"with timezone $timeZoneId and evalMode $evalMode" +
               reason.map(str => s" ($str)").getOrElse("")
 
           castSupport match {
             case Compatible(_) =>
-              castToProto(timeZoneId, dt, childExpr, actualEvalModeStr)
+              castToProto(timeZoneId, dt, childExpr, evalMode)
             case Incompatible(reason) =>
               if (CometConf.COMET_CAST_ALLOW_INCOMPATIBLE.get()) {
                 logWarning(getIncompatMessage(reason))
-                castToProto(timeZoneId, dt, childExpr, actualEvalModeStr)
+                castToProto(timeZoneId, dt, childExpr, evalMode)
               } else {
                 withInfo(
                   expr,
@@ -677,7 +671,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
               withInfo(
                 expr,
                 s"Unsupported cast from ${child.dataType} to $dt " +
-                  s"with timezone $timeZoneId and evalMode $actualEvalModeStr")
+                  s"with timezone $timeZoneId and evalMode $evalMode")
               None
           }
         } else {
@@ -701,17 +695,10 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
 
         case UnaryExpression(child) if expr.prettyName == "trycast" =>
           val timeZoneId = SQLConf.get.sessionLocalTimeZone
-          handleCast(child, inputs, expr.dataType, Some(timeZoneId), "TRY")
+          handleCast(child, inputs, expr.dataType, Some(timeZoneId), CometEvalMode.TRY)
 
-        case Cast(child, dt, timeZoneId, evalMode) =>
-          val evalModeStr = if (evalMode.isInstanceOf[Boolean]) {
-            // Spark 3.2 & 3.3 has ansiEnabled boolean
-            if (evalMode.asInstanceOf[Boolean]) "ANSI" else "LEGACY"
-          } else {
-            // Spark 3.4+ has EvalMode enum with values LEGACY, ANSI, and TRY
-            evalMode.toString
-          }
-          handleCast(child, inputs, dt, timeZoneId, evalModeStr)
+        case c @ Cast(child, dt, timeZoneId, _) =>
+          handleCast(child, inputs, dt, timeZoneId, evalMode(c))
 
         case add @ Add(left, right, _) if supportedDataType(left.dataType) =>
           val leftExpr = exprToProtoInternal(left, inputs)
@@ -2006,7 +1993,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           // TODO: Remove this once we have new DataFusion release which includes
           // the fix: https://github.com/apache/arrow-datafusion/pull/9459
           if (childExpr.isDefined) {
-            castToProto(None, a.dataType, childExpr, "LEGACY")
+            castToProto(None, a.dataType, childExpr, CometEvalMode.LEGACY)
           } else {
             withInfo(expr, a.children: _*)
             None
diff --git a/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala
index f5a578f82..2c6f6ccf4 100644
--- a/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala
+++ b/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala
@@ -18,6 +18,7 @@
  */
 package org.apache.comet.shims
 
+import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
 
 /**
@@ -27,7 +28,10 @@ trait CometExprShim {
     /**
      * Returns a tuple of expressions for the `unhex` function.
      */
-    def unhexSerde(unhex: Unhex): (Expression, Expression) = {
+    protected def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    protected def evalMode(c: Cast): CometEvalMode.Value = CometEvalMode.fromBoolean(c.ansiEnabled)
 }
+
diff --git a/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala
index f5a578f82..150656c23 100644
--- a/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala
+++ b/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala
@@ -18,6 +18,7 @@
  */
 package org.apache.comet.shims
 
+import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
 
 /**
@@ -27,7 +28,9 @@ trait CometExprShim {
     /**
      * Returns a tuple of expressions for the `unhex` function.
      */
-    def unhexSerde(unhex: Unhex): (Expression, Expression) = {
+    protected def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    protected def evalMode(c: Cast): CometEvalMode.Value = CometEvalMode.fromBoolean(c.ansiEnabled)
 }
diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
index 3f2301f0a..5f4e3fba2 100644
--- a/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
+++ b/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
@@ -18,6 +18,7 @@
  */
 package org.apache.comet.shims
 
+import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
 
 /**
@@ -27,7 +28,19 @@ trait CometExprShim {
     /**
      * Returns a tuple of expressions for the `unhex` function.
      */
-    def unhexSerde(unhex: Unhex): (Expression, Expression) = {
+    protected def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(unhex.failOnError))
     }
+
+    protected def evalMode(c: Cast): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 }
+
+object CometEvalModeUtil {
+    def fromSparkEvalMode(evalMode: EvalMode.Value): CometEvalMode.Value = evalMode match {
+        case EvalMode.LEGACY => CometEvalMode.LEGACY
+        case EvalMode.TRY => CometEvalMode.TRY
+        case EvalMode.ANSI => CometEvalMode.ANSI
+    }
+}
+
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 01f923206..5f4e3fba2 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
@@ -18,6 +18,7 @@
  */
 package org.apache.comet.shims
 
+import org.apache.comet.expressions.CometEvalMode
 import org.apache.spark.sql.catalyst.expressions._
 
 /**
@@ -30,4 +31,16 @@ trait CometExprShim {
     protected def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(unhex.failOnError))
     }
+
+    protected def evalMode(c: Cast): CometEvalMode.Value =
+        CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 }
+
+object CometEvalModeUtil {
+    def fromSparkEvalMode(evalMode: EvalMode.Value): CometEvalMode.Value = evalMode match {
+        case EvalMode.LEGACY => CometEvalMode.LEGACY
+        case EvalMode.TRY => CometEvalMode.TRY
+        case EvalMode.ANSI => CometEvalMode.ANSI
+    }
+}
+
diff --git a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
index fd2218965..25343f933 100644
--- a/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCastSuite.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataType, DataTypes, DecimalType}
 
-import org.apache.comet.expressions.{CometCast, Compatible}
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible}
 
 class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
@@ -76,7 +76,7 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
             } else {
               val testIgnored =
                 tags.get(expectedTestName).exists(s => s.contains("org.scalatest.Ignore"))
-              CometCast.isSupported(fromType, toType, None, "LEGACY") match {
+              CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) match {
                 case Compatible(_) =>
                   if (testIgnored) {
                     fail(