feat: [ANSI] Ansi sql error messages

Author: parthchandra

common/src/main/java/org/apache/comet/exceptions/CometQueryExecutionException.java

@@ -0,0 +1,62 @@
+/*
+ * 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.exceptions;
+
+import org.apache.comet.CometNativeException;
+
+/**
+ * Exception thrown from Comet native execution containing JSON-encoded error information. The
+ * message contains a JSON object with the following structure:
+ *
+ * <pre>
+ * {
+ *   "errorType": "DivideByZero",
+ *   "errorClass": "DIVIDE_BY_ZERO",
+ *   "params": { ... },
+ *   "context": { "sqlText": "...", "startOffset": 0, "stopOffset": 10 },
+ *   "hint": "Use `try_divide` to tolerate divisor being 0"
+ * }
+ * </pre>
+ *
+ * CometExecIterator parses this JSON and converts it to the appropriate Spark exception by calling
+ * the corresponding QueryExecutionErrors.* method.
+ */
+public final class CometQueryExecutionException extends CometNativeException {
+
+  /**
+   * Creates a new CometQueryExecutionException with a JSON-encoded error message.
+   *
+   * @param jsonMessage JSON string containing error information
+   */
+  public CometQueryExecutionException(String jsonMessage) {
+    super(jsonMessage);
+  }
+
+  /**
+   * Returns true if the message appears to be JSON-formatted. This is used to distinguish between
+   * JSON-encoded errors and legacy error messages.
+   *
+   * @return true if message starts with '{' and ends with '}'
+   */
+  public boolean isJsonMessage() {
+    String msg = getMessage();
+    return msg != null && msg.trim().startsWith("{") && msg.trim().endsWith("}");
+  }
+}

native/Cargo.lock

@@ -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, JValue};
+use jni::objects::{GlobalRef, JThrowable};
 use jni::JNIEnv;
 use lazy_static::lazy_static;
 use parquet::errors::ParquetError;
@@ -223,9 +223,9 @@ impl jni::errors::ToException for CometError {
                 class: "java/lang/NullPointerException".to_string(),
                 msg: self.to_string(),
             },
-            CometError::Spark { .. } => Exception {
-                class: "org/apache/spark/SparkException".to_string(),
-                msg: self.to_string(),
+            CometError::Spark(spark_err) => Exception {
+                class: spark_err.exception_class().to_string(),
+                msg: spark_err.to_string(),
             },
             CometError::NumberIntFormat { source: s } => Exception {
                 class: "java/lang/NumberFormatException".to_string(),
@@ -392,33 +392,37 @@ fn throw_exception(env: &mut JNIEnv, error: &CometError, backtrace: Option<Strin
                         throwable,
                     },
             } => env.throw(<&JThrowable>::from(throwable.as_obj())),
+            // Handle DataFusion errors containing SparkError - serialize to JSON
             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)
+            } => {
+                // Try SparkErrorWithContext first (includes context)
+                if let Some(spark_error_with_ctx) =
+                    e.downcast_ref::<datafusion_comet_spark_expr::SparkErrorWithContext>()
+                {
+                    let json_message = spark_error_with_ctx.to_json();
+                    env.throw_new(
+                        "org/apache/comet/exceptions/CometQueryExecutionException",
+                        json_message,
+                    )
+                } else if let Some(spark_error) = e.downcast_ref::<SparkError>() {
+                    // Fall back to plain SparkError (no context)
+                    throw_spark_error_as_json(env, spark_error)
+                } else {
+                    // Not a SparkError, use generic exception
+                    let exception = error.to_exception();
+                    match backtrace {
+                        Some(backtrace_string) => env.throw_new(
+                            exception.class,
+                            to_stacktrace_string(exception.msg, backtrace_string).unwrap(),
+                        ),
+                        _ => env.throw_new(exception.class, exception.msg),
                     }
-                    _ => unreachable!(),
                 }
             }
+            // Handle direct SparkError - serialize to JSON
+            CometError::Spark(spark_error) => throw_spark_error_as_json(env, spark_error),
             _ => {
                 let exception = error.to_exception();
                 match backtrace {
@@ -434,6 +438,21 @@ fn throw_exception(env: &mut JNIEnv, error: &CometError, backtrace: Option<Strin
     }
 }
 
+/// Throws a CometQueryExecutionException with JSON-encoded SparkError
+fn throw_spark_error_as_json(
+    env: &mut JNIEnv,
+    spark_error: &SparkError,
+) -> jni::errors::Result<()> {
+    // Serialize error to JSON
+    let json_message = spark_error.to_json();
+
+    // Throw CometQueryExecutionException with JSON message
+    env.throw_new(
+        "org/apache/comet/exceptions/CometQueryExecutionException",
+        json_message,
+    )
+}
+
 #[derive(Debug, Error)]
 enum StacktraceError {
     #[error("Unable to initialize message: {0}")]

native/core/src/errors.rs

@@ -17,6 +17,122 @@
 
 //! Arithmetic expression builders
 
+use std::any::Any;
+use std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
+
+use arrow::datatypes::{DataType, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion::common::DataFusionError;
+use datafusion::logical_expr::ColumnarValue;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::{QueryContext, SparkError, SparkErrorWithContext};
+
+/// Wrapper expression that catches and wraps SparkError with QueryContext
+/// for binary arithmetic operations.
+#[derive(Debug)]
+pub struct CheckedBinaryExpr {
+    /// The underlying physical expression (typically a ScalarFunctionExpr)
+    child: Arc<dyn PhysicalExpr>,
+    /// Optional query context to attach to errors
+    query_context: Option<Arc<QueryContext>>,
+}
+
+impl CheckedBinaryExpr {
+    pub fn new(child: Arc<dyn PhysicalExpr>, query_context: Option<Arc<QueryContext>>) -> Self {
+        Self {
+            child,
+            query_context,
+        }
+    }
+}
+
+impl Display for CheckedBinaryExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "CheckedBinaryExpr({})", self.child)
+    }
+}
+
+impl PartialEq for CheckedBinaryExpr {
+    fn eq(&self, other: &Self) -> bool {
+        self.child.eq(&other.child)
+    }
+}
+
+impl Eq for CheckedBinaryExpr {}
+
+impl PartialEq<dyn Any> for CheckedBinaryExpr {
+    fn eq(&self, other: &dyn Any) -> bool {
+        other
+            .downcast_ref::<Self>()
+            .map(|x| self.eq(x))
+            .unwrap_or(false)
+    }
+}
+
+impl Hash for CheckedBinaryExpr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.child.hash(state);
+    }
+}
+
+impl PhysicalExpr for CheckedBinaryExpr {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.child.fmt_sql(f)
+    }
+
+    fn data_type(&self, input_schema: &Schema) -> datafusion::common::Result<DataType> {
+        self.child.data_type(input_schema)
+    }
+
+    fn nullable(&self, input_schema: &Schema) -> datafusion::common::Result<bool> {
+        self.child.nullable(input_schema)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result<ColumnarValue> {
+        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::<SparkError>() {
+                    let wrapped = SparkErrorWithContext::with_context(
+                        spark_err.clone(),
+                        self.query_context.as_ref().unwrap().clone(),
+                    );
+                    Err(DataFusionError::External(Box::new(wrapped)))
+                } else {
+                    Err(DataFusionError::External(e))
+                }
+            }
+            other => other,
+        }
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.child]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> datafusion::common::Result<Arc<dyn PhysicalExpr>> {
+        match children.len() {
+            1 => Ok(Arc::new(CheckedBinaryExpr::new(
+                Arc::clone(&children[0]),
+                self.query_context.clone(),
+            ))),
+            _ => Err(DataFusionError::Internal(
+                "CheckedBinaryExpr should have exactly one child".to_string(),
+            )),
+        }
+    }
+}
+
 /// Macro to generate arithmetic expression builders that need eval_mode handling
 #[macro_export]
 macro_rules! arithmetic_expr_builder {
@@ -37,6 +153,7 @@ macro_rules! arithmetic_expr_builder {
                 let eval_mode =
                     $crate::execution::planner::from_protobuf_eval_mode(expr.eval_mode)?;
                 planner.create_binary_expr(
+                    spark_expr, // Pass the full spark_expr for query_context lookup
                     expr.left.as_ref().unwrap(),
                     expr.right.as_ref().unwrap(),
                     expr.return_type.as_ref(),
@@ -53,7 +170,6 @@ use std::sync::Arc;
 
 use arrow::datatypes::SchemaRef;
 use datafusion::logical_expr::Operator as DataFusionOperator;
-use datafusion::physical_expr::PhysicalExpr;
 use datafusion_comet_proto::spark_expression::Expr;
 use datafusion_comet_spark_expr::{create_modulo_expr, create_negate_expr, EvalMode};
 
@@ -95,6 +211,7 @@ impl ExpressionBuilder for IntegralDivideBuilder {
         let expr = extract_expr!(spark_expr, IntegralDivide);
         let eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
         planner.create_binary_expr_with_options(
+            spark_expr, // Pass the full spark_expr for query_context lookup
             expr.left.as_ref().unwrap(),
             expr.right.as_ref().unwrap(),
             expr.return_type.as_ref(),