fix: Spark 4.0 SparkArithmeticException test

common/src/main/scala/org/apache/spark/sql/comet/CastOverflowException.scala

@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import org.apache.spark.sql.comet.shims.ShimCastOverflowException
+
+class CastOverflowException(t: String, from: String, to: String)
+    extends ShimCastOverflowException(t, from, to) {}

common/src/main/spark-3.3/org/apache/spark/sql/comet/shims/ShimCastOverflowException.scala

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.shims
+
+import org.apache.spark.SparkArithmeticException
+import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLConf
+import org.apache.spark.sql.internal.SQLConf
+
+// TODO: Only the Spark 3.3 version of this class is different from the others.
+//       Remove this class after dropping Spark 3.3 support.
+class ShimCastOverflowException(t: String, from: String, to: String)
+  extends SparkArithmeticException(
+    "CAST_OVERFLOW",
+    Array(t, s""""$from"""", s""""$to"""", toSQLConf(SQLConf.ANSI_ENABLED.key))
+  ) {}

common/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCastOverflowException.scala

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.shims
+
+import org.apache.spark.SparkArithmeticException
+import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLConf
+import org.apache.spark.sql.internal.SQLConf
+
+// TODO: Only the Spark 3.3 version of this class is different from the others.
+//       Remove this class after dropping Spark 3.3 support.
+class ShimCastOverflowException(t: String, from: String, to: String)
+  extends SparkArithmeticException(
+    "CAST_OVERFLOW",
+    Map(
+      "value" -> t,
+      "sourceType" -> s""""$from"""",
+      "targetType" -> s""""$to"""",
+      "ansiConfig" -> toSQLConf(SQLConf.ANSI_ENABLED.key)),
+    Array.empty,
+    "") {}

common/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCastOverflowException.scala

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.shims
+
+import org.apache.spark.SparkArithmeticException
+import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLConf
+import org.apache.spark.sql.internal.SQLConf
+
+// TODO: Only the Spark 3.3 version of this class is different from the others.
+//       Remove this class after dropping Spark 3.3 support.
+class ShimCastOverflowException(t: String, from: String, to: String)
+  extends SparkArithmeticException(
+    "CAST_OVERFLOW",
+    Map(
+      "value" -> t,
+      "sourceType" -> s""""$from"""",
+      "targetType" -> s""""$to"""",
+      "ansiConfig" -> toSQLConf(SQLConf.ANSI_ENABLED.key)),
+    Array.empty,
+    "") {}

common/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCastOverflowException.scala

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.shims
+
+import org.apache.spark.SparkArithmeticException
+import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLConf
+import org.apache.spark.sql.internal.SQLConf
+
+// TODO: Only the Spark 3.3 version of this class is different from the others.
+//       Remove this class after dropping Spark 3.3 support.
+class ShimCastOverflowException(t: String, from: String, to: String)
+  extends SparkArithmeticException(
+    "CAST_OVERFLOW",
+    Map(
+      "value" -> t,
+      "sourceType" -> s""""$from"""",
+      "targetType" -> s""""$to"""",
+      "ansiConfig" -> toSQLConf(SQLConf.ANSI_ENABLED.key)),
+    Array.empty,
+    "") {}

dev/diffs/4.0.0-preview1.diff

@@ -846,7 +846,7 @@ index 34c6c49bc49..f5dea07a213 100644
    protected val baseResourcePath = {
      // use the same way as `SQLQueryTestSuite` to get the resource path
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
-index 56c364e2084..11779ee3b4b 100644
+index 56c364e2084..fc3abd7cdc4 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
 @@ -1510,7 +1510,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
@@ -889,16 +889,6 @@ index 56c364e2084..11779ee3b4b 100644
      withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
        SQLConf.ANSI_ENABLED.key -> "true") {
        withTable("t") {
-@@ -4639,7 +4643,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
-       Row(1, 2, 3, 1, 2, 3, 1, 1))
-   }
-
--  test("SPARK-40389: Don't eliminate a cast which can cause overflow") {
-+  test("SPARK-40389: Don't eliminate a cast which can cause overflow",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/551")) {
-     withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
-       withTable("dt") {
-         sql("create table dt using parquet as select 9000000000BD as d")
 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 68f14f13bbd..4b8e967102f 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala

native/core/src/errors.rs

@@ -39,7 +39,7 @@ use jni::sys::{jboolean, jbyte, jchar, jdouble, jfloat, jint, jlong, jobject, js
 
 use crate::execution::operators::ExecutionError;
 use datafusion_comet_spark_expr::SparkError;
-use jni::objects::{GlobalRef, JThrowable};
+use jni::objects::{GlobalRef, JThrowable, JValue};
 use jni::JNIEnv;
 use lazy_static::lazy_static;
 use parquet::errors::ParquetError;
@@ -383,6 +383,33 @@ fn throw_exception(env: &mut JNIEnv, error: &CometError, backtrace: Option<Strin
                         throwable,
                     },
             } => env.throw(<&JThrowable>::from(throwable.as_obj())),
+            CometError::DataFusion {
+                msg: _,
+                source: DataFusionError::External(e),
+            } if matches!(e.downcast_ref(), Some(SparkError::CastOverFlow { .. })) => {
+                match e.downcast_ref() {
+                    Some(SparkError::CastOverFlow {
+                        value,
+                        from_type,
+                        to_type,
+                    }) => {
+                        let throwable: JThrowable = env
+                            .new_object(
+                                "org/apache/spark/sql/comet/CastOverflowException",
+                                "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V",
+                                &[
+                                    JValue::Object(&env.new_string(value).unwrap()),
+                                    JValue::Object(&env.new_string(from_type).unwrap()),
+                                    JValue::Object(&env.new_string(to_type).unwrap()),
+                                ],
+                            )
+                            .unwrap()
+                            .into();
+                        env.throw(throwable)
+                    }
+                    _ => unreachable!(),
+                }
+            }
             _ => {
                 let exception = error.to_exception();
                 match backtrace {

spark/src/test/scala/org/apache/comet/CometCastSuite.scala

@@ -988,8 +988,11 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
           case (Some(sparkException), Some(cometException)) =>
             // both systems threw an exception so we make sure they are the same
             val sparkMessage =
-              if (sparkException.getCause != null) sparkException.getCause.getMessage else null
-            val cometMessage = cometException.getCause.getMessage
+              if (sparkException.getCause != null) sparkException.getCause.getMessage
+              else sparkException.getMessage
+            val cometMessage =
+              if (cometException.getCause != null) cometException.getCause.getMessage
+              else cometException.getMessage
             if (CometSparkSessionExtensions.isSpark40Plus) {
               // for Spark 4 we expect to sparkException carries the message
               assert(