diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index fb3d7371a77..1642013f1c0 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -1088,25 +1088,27 @@ index 525d97e4998..aded8906d75 100644 val hiveVersion = "2.3.9" // transitive=false, only download specified jar diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -index 2dabcf01be7..8fcec0d1ce4 100644 +index 2dabcf01be7..9c326b3732d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -@@ -491,8 +491,23 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper +@@ -491,8 +491,25 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper s"Schema did not match for query #$i\n${expected.sql}: $output") { output.schema } - assertResult(expected.output, s"Result did not match" + - s" for query #$i\n${expected.sql}") { output.output } -+ // Comet may surface errors as `CometNativeException` instead of the matching Spark -+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a -+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it). -+ // Same category, different surface. Collapse both sides to a placeholder when this happens -+ // so the literal compare passes. ++ // With parquet filter pushdown, DataFusion 54/55 formats an error raised while ++ // evaluating a pushed-down predicate into a string, so a divide-by-zero from a ++ // `ScalaUDF` dispatched into the native plan surfaces as `CometNativeException` ++ // instead of `SparkArithmeticException`. Collapse both sides to a placeholder in that ++ // case only. Remove once Comet uses a DataFusion release that includes ++ // https://github.com/apache/datafusion/pull/24638. + // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517 + val (expectedOut, actualOut) = if (isCometEnabled && + expected.output.startsWith("org.apache.spark.SparkArithmeticException") && + expected.output.contains("\"DIVIDE_BY_ZERO\"") && + output.output.startsWith("org.apache.comet.CometNativeException") && ++ output.output.contains("Error evaluating filter predicate") && + output.output.contains("DivideByZero")) { + ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]") + } else { diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index e9baa8b4650..e0dbd3b5568 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -1277,24 +1277,26 @@ index 3cf2bfd17ab..11d3ca4d5a9 100644 val hiveVersion = "2.3.9" // transitive=false, only download specified jar diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -index 71af1fd69c3..81a04c93c9c 100644 +index 71af1fd69c3..8a8b49f007c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -@@ -872,9 +872,24 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper +@@ -872,9 +872,26 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper s"Schema did not match for query #$i\n${expected.sql}: $output") { output.schema } - assertResult(expected.output, s"Result did not match" + -+ // Comet may surface errors as `CometNativeException` instead of the matching Spark -+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a -+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it). -+ // Same category, different surface. Collapse both sides to a placeholder when this happens -+ // so the literal compare passes. ++ // With parquet filter pushdown, DataFusion 54/55 formats an error raised while ++ // evaluating a pushed-down predicate into a string, so a divide-by-zero from a ++ // `ScalaUDF` dispatched into the native plan surfaces as `CometNativeException` ++ // instead of `SparkArithmeticException`. Collapse both sides to a placeholder in that ++ // case only. Remove once Comet uses a DataFusion release that includes ++ // https://github.com/apache/datafusion/pull/24638. + // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517 + val (expectedOut, actualOut) = if (isCometEnabled && + expected.output.startsWith("org.apache.spark.SparkArithmeticException") && + expected.output.contains("\"DIVIDE_BY_ZERO\"") && + output.output.startsWith("org.apache.comet.CometNativeException") && ++ output.output.contains("Error evaluating filter predicate") && + output.output.contains("DivideByZero")) { + ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]") + } else { diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index ede5e39dd36..f2752beca12 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -1447,24 +1447,26 @@ index f294ff81021..f019c6de8e9 100644 df.collect() .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -index 575a4ae69d1..129d9f27232 100644 +index 575a4ae69d1..6160539cd97 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -@@ -679,9 +679,24 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper +@@ -679,9 +679,26 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper s"Schema did not match for query #$i\n${expected.sql}: $output") { output.schema } - assertResult(expected.output, s"Result did not match" + -+ // Comet may surface errors as `CometNativeException` instead of the matching Spark -+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a -+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it). -+ // Same category, different surface. Collapse both sides to a placeholder when this happens -+ // so the literal compare passes. ++ // With parquet filter pushdown, DataFusion 54/55 formats an error raised while ++ // evaluating a pushed-down predicate into a string, so a divide-by-zero from a ++ // `ScalaUDF` dispatched into the native plan surfaces as `CometNativeException` ++ // instead of `SparkArithmeticException`. Collapse both sides to a placeholder in that ++ // case only. Remove once Comet uses a DataFusion release that includes ++ // https://github.com/apache/datafusion/pull/24638. + // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517 + val (expectedOut, actualOut) = if (isCometEnabled && + expected.output.startsWith("org.apache.spark.SparkArithmeticException") && + expected.output.contains("\"DIVIDE_BY_ZERO\"") && + output.output.startsWith("org.apache.comet.CometNativeException") && ++ output.output.contains("Error evaluating filter predicate") && + output.output.contains("DivideByZero")) { + ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]") + } else { diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index a76ac1fd0a8..d6101e51410 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -1605,7 +1605,7 @@ index 74cdee49e55..da1e9bd966a 100644 df.collect() .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -index 23f0144dcec..40d536bb23a 100644 +index 23f0144dcec..e4c2ff6e7ac 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -166,7 +166,11 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper @@ -1621,21 +1621,23 @@ index 23f0144dcec..40d536bb23a 100644 ) ++ otherIgnoreList /** List of test cases that require TPCDS table schemas to be loaded. */ private def requireTPCDSCases: Seq[String] = Seq("pipe-operators.sql") -@@ -682,9 +686,24 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper +@@ -682,9 +686,26 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession with SQLHelper s"Schema did not match for query #$i\n${expected.sql}: $output") { output.schema } - assertResult(expected.output, s"Result did not match" + -+ // Comet may surface errors as `CometNativeException` instead of the matching Spark -+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a -+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it). -+ // Same category, different surface. Collapse both sides to a placeholder when this happens -+ // so the literal compare passes. ++ // With parquet filter pushdown, DataFusion 54/55 formats an error raised while ++ // evaluating a pushed-down predicate into a string, so a divide-by-zero from a ++ // `ScalaUDF` dispatched into the native plan surfaces as `CometNativeException` ++ // instead of `SparkArithmeticException`. Collapse both sides to a placeholder in that ++ // case only. Remove once Comet uses a DataFusion release that includes ++ // https://github.com/apache/datafusion/pull/24638. + // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517 + val (expectedOut, actualOut) = if (isCometEnabled && + expected.output.startsWith("org.apache.spark.SparkArithmeticException") && + expected.output.contains("\"DIVIDE_BY_ZERO\"") && + output.output.startsWith("org.apache.comet.CometNativeException") && ++ output.output.contains("Error evaluating filter predicate") && + output.output.contains("DivideByZero")) { + ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]") + } else { diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 37fc70b53db..2a082328cb9 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -1654,7 +1654,7 @@ index da6f6aca2ad..64c981ad3d7 100644 df.collect() .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala -index 395cb67f441..33ac6ed19af 100644 +index 395cb67f441..e824859d597 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -190,6 +190,33 @@ class SQLQueryTestSuite extends SharedSparkSession with SQLHelper @@ -1691,21 +1691,23 @@ index 395cb67f441..33ac6ed19af 100644 /** List of test cases that require TPCDS table schemas to be loaded. */ private def requireTPCDSCases: Seq[String] = Seq("pipe-operators.sql") /** List of TPCDS table names and schemas to load from the [[TPCDSSchema]] base class. */ -@@ -829,9 +856,24 @@ class SQLQueryTestSuite extends SharedSparkSession with SQLHelper +@@ -829,9 +856,26 @@ class SQLQueryTestSuite extends SharedSparkSession with SQLHelper s"Schema did not match for query #$i\n${expected.sql}: $output") { output.schema } - assertResult(expected.output, s"Result did not match" + -+ // Comet may surface errors as `CometNativeException` instead of the matching Spark -+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a -+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it). -+ // Same category, different surface. Collapse both sides to a placeholder when this happens -+ // so the literal compare passes. ++ // With parquet filter pushdown, DataFusion 54/55 formats an error raised while ++ // evaluating a pushed-down predicate into a string, so a divide-by-zero from a ++ // `ScalaUDF` dispatched into the native plan surfaces as `CometNativeException` ++ // instead of `SparkArithmeticException`. Collapse both sides to a placeholder in that ++ // case only. Remove once Comet uses a DataFusion release that includes ++ // https://github.com/apache/datafusion/pull/24638. + // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517 + val (expectedOut, actualOut) = if (isCometEnabled && + expected.output.startsWith("org.apache.spark.SparkArithmeticException") && + expected.output.contains("\"DIVIDE_BY_ZERO\"") && + output.output.startsWith("org.apache.comet.CometNativeException") && ++ output.output.contains("Error evaluating filter predicate") && + output.output.contains("DivideByZero")) { + ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]") + } else { diff --git a/native/common/src/error.rs b/native/common/src/error.rs index a33613f3099..7c233966fd8 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -1021,6 +1021,18 @@ impl std::fmt::Display for SparkErrorWithContext { impl std::error::Error for SparkErrorWithContext {} +/// Iterates over `err` and every error reachable through `Error::source()`. +/// +/// DataFusion, Arrow and Parquet wrap an external error in several layers before it reaches +/// the JNI boundary (`External`, `Context`, `Shared`, `Diagnostic`, `ArrowError::ExternalError`, +/// `ParquetError::External`, ...). All of them expose the wrapped error through `source()`, so +/// callers looking for a `SparkError` should search this chain rather than match on variants. +pub fn error_chain<'a>( + err: &'a (dyn std::error::Error + 'static), +) -> impl Iterator { + std::iter::successors(Some(err), |e| e.source()) +} + impl From for SparkErrorWithContext { fn from(error: SparkError) -> Self { SparkErrorWithContext::new(error) @@ -1058,6 +1070,53 @@ pub fn unwrap_arrow_external_error(error: ArrowError) -> DataFusionError { mod tests { use super::*; + fn find_spark_error(err: &DataFusionError) -> Option<&SparkError> { + error_chain(err).find_map(|e| e.downcast_ref::()) + } + + #[test] + fn test_error_chain_finds_spark_error_through_wrappers() { + let external = || DataFusionError::from(SparkError::DivideByZero); + let cases = [ + ("External", external()), + ("Context", external().context("evaluating expression")), + ( + "External(External)", + DataFusionError::External(Box::new(external())), + ), + ("Shared", DataFusionError::Shared(Arc::new(external()))), + ( + "Diagnostic", + external().with_diagnostic(datafusion::common::Diagnostic::new_error( + "divide by zero", + None, + )), + ), + ( + "ArrowError(ExternalError)", + DataFusionError::ArrowError( + Box::new(ArrowError::ExternalError(Box::new( + external().context("Error evaluating filter predicate"), + ))), + None, + ), + ), + ]; + for (name, err) in &cases { + assert!( + matches!(find_spark_error(err), Some(SparkError::DivideByZero)), + "SparkError not found through {name}: {err:?}" + ); + } + } + + #[test] + fn test_error_chain_without_spark_error() { + let err = DataFusionError::Execution("boom".to_string()).context("outer"); + assert!(find_spark_error(&err).is_none()); + assert_eq!(error_chain(&err).count(), 2); + } + #[test] fn test_divide_by_zero_json() { let error = SparkError::DivideByZero; diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index b1b155aba82..cf53243810d 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -28,8 +28,8 @@ mod utf8; mod utils; pub use error::{ - decimal_overflow_error, unwrap_arrow_external_error, SparkError, SparkErrorWithContext, - SparkResult, + decimal_overflow_error, error_chain, unwrap_arrow_external_error, SparkError, + SparkErrorWithContext, SparkResult, }; pub use ffi_offsets::zero_offsets; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; diff --git a/native/jni-bridge/src/errors.rs b/native/jni-bridge/src/errors.rs index a89742a9e85..04c20611f3c 100644 --- a/native/jni-bridge/src/errors.rs +++ b/native/jni-bridge/src/errors.rs @@ -19,7 +19,7 @@ use arrow::error::ArrowError; use datafusion::common::DataFusionError; -use datafusion_comet_common::{SparkError, SparkErrorWithContext}; +use datafusion_comet_common::{error_chain, SparkError, SparkErrorWithContext}; use jni::errors::{Exception, ToException}; use regex::Regex; @@ -736,17 +736,14 @@ fn jvm_exception(error: &CometError) -> JvmException<'_> { message: msg, throwable, }, - // Handle DataFusion errors containing SparkError or SparkErrorWithContext - CometError::DataFusion { - msg: _, - source: df_error @ DataFusionError::External(e), - } => { - if let Some(CometError::JavaException { + // Handle DataFusion errors carrying a Spark-typed payload at any depth of the + // `source()` chain (see `extract_spark_payload`). + CometError::DataFusion { msg: _, source } => match extract_spark_payload(source) { + Some(SparkPayload::JavaException { class, msg, throwable, - }) = e.downcast_ref::() - { + }) => { // A Java exception captured inside a JVM UDF kernel (e.g. Spark codegen // raising INVALID_REGEXP_REPLACE). Re-throw the original throwable so callers // see the exact Spark exception type rather than a wrapped CometNativeException. @@ -755,58 +752,78 @@ fn jvm_exception(error: &CometError) -> JvmException<'_> { message: msg, throwable, } - } else if let Some(spark_error_with_ctx) = e.downcast_ref::() { - JvmException::Spark { - json: spark_error_with_ctx.to_json(), - } - } else if let Some(spark_error) = e.downcast_ref::() { - spark_exception(spark_error) - } else if let Some(spark_error) = try_classify_file_read_error(df_error) { - spark_exception(&spark_error) - } else { - // Not a SparkError, use generic exception - natural_exception(error) - } - } - // Typed file-read errors (corrupt/truncated parquet, object_store) raised by the native - // scan -- classified by DataFusionError variant, not message text -- surfaced as - // FAILED_READ_FILE / FileNotFound via the structured SparkError channel. Anything else - // falls back to generic handling. - CometError::DataFusion { msg: _, source } => { - if let Some(spark_error) = parquet_external_spark_error(source) { - spark_exception(spark_error) - } else if let Some(spark_error) = try_classify_file_read_error(source) { - spark_exception(&spark_error) - } else { - generic_exception(error) } - } + Some(SparkPayload::WithContext(ctx)) => JvmException::Spark { + json: ctx.to_json(), + }, + Some(SparkPayload::Bare(spark_error)) => spark_exception(spark_error), + // Typed file-read errors (corrupt/truncated parquet, object_store) raised by the + // native scan -- classified by DataFusionError variant, not message text -- surfaced + // as FAILED_READ_FILE / FileNotFound via the structured SparkError channel. + None => match try_classify_file_read_error(source) { + Some(spark_error) => spark_exception(&spark_error), + None if matches!(source, DataFusionError::External(_)) => natural_exception(error), + None => generic_exception(error), + }, + }, // Handle direct SparkError - serialize to JSON CometError::Spark(spark_error) => spark_exception(spark_error), _ => generic_exception(error), } } +/// Payload recovered from a DataFusionError chain. +enum SparkPayload<'a> { + JavaException { + class: &'a str, + msg: &'a str, + throwable: &'a Global>, + }, + WithContext(&'a SparkErrorWithContext), + Bare(&'a SparkError), +} + +/// Returns the first Spark-typed payload found in `err`'s `source()` chain, or `None`. +/// +/// The payload is usually wrapped by DataFusion (`External`, `Context`, `Shared`, ...) and, for +/// predicates pushed into the parquet reader, also by Arrow and Parquet errors. Each layer is +/// checked for a Java throwable first, then a `SparkErrorWithContext`, then a bare `SparkError`. +fn extract_spark_payload(err: &DataFusionError) -> Option> { + error_chain(err).find_map(|e| { + if let Some(CometError::JavaException { + class, + msg, + throwable, + }) = e.downcast_ref::() + { + Some(SparkPayload::JavaException { + class, + msg, + throwable, + }) + } else if let Some(ctx) = e.downcast_ref::() { + Some(SparkPayload::WithContext(ctx)) + } else { + e.downcast_ref::().map(SparkPayload::Bare) + } + }) +} + /// Walks the cause chain for a `CometError` variant that owns its JVM exception class, so class and /// message survive DataFusion's `Context` / `Shared` / `External` wrappers. The class comes from /// `to_exception`, keeping one variant-to-class table. Classification is by variant only -- an /// error that merely mentions the same words keeps its normal classification. fn typed_jvm_exception(error: &(dyn std::error::Error + 'static)) -> Option { - let mut cause = Some(error); - while let Some(error) = cause { - if let Some( + error_chain(error).find_map(|e| match e.downcast_ref::() { + Some( typed @ (CometError::ShuffleSizeLimit(_) | CometError::IllegalState(_) | CometError::NoSuchElement(_) | CometError::NativeUdfNotAllowed(_) | CometError::NativeUdfAbi(_)), - ) = error.downcast_ref::() - { - return Some(typed.to_exception()); - } - cause = error.source(); - } - None + ) => Some(typed.to_exception()), + _ => None, + }) } /// Generic fallback for an error that isn't a structured `SparkError`. Recognises a file-not-found @@ -882,23 +899,6 @@ fn spark_exception(spark_error: &SparkError) -> JvmException<'static> { } } -/// A `SparkError` the Parquet reader raised on open arrives as -/// `DataFusionError::ParquetError(ParquetError::External(spark_error))`. Unwrap it so the error -/// keeps its own JVM exception class instead of being classified as a file read failure. -/// `Context` and `Shared` wrappers are looked through, as `try_classify_file_read_error` does. -fn parquet_external_spark_error(error: &DataFusionError) -> Option<&SparkError> { - use datafusion::common::DataFusionError as DFE; - match error { - DFE::ParquetError(pe) => match pe.as_ref() { - ParquetError::External(inner) => inner.downcast_ref::(), - _ => None, - }, - DFE::Context(_, inner) => parquet_external_spark_error(inner), - DFE::Shared(inner) => parquet_external_spark_error(inner), - _ => None, - } -} - /// Classify a `DataFusionError` as a per-file read failure by TYPED variant (not message text), /// returning `SparkError::CannotReadFile` if so. This is the structured replacement for the /// previous JVM-side substring matching on error prose. @@ -1963,30 +1963,35 @@ mod tests { /// wrapper and through `Context` and `Shared` wrappers, while an ordinary reader error does /// not match. #[test] - fn parquet_external_spark_error_keeps_its_type() { + fn extract_spark_payload_through_parquet_external() { let raised = DataFusionError::ParquetError(Box::new(ParquetError::External(Box::new( SparkError::ParquetMissingFieldIds { file_path: "a.parquet".to_string(), }, )))); assert!(matches!( - parquet_external_spark_error(&raised), - Some(SparkError::ParquetMissingFieldIds { file_path }) if file_path == "a.parquet" + extract_spark_payload(&raised), + Some(SparkPayload::Bare(SparkError::ParquetMissingFieldIds { file_path })) + if file_path == "a.parquet" )); let wrapped = DataFusionError::Context("open".to_string(), Box::new(raised)); assert!(matches!( - parquet_external_spark_error(&wrapped), - Some(SparkError::ParquetMissingFieldIds { .. }) + extract_spark_payload(&wrapped), + Some(SparkPayload::Bare( + SparkError::ParquetMissingFieldIds { .. } + )) )); let shared = DataFusionError::Shared(Arc::new(wrapped)); assert!(matches!( - parquet_external_spark_error(&shared), - Some(SparkError::ParquetMissingFieldIds { .. }) + extract_spark_payload(&shared), + Some(SparkPayload::Bare( + SparkError::ParquetMissingFieldIds { .. } + )) )); let corrupt = DataFusionError::ParquetError(Box::new(ParquetError::General( "corrupt footer".to_string(), ))); - assert!(parquet_external_spark_error(&corrupt).is_none()); + assert!(extract_spark_payload(&corrupt).is_none()); } #[test] @@ -2189,6 +2194,49 @@ mod tests { ); } + #[test] + fn extract_spark_payload_through_parquet_row_filter() { + // Shape produced by a pushed-down parquet row filter once DataFusion keeps the typed + // error (apache/datafusion#24638): the parquet reader wraps the ArrowError it gets + // from the predicate, which wraps the DataFusionError from the expression. + let predicate_err = + DataFusionError::from(SparkErrorWithContext::new(SparkError::DivideByZero)) + .context("Error evaluating filter predicate"); + let e = DataFusionError::ParquetError(Box::new(parquet::errors::ParquetError::External( + Box::new(ArrowError::ExternalError(Box::new(predicate_err))), + ))); + match extract_spark_payload(&e) { + Some(SparkPayload::WithContext(ctx)) => { + assert!(matches!(ctx.error, SparkError::DivideByZero)) + } + _ => panic!("expected SparkErrorWithContext in {e:?}"), + } + } + + #[test] + fn extract_spark_payload_keeps_context_through_wrappers() { + // A `SparkErrorWithContext` under a wrapper is returned as is, so its SQL context is kept. + let e = DataFusionError::from(SparkErrorWithContext::new(SparkError::DivideByZero)) + .context("outer"); + assert!(matches!( + extract_spark_payload(&e), + Some(SparkPayload::WithContext(_)) + )); + } + + #[test] + fn extract_spark_payload_ignores_stringified_errors() { + // DataFusion 54/55 formats a failing row-filter predicate into a string, so the typed + // payload cannot be recovered and the error falls through to generic handling. + let e = DataFusionError::ArrowError( + Box::new(ArrowError::ComputeError( + "Error evaluating filter predicate: External(DivideByZero)".to_string(), + )), + None, + ); + assert!(extract_spark_payload(&e).is_none()); + } + #[test] fn classify_non_file_errors_are_not_file_read() { // Execution / Internal errors (and object_store Generic config errors, which arrive as diff --git a/native/spark-expr/src/math_funcs/checked_binary_expr.rs b/native/spark-expr/src/math_funcs/checked_binary_expr.rs index 5602b390eee..d269d3baefc 100644 --- a/native/spark-expr/src/math_funcs/checked_binary_expr.rs +++ b/native/spark-expr/src/math_funcs/checked_binary_expr.rs @@ -25,6 +25,7 @@ use arrow::record_batch::RecordBatch; use datafusion::common::DataFusionError; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_common::error_chain; use crate::{QueryContext, SparkError, SparkErrorWithContext}; @@ -94,19 +95,20 @@ impl PhysicalExpr for CheckedBinaryExpr { } fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result { - let result = self.child.evaluate(batch); - - // If there's an error and we have query_context, wrap it - match result { - Err(DataFusionError::External(e)) if self.query_context.is_some() => { - if let Some(spark_err) = e.downcast_ref::() { - let wrapped = SparkErrorWithContext::with_context( - spark_err.clone(), - Arc::clone(self.query_context.as_ref().unwrap()), - ); - Err(DataFusionError::External(Box::new(wrapped))) - } else { - Err(DataFusionError::External(e)) + match self.child.evaluate(batch) { + Err(e) if self.query_context.is_some() => { + let spark_err = error_chain(&e) + .find_map(|e| e.downcast_ref::()) + .cloned(); + match spark_err { + Some(spark_err) => { + let wrapped = SparkErrorWithContext::with_context( + spark_err, + Arc::clone(self.query_context.as_ref().unwrap()), + ); + Err(DataFusionError::External(Box::new(wrapped))) + } + None => Err(e), } } other => other, @@ -132,3 +134,129 @@ impl PhysicalExpr for CheckedBinaryExpr { } } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Field; + + /// Child expression whose `evaluate` always fails with `make_err()`. Instances compare by + /// identity, which is all the `PhysicalExpr` bounds need here. + #[derive(Debug)] + struct FailingExpr { + make_err: fn() -> DataFusionError, + } + + impl Display for FailingExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FailingExpr") + } + } + + impl PartialEq for FailingExpr { + fn eq(&self, other: &Self) -> bool { + std::ptr::eq(self, other) + } + } + + impl Eq for FailingExpr {} + + impl PartialEq for FailingExpr { + fn eq(&self, other: &dyn Any) -> bool { + other + .downcast_ref::() + .map(|x| self.eq(x)) + .unwrap_or(false) + } + } + + impl Hash for FailingExpr { + fn hash(&self, state: &mut H) { + std::ptr::hash(self, state); + } + } + + impl PhysicalExpr for FailingExpr { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FailingExpr") + } + + fn data_type(&self, _: &Schema) -> datafusion::common::Result { + Ok(DataType::Int32) + } + + fn nullable(&self, _: &Schema) -> datafusion::common::Result { + Ok(true) + } + + fn evaluate(&self, _: &RecordBatch) -> datafusion::common::Result { + Err((self.make_err)()) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> datafusion::common::Result> { + Ok(self) + } + } + + fn query_context(sql: &str) -> Arc { + Arc::new(QueryContext::new(sql.to_string(), 0, 4, None, None, 1, 0)) + } + + /// Evaluates a `CheckedBinaryExpr` carrying `ctx` over a child that fails with `make_err()`. + fn evaluate_err(make_err: fn() -> DataFusionError, ctx: &Arc) -> DataFusionError { + let expr = + CheckedBinaryExpr::new(Arc::new(FailingExpr { make_err }), Some(Arc::clone(ctx))); + let batch = RecordBatch::new_empty(Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Int32, + true, + )]))); + expr.evaluate(&batch).expect_err("child fails") + } + + /// The `SparkErrorWithContext` the result carries directly in `External`. + fn with_context(err: &DataFusionError) -> &SparkErrorWithContext { + match err { + DataFusionError::External(e) => e + .downcast_ref::() + .unwrap_or_else(|| panic!("expected SparkErrorWithContext, got {err:?}")), + other => panic!("expected External, got {other:?}"), + } + } + + #[test] + fn attaches_query_context_to_wrapped_spark_error() { + let ctx = query_context("1 / 0"); + let err = evaluate_err( + || DataFusionError::from(SparkError::DivideByZero).context("evaluating child"), + &ctx, + ); + let wrapped = with_context(&err); + assert!(matches!(wrapped.error, SparkError::DivideByZero)); + assert!(Arc::ptr_eq(wrapped.context.as_ref().unwrap(), &ctx)); + } + + #[test] + fn keeps_inner_query_context() { + let outer = query_context("outer"); + let err = evaluate_err( + || { + DataFusionError::from(SparkErrorWithContext::with_context( + SparkError::DivideByZero, + query_context("inner"), + )) + }, + &outer, + ); + let wrapped = with_context(&err); + assert!(matches!(wrapped.error, SparkError::DivideByZero)); + assert_eq!(wrapped.context.as_ref().unwrap().sql_text.as_str(), "inner"); + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 63343c23c7c..8771e6dbc3a 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -3095,6 +3095,35 @@ class CometCodegenSuite } } + test("divide-by-zero through dispatched ScalaUDF surfaces SparkArithmeticException") { + // Native arithmetic on the result of a dispatched ScalaUDF must surface the same Spark + // exception and error class as Spark, not CometNativeException. + checkDispatchedUdfError("SELECT 1 / identity_int(a) FROM t", "DIVIDE_BY_ZERO") + } + + test("integer overflow through dispatched ScalaUDF surfaces SparkArithmeticException") { + // A second SparkError variant through the same path. + checkDispatchedUdfError("SELECT identity_int(a) + 2147483647 FROM t", "ARITHMETIC_OVERFLOW") + } + + /** + * Runs `query` over `t (a INT)` holding 0 and 1 with ANSI enabled, where the query applies + * native arithmetic to the result of the dispatched ScalaUDF `identity_int`, and checks that + * Comet fails with the same exception and `errorClass` as Spark. + */ + private def checkDispatchedUdfError(query: String, errorClass: String): Unit = { + spark.udf.register("identity_int", (i: java.lang.Integer) => i) + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("t") { + sql("CREATE TABLE t (a INT) USING parquet") + sql("INSERT INTO t VALUES (0), (1)") + assertCodegenRan { + checkSparkError(sql(query), errorClass) + } + } + } + } + test("Invoke routes through the codegen dispatcher (#5575)") { // `Invoke` has no serde of its own beyond the catch-all, so this also pins the registration: // reaching a `JvmScalarUdf` proto means `QueryPlanSerde` resolved `CometInvoke`. The target is diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 7cd9923cc45..b8a645a9eb6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -3610,4 +3610,27 @@ class CometJoinSuite extends CometTestBase { } } } + + test("Spark error on the hash join build side keeps its exception class") { + // The aggregate output is already partitioned on the join key, so the build side runs in + // the same native plan as the join. DataFusion hands the build-side error to the probe + // side wrapped in `DataFusionError::Shared`, and the JNI bridge has to look through it. + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTable("ta", "tb") { + sql("CREATE TABLE ta (k INT) USING parquet") + sql("INSERT INTO ta VALUES (1), (2)") + sql("CREATE TABLE tb (k INT, v INT) USING parquet") + sql("INSERT INTO tb VALUES (1, 0), (2, 0)") + val df = sql("""SELECT /*+ SHUFFLE_HASH(b) */ a.k, b.x + |FROM ta a JOIN (SELECT k, 1 / sum(v) AS x FROM tb GROUP BY k) b + |ON a.k = b.k""".stripMargin) + checkSparkError(df, "DIVIDE_BY_ZERO") + val joins = df.queryExecution.executedPlan.collect { case j: CometHashJoinExec => j } + assert(joins.size == 1, df.queryExecution.executedPlan) + } + } + } }